Merge branch 'sea-migration' into sea-norm-cols

Author: varun-edachali-dbx

examples/experimental/tests/test_sea_sync_query.py

@@ -72,6 +72,13 @@ def test_sea_sync_query_with_cloud_fetch():
             f"{actual_row_count} rows retrieved against {requested_row_count} requested"
         )
 
+        # Verify total row count
+        if actual_row_count != requested_row_count:
+            logger.error(
+                f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
+            )
+            return False
+
         # Close resources
         cursor.close()
         connection.close()
@@ -132,15 +139,27 @@ def test_sea_sync_query_without_cloud_fetch():
         # For non-cloud fetch, use a smaller row count to avoid exceeding inline limits
         requested_row_count = 100
         cursor = connection.cursor()
-        logger.info("Executing synchronous query without cloud fetch: SELECT 100 rows")
+        logger.info(
+            f"Executing synchronous query without cloud fetch: SELECT {requested_row_count} rows"
+        )
         cursor.execute(
             "SELECT id, 'test_value_' || CAST(id as STRING) as test_value FROM range(1, 101)"
         )
 
         results = [cursor.fetchone()]
         results.extend(cursor.fetchmany(10))
         results.extend(cursor.fetchall())
-        logger.info(f"{len(results)} rows retrieved against 100 requested")
+        actual_row_count = len(results)
+        logger.info(
+            f"{actual_row_count} rows retrieved against {requested_row_count} requested"
+        )
+
+        # Verify total row count
+        if actual_row_count != requested_row_count:
+            logger.error(
+                f"FAIL: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
+            )
+            return False
 
         # Close resources
         cursor.close()

src/databricks/sql/backend/databricks_client.py

@@ -96,7 +96,7 @@ def execute_command(
             max_rows: Maximum number of rows to fetch in a single fetch batch
             max_bytes: Maximum number of bytes to fetch in a single fetch batch
             lz4_compression: Whether to use LZ4 compression for result data
-            cursor: The cursor object that will handle the results
+            cursor: The cursor object that will handle the results. The command id is set in this cursor.
             use_cloud_fetch: Whether to use cloud fetch for retrieving large result sets
             parameters: List of parameters to bind to the query
             async_op: Whether to execute the command asynchronously
@@ -282,7 +282,9 @@ def get_tables(
             max_bytes: Maximum number of bytes to fetch in a single batch
             cursor: The cursor object that will handle the results
             catalog_name: Optional catalog name pattern to filter by
+                if catalog_name is None, we fetch across all catalogs
             schema_name: Optional schema name pattern to filter by
+                if schema_name is None, we fetch across all schemas
             table_name: Optional table name pattern to filter by
             table_types: Optional list of table types to filter by (e.g., ['TABLE', 'VIEW'])
 
@@ -321,6 +323,7 @@ def get_columns(
             catalog_name: Optional catalog name pattern to filter by
             schema_name: Optional schema name pattern to filter by
             table_name: Optional table name pattern to filter by
+                if table_name is None, we fetch across all tables
             column_name: Optional column name pattern to filter by
 
         Returns: