formatting (black)

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

Author: varun-edachali-dbx

examples/experimental/sea_connector_test.py

@@ -22,90 +22,99 @@
     "test_sea_metadata",
 ]
 
+
 def load_test_function(module_name: str) -> Callable:
     """Load a test function from a module."""
     module_path = os.path.join(
-        os.path.dirname(os.path.abspath(__file__)),
-        "tests",
-        f"{module_name}.py"
+        os.path.dirname(os.path.abspath(__file__)), "tests", f"{module_name}.py"
     )
-    
+
     spec = importlib.util.spec_from_file_location(module_name, module_path)
     module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(module)
-    
+
     # Get the main test function (assuming it starts with "test_")
     for name in dir(module):
         if name.startswith("test_") and callable(getattr(module, name)):
             # For sync and async query modules, we want the main function that runs both tests
             if name == f"test_sea_{module_name.replace('test_sea_', '')}_exec":
                 return getattr(module, name)
-    
+
     # Fallback to the first test function found
     for name in dir(module):
         if name.startswith("test_") and callable(getattr(module, name)):
             return getattr(module, name)
-    
+
     raise ValueError(f"No test function found in module {module_name}")
 
+
 def run_tests() -> List[Tuple[str, bool]]:
     """Run all tests and return results."""
     results = []
-    
+
     for module_name in TEST_MODULES:
         try:
             test_func = load_test_function(module_name)
             logger.info(f"\n{'=' * 50}")
             logger.info(f"Running test: {module_name}")
             logger.info(f"{'-' * 50}")
-            
+
             success = test_func()
             results.append((module_name, success))
-            
+
             status = "✅ PASSED" if success else "❌ FAILED"
             logger.info(f"Test {module_name}: {status}")
-            
+
         except Exception as e:
             logger.error(f"Error loading or running test {module_name}: {str(e)}")
             import traceback
+
             logger.error(traceback.format_exc())
             results.append((module_name, False))
-    
+
     return results
 
+
 def print_summary(results: List[Tuple[str, bool]]) -> None:
     """Print a summary of test results."""
     logger.info(f"\n{'=' * 50}")
     logger.info("TEST SUMMARY")
     logger.info(f"{'-' * 50}")
-    
+
     passed = sum(1 for _, success in results if success)
     total = len(results)
-    
+
     for module_name, success in results:
         status = "✅ PASSED" if success else "❌ FAILED"
         logger.info(f"{status} - {module_name}")
-    
+
     logger.info(f"{'-' * 50}")
     logger.info(f"Total: {total} | Passed: {passed} | Failed: {total - passed}")
     logger.info(f"{'=' * 50}")
 
+
 if __name__ == "__main__":
     # Check if required environment variables are set
-    required_vars = ["DATABRICKS_SERVER_HOSTNAME", "DATABRICKS_HTTP_PATH", "DATABRICKS_TOKEN"]
+    required_vars = [
+        "DATABRICKS_SERVER_HOSTNAME",
+        "DATABRICKS_HTTP_PATH",
+        "DATABRICKS_TOKEN",
+    ]
     missing_vars = [var for var in required_vars if not os.environ.get(var)]
-    
+
     if missing_vars:
-        logger.error(f"Missing required environment variables: {', '.join(missing_vars)}")
+        logger.error(
+            f"Missing required environment variables: {', '.join(missing_vars)}"
+        )
         logger.error("Please set these variables before running the tests.")
         sys.exit(1)
-    
+
     # Run all tests
     results = run_tests()
-    
+
     # Print summary
     print_summary(results)
-    
+
     # Exit with appropriate status code
     all_passed = all(success for _, success in results)
-    sys.exit(0 if all_passed else 1)
+    sys.exit(0 if all_passed else 1)

examples/experimental/tests/__init__.py

@@ -1 +0,0 @@
-# This file makes the tests directory a Python package

examples/experimental/tests/test_sea_async_query.py

@@ -33,7 +33,9 @@ def test_sea_async_query_with_cloud_fetch():
 
     try:
         # Create connection with cloud fetch enabled
-        logger.info("Creating connection for asynchronous query execution with cloud fetch enabled")
+        logger.info(
+            "Creating connection for asynchronous query execution with cloud fetch enabled"
+        )
         connection = Connection(
             server_hostname=server_hostname,
             http_path=http_path,
@@ -51,30 +53,39 @@ def test_sea_async_query_with_cloud_fetch():
 
         # Execute a simple query asynchronously
         cursor = connection.cursor()
-        logger.info("Executing asynchronous query with cloud fetch: SELECT 1 as test_value")
+        logger.info(
+            "Executing asynchronous query with cloud fetch: SELECT 1 as test_value"
+        )
         cursor.execute_async("SELECT 1 as test_value")
-        logger.info("Asynchronous query submitted successfully with cloud fetch enabled")
-        
+        logger.info(
+            "Asynchronous query submitted successfully with cloud fetch enabled"
+        )
+
         # Check query state
         logger.info("Checking query state...")
         while cursor.is_query_pending():
             logger.info("Query is still pending, waiting...")
             time.sleep(1)
-            
+
         logger.info("Query is no longer pending, getting results...")
         cursor.get_async_execution_result()
-        logger.info("Successfully retrieved asynchronous query results with cloud fetch enabled")
-        
+        logger.info(
+            "Successfully retrieved asynchronous query results with cloud fetch enabled"
+        )
+
         # Close resources
         cursor.close()
         connection.close()
         logger.info("Successfully closed SEA session")
-        
+
         return True
 
     except Exception as e:
-        logger.error(f"Error during SEA asynchronous query execution test with cloud fetch: {str(e)}")
+        logger.error(
+            f"Error during SEA asynchronous query execution test with cloud fetch: {str(e)}"
+        )
         import traceback
+
         logger.error(traceback.format_exc())
         return False
 
@@ -100,7 +111,9 @@ def test_sea_async_query_without_cloud_fetch():
 
     try:
         # Create connection with cloud fetch disabled
-        logger.info("Creating connection for asynchronous query execution with cloud fetch disabled")
+        logger.info(
+            "Creating connection for asynchronous query execution with cloud fetch disabled"
+        )
         connection = Connection(
             server_hostname=server_hostname,
             http_path=http_path,
@@ -119,30 +132,39 @@ def test_sea_async_query_without_cloud_fetch():
 
         # Execute a simple query asynchronously
         cursor = connection.cursor()
-        logger.info("Executing asynchronous query without cloud fetch: SELECT 1 as test_value")
+        logger.info(
+            "Executing asynchronous query without cloud fetch: SELECT 1 as test_value"
+        )
         cursor.execute_async("SELECT 1 as test_value")
-        logger.info("Asynchronous query submitted successfully with cloud fetch disabled")
-        
+        logger.info(
+            "Asynchronous query submitted successfully with cloud fetch disabled"
+        )
+
         # Check query state
         logger.info("Checking query state...")
         while cursor.is_query_pending():
             logger.info("Query is still pending, waiting...")
             time.sleep(1)
-            
+
         logger.info("Query is no longer pending, getting results...")
         cursor.get_async_execution_result()
-        logger.info("Successfully retrieved asynchronous query results with cloud fetch disabled")
-        
+        logger.info(
+            "Successfully retrieved asynchronous query results with cloud fetch disabled"
+        )
+
         # Close resources
         cursor.close()
         connection.close()
         logger.info("Successfully closed SEA session")
-        
+
         return True
 
     except Exception as e:
-        logger.error(f"Error during SEA asynchronous query execution test without cloud fetch: {str(e)}")
+        logger.error(
+            f"Error during SEA asynchronous query execution test without cloud fetch: {str(e)}"
+        )
         import traceback
+
         logger.error(traceback.format_exc())
         return False
 
@@ -152,14 +174,18 @@ def test_sea_async_query_exec():
     Run both asynchronous query tests and return overall success.
     """
     with_cloud_fetch_success = test_sea_async_query_with_cloud_fetch()
-    logger.info(f"Asynchronous query with cloud fetch: {'✅ PASSED' if with_cloud_fetch_success else '❌ FAILED'}")
-    
+    logger.info(
+        f"Asynchronous query with cloud fetch: {'✅ PASSED' if with_cloud_fetch_success else '❌ FAILED'}"
+    )
+
     without_cloud_fetch_success = test_sea_async_query_without_cloud_fetch()
-    logger.info(f"Asynchronous query without cloud fetch: {'✅ PASSED' if without_cloud_fetch_success else '❌ FAILED'}")
-    
+    logger.info(
+        f"Asynchronous query without cloud fetch: {'✅ PASSED' if without_cloud_fetch_success else '❌ FAILED'}"
+    )
+
     return with_cloud_fetch_success and without_cloud_fetch_success
 
 
 if __name__ == "__main__":
     success = test_sea_async_query_exec()
-    sys.exit(0 if success else 1)
+    sys.exit(0 if success else 1)

examples/experimental/tests/test_sea_metadata.py

@@ -28,9 +28,11 @@ def test_sea_metadata():
             "Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
         )
         return False
-        
+
     if not catalog:
-        logger.error("DATABRICKS_CATALOG environment variable is required for metadata tests.")
+        logger.error(
+            "DATABRICKS_CATALOG environment variable is required for metadata tests."
+        )
         return False
 
     try:
@@ -55,37 +57,42 @@ def test_sea_metadata():
         logger.info("Fetching catalogs...")
         cursor.catalogs()
         logger.info("Successfully fetched catalogs")
-        
+
         # Test schemas
         logger.info(f"Fetching schemas for catalog '{catalog}'...")
         cursor.schemas(catalog_name=catalog)
         logger.info("Successfully fetched schemas")
-        
+
         # Test tables
         logger.info(f"Fetching tables for catalog '{catalog}', schema 'default'...")
         cursor.tables(catalog_name=catalog, schema_name="default")
         logger.info("Successfully fetched tables")
-        
+
         # Test columns for a specific table
         # Using a common table that should exist in most environments
-        logger.info(f"Fetching columns for catalog '{catalog}', schema 'default', table 'information_schema'...")
-        cursor.columns(catalog_name=catalog, schema_name="default", table_name="information_schema")
+        logger.info(
+            f"Fetching columns for catalog '{catalog}', schema 'default', table 'information_schema'..."
+        )
+        cursor.columns(
+            catalog_name=catalog, schema_name="default", table_name="information_schema"
+        )
         logger.info("Successfully fetched columns")
-        
+
         # Close resources
         cursor.close()
         connection.close()
         logger.info("Successfully closed SEA session")
-        
+
         return True
 
     except Exception as e:
         logger.error(f"Error during SEA metadata test: {str(e)}")
         import traceback
+
         logger.error(traceback.format_exc())
         return False
 
 
 if __name__ == "__main__":
     success = test_sea_metadata()
-    sys.exit(0 if success else 1)
+    sys.exit(0 if success else 1)

examples/experimental/tests/test_sea_session.py

@@ -55,16 +55,17 @@ def test_sea_session():
         logger.info("Closing the SEA session...")
         connection.close()
         logger.info("Successfully closed SEA session")
-        
+
         return True
 
     except Exception as e:
         logger.error(f"Error testing SEA session: {str(e)}")
         import traceback
+
         logger.error(traceback.format_exc())
         return False
 
 
 if __name__ == "__main__":
     success = test_sea_session()
-    sys.exit(0 if success else 1)
+    sys.exit(0 if success else 1)