access ssl_options through connection

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

examples/experimental/test_sea_multi_chunk.py

@@ -21,22 +21,22 @@
 def test_sea_multi_chunk_with_cloud_fetch(requested_row_count=5000):
     """
     Test executing a query that generates multiple chunks using cloud fetch.
-    
+
     Args:
         requested_row_count: Number of rows to request in the query
-        
+
     Returns:
         bool: True if the test passed, False otherwise
     """
     server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
     http_path = os.environ.get("DATABRICKS_HTTP_PATH")
     access_token = os.environ.get("DATABRICKS_TOKEN")
     catalog = os.environ.get("DATABRICKS_CATALOG")
-    
+
     # Create output directory for test results
     output_dir = Path("test_results")
     output_dir.mkdir(exist_ok=True)
-    
+
     # Files to store results
     rows_file = output_dir / "cloud_fetch_rows.csv"
     stats_file = output_dir / "cloud_fetch_stats.json"
@@ -50,9 +50,7 @@ def test_sea_multi_chunk_with_cloud_fetch(requested_row_count=5000):
 
     try:
         # Create connection with cloud fetch enabled
-        logger.info(
-            "Creating connection for query execution with cloud fetch enabled"
-        )
+        logger.info("Creating connection for query execution with cloud fetch enabled")
         connection = Connection(
             server_hostname=server_hostname,
             http_path=http_path,
@@ -76,46 +74,50 @@ def test_sea_multi_chunk_with_cloud_fetch(requested_row_count=5000):
             concat('value_', repeat('a', 10000)) as test_value
         FROM range(1, {requested_row_count} + 1) AS t(id)
         """
-        
-        logger.info(f"Executing query with cloud fetch to generate {requested_row_count} rows")
+
+        logger.info(
+            f"Executing query with cloud fetch to generate {requested_row_count} rows"
+        )
         start_time = time.time()
         cursor.execute(query)
-        
+
         # Fetch all rows
         rows = cursor.fetchall()
         actual_row_count = len(rows)
         end_time = time.time()
         execution_time = end_time - start_time
-        
+
         logger.info(f"Query executed in {execution_time:.2f} seconds")
-        logger.info(f"Requested {requested_row_count} rows, received {actual_row_count} rows")
-        
+        logger.info(
+            f"Requested {requested_row_count} rows, received {actual_row_count} rows"
+        )
+
         # Write rows to CSV file for inspection
         logger.info(f"Writing rows to {rows_file}")
-        with open(rows_file, 'w', newline='') as f:
+        with open(rows_file, "w", newline="") as f:
             writer = csv.writer(f)
-            writer.writerow(['id', 'value_length'])  # Header
-            
+            writer.writerow(["id", "value_length"])  # Header
+
             # Extract IDs to check for duplicates and missing values
             row_ids = []
             for row in rows:
                 row_id = row[0]
                 value_length = len(row[1])
                 writer.writerow([row_id, value_length])
                 row_ids.append(row_id)
-        
+
         # Verify row count
         success = actual_row_count == requested_row_count
-        
+
         # Check for duplicate IDs
         unique_ids = set(row_ids)
         duplicate_count = len(row_ids) - len(unique_ids)
-        
+
         # Check for missing IDs
         expected_ids = set(range(1, requested_row_count + 1))
         missing_ids = expected_ids - unique_ids
         extra_ids = unique_ids - expected_ids
-        
+
         # Write statistics to JSON file
         stats = {
             "requested_row_count": requested_row_count,
@@ -124,56 +126,64 @@ def test_sea_multi_chunk_with_cloud_fetch(requested_row_count=5000):
             "duplicate_count": duplicate_count,
             "missing_ids_count": len(missing_ids),
             "extra_ids_count": len(extra_ids),
-            "missing_ids": list(missing_ids)[:100] if missing_ids else [],  # Limit to first 100 for readability
-            "extra_ids": list(extra_ids)[:100] if extra_ids else [],  # Limit to first 100 for readability
-            "success": success and duplicate_count == 0 and len(missing_ids) == 0 and len(extra_ids) == 0
+            "missing_ids": list(missing_ids)[:100]
+            if missing_ids
+            else [],  # Limit to first 100 for readability
+            "extra_ids": list(extra_ids)[:100]
+            if extra_ids
+            else [],  # Limit to first 100 for readability
+            "success": success
+            and duplicate_count == 0
+            and len(missing_ids) == 0
+            and len(extra_ids) == 0,
         }
-        
-        with open(stats_file, 'w') as f:
+
+        with open(stats_file, "w") as f:
             json.dump(stats, f, indent=2)
-        
+
         # Log detailed results
         if duplicate_count > 0:
             logger.error(f"❌ FAILED: Found {duplicate_count} duplicate row IDs")
             success = False
         else:
             logger.info("✅ PASSED: No duplicate row IDs found")
-            
+
         if missing_ids:
             logger.error(f"❌ FAILED: Missing {len(missing_ids)} expected row IDs")
             if len(missing_ids) <= 10:
                 logger.error(f"Missing IDs: {sorted(list(missing_ids))}")
             success = False
         else:
             logger.info("✅ PASSED: All expected row IDs present")
-            
+
         if extra_ids:
             logger.error(f"❌ FAILED: Found {len(extra_ids)} unexpected row IDs")
             if len(extra_ids) <= 10:
                 logger.error(f"Extra IDs: {sorted(list(extra_ids))}")
             success = False
         else:
             logger.info("✅ PASSED: No unexpected row IDs found")
-        
+
         if actual_row_count == requested_row_count:
             logger.info("✅ PASSED: Row count matches requested count")
         else:
-            logger.error(f"❌ FAILED: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}")
+            logger.error(
+                f"❌ FAILED: Row count mismatch. Expected {requested_row_count}, got {actual_row_count}"
+            )
             success = False
-        
+
         # Close resources
         cursor.close()
         connection.close()
         logger.info("Successfully closed SEA session")
-        
+
         logger.info(f"Test results written to {rows_file} and {stats_file}")
         return success
 
     except Exception as e:
-        logger.error(
-            f"Error during SEA multi-chunk test with cloud fetch: {str(e)}"
-        )
+        logger.error(f"Error during SEA multi-chunk test with cloud fetch: {str(e)}")
         import traceback
+
         logger.error(traceback.format_exc())
         return False
 
@@ -193,31 +203,33 @@ def main():
         )
         logger.error("Please set these variables before running the tests.")
         sys.exit(1)
-    
+
     # Get row count from command line or use default
     requested_row_count = 10000
-    
+
     if len(sys.argv) > 1:
         try:
             requested_row_count = int(sys.argv[1])
         except ValueError:
             logger.error(f"Invalid row count: {sys.argv[1]}")
             logger.error("Please provide a valid integer for row count.")
             sys.exit(1)
-    
+
     logger.info(f"Testing with {requested_row_count} rows")
-    
+
     # Run the multi-chunk test with cloud fetch
     success = test_sea_multi_chunk_with_cloud_fetch(requested_row_count)
-    
+
     # Report results
     if success:
-        logger.info("✅ TEST PASSED: Multi-chunk cloud fetch test completed successfully")
+        logger.info(
+            "✅ TEST PASSED: Multi-chunk cloud fetch test completed successfully"
+        )
         sys.exit(0)
     else:
         logger.error("❌ TEST FAILED: Multi-chunk cloud fetch test encountered errors")
         sys.exit(1)
 
 
 if __name__ == "__main__":
-    main()
+    main()

examples/experimental/tests/test_sea_async_query.py

@@ -77,24 +77,29 @@ def test_sea_async_query_with_cloud_fetch():
 
         logger.info("Query is no longer pending, getting results...")
         cursor.get_async_execution_result()
-        
+
         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"
+        )
 
         logger.info(
             f"Requested {requested_row_count} rows, received {actual_row_count} rows"
         )
-        
+
         # 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
-        
-        logger.info("PASS: Received correct number of rows with cloud fetch and all fetch methods work correctly")
+
+        logger.info(
+            "PASS: Received correct number of rows with cloud fetch and all fetch methods work correctly"
+        )
 
         # Close resources
         cursor.close()
@@ -182,20 +187,25 @@ def test_sea_async_query_without_cloud_fetch():
         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"
+        )
 
         logger.info(
             f"Requested {requested_row_count} rows, received {actual_row_count} rows"
         )
-        
+
         # 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
 
-        logger.info("PASS: Received correct number of rows without cloud fetch and all fetch methods work correctly")
+        logger.info(
+            "PASS: Received correct number of rows without cloud fetch and all fetch methods work correctly"
+        )
 
         # Close resources
         cursor.close()

examples/experimental/tests/test_sea_sync_query.py

@@ -62,10 +62,14 @@ def test_sea_sync_query_with_cloud_fetch():
         logger.info(
             f"Executing synchronous query with cloud fetch to generate {requested_row_count} rows"
         )
+        cursor.execute(query)
         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"
+        )
 
         # Close resources
         cursor.close()

src/databricks/sql/result_set.py

@@ -472,23 +472,14 @@ def __init__(
             result_data: Result data from SEA response (optional)
             manifest: Manifest from SEA response (optional)
         """
-        # Extract and store SEA-specific properties
-        self.statement_id = (
-            execute_response.command_id.to_sea_statement_id()
-            if execute_response.command_id
-            else None
-        )
-
-        # Build the results queue
-        results_queue = None
 
         results_queue = None
         if result_data:
             results_queue = SeaResultSetQueueFactory.build_queue(
                 result_data,
                 manifest,
                 str(execute_response.command_id.to_sea_statement_id()),
-                ssl_options=self.connection.session.ssl_options,
+                ssl_options=connection.session.ssl_options,
                 description=execute_response.description,
                 max_download_threads=sea_client.max_download_threads,
                 sea_client=sea_client,
@@ -513,6 +504,21 @@ def __init__(
         # Initialize queue for result data if not provided
         self.results = results_queue or JsonQueue([])
 
+    def _convert_json_table(self, rows):
+        """
+        Convert raw data rows to Row objects with named columns based on description.
+        Args:
+            rows: List of raw data rows
+        Returns:
+            List of Row objects with named columns
+        """
+        if not self.description or not rows:
+            return rows
+
+        column_names = [col[0] for col in self.description]
+        ResultRow = Row(*column_names)
+        return [ResultRow(*row) for row in rows]
+
     def fetchmany_json(self, size: int):
         """
         Fetch the next set of rows as a columnar table.
@@ -586,7 +592,7 @@ def fetchone(self) -> Optional[Row]:
             A single Row object or None if no more rows are available
         """
         if isinstance(self.results, JsonQueue):
-            res = self.fetchmany_json(1)
+            res = self._convert_json_table(self.fetchmany_json(1))
         else:
             res = self._convert_arrow_table(self.fetchmany_arrow(1))
 
@@ -606,7 +612,7 @@ def fetchmany(self, size: int) -> List[Row]:
             ValueError: If size is negative
         """
         if isinstance(self.results, JsonQueue):
-            return self.fetchmany_json(size)
+            return self._convert_json_table(self.fetchmany_json(size))
         else:
             return self._convert_arrow_table(self.fetchmany_arrow(size))
 
@@ -618,6 +624,6 @@ def fetchall(self) -> List[Row]:
             List of Row objects containing all remaining rows
         """
         if isinstance(self.results, JsonQueue):
-            return self.fetchall_json()
+            return self._convert_json_table(self.fetchall_json())
         else:
             return self._convert_arrow_table(self.fetchall_arrow())