Skip to content

Commit 45585d4

Browse files
Introduce manual SEA test scripts for Exec Phase (#589)
* [squash from exec-sea] bring over execution phase changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess test Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * add docstring Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remvoe exec func in sea backend Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess files Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess models Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove excess sea backend tests Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * cleanup Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * re-introduce get_schema_desc Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove SeaResultSet Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * clean imports and attributes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * pass CommandId to ExecResp Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove changes in types Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * add back essential types (ExecResponse, from_sea_state) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * fix fetch types Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * excess imports Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * reduce diff by maintaining logs Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * fix int test types Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * [squashed from exec-sea] init execution func Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove irrelevant changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove ResultSetFilter functionality Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove more irrelevant changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove more irrelevant changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * even more irrelevant changes Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove sea response as init option Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * exec test example scripts Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * formatting (black) Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * change to valid table name Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * remove un-necessary changes covered by #588 Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * simplify test module Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * logging -> debug level Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * change table name in log Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> * add basic documentation on env vars to be set Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com> --------- Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 677a7b0 commit 45585d4

File tree

6 files changed

+632
-56
lines changed

6 files changed

+632
-56
lines changed
Lines changed: 111 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,66 +1,121 @@
1+
"""
2+
Main script to run all SEA connector tests.
3+
4+
This script runs all the individual test modules and displays
5+
a summary of test results with visual indicators.
6+
7+
In order to run the script, the following environment variables need to be set:
8+
- DATABRICKS_SERVER_HOSTNAME: The hostname of the Databricks server
9+
- DATABRICKS_HTTP_PATH: The HTTP path of the Databricks server
10+
- DATABRICKS_TOKEN: The token to use for authentication
11+
"""
12+
113
import os
214
import sys
315
import logging
4-
from databricks.sql.client import Connection
16+
import subprocess
17+
from typing import List, Tuple
518

619
logging.basicConfig(level=logging.DEBUG)
720
logger = logging.getLogger(__name__)
821

9-
def test_sea_session():
10-
"""
11-
Test opening and closing a SEA session using the connector.
12-
13-
This function connects to a Databricks SQL endpoint using the SEA backend,
14-
opens a session, and then closes it.
15-
16-
Required environment variables:
17-
- DATABRICKS_SERVER_HOSTNAME: Databricks server hostname
18-
- DATABRICKS_HTTP_PATH: HTTP path for the SQL endpoint
19-
- DATABRICKS_TOKEN: Personal access token for authentication
20-
"""
21-
22-
server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
23-
http_path = os.environ.get("DATABRICKS_HTTP_PATH")
24-
access_token = os.environ.get("DATABRICKS_TOKEN")
25-
catalog = os.environ.get("DATABRICKS_CATALOG")
26-
27-
if not all([server_hostname, http_path, access_token]):
28-
logger.error("Missing required environment variables.")
29-
logger.error("Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN.")
30-
sys.exit(1)
31-
32-
logger.info(f"Connecting to {server_hostname}")
33-
logger.info(f"HTTP Path: {http_path}")
34-
if catalog:
35-
logger.info(f"Using catalog: {catalog}")
36-
37-
try:
38-
logger.info("Creating connection with SEA backend...")
39-
connection = Connection(
40-
server_hostname=server_hostname,
41-
http_path=http_path,
42-
access_token=access_token,
43-
catalog=catalog,
44-
schema="default",
45-
use_sea=True,
46-
user_agent_entry="SEA-Test-Client" # add custom user agent
22+
TEST_MODULES = [
23+
"test_sea_session",
24+
"test_sea_sync_query",
25+
"test_sea_async_query",
26+
"test_sea_metadata",
27+
]
28+
29+
30+
def run_test_module(module_name: str) -> bool:
31+
"""Run a test module and return success status."""
32+
module_path = os.path.join(
33+
os.path.dirname(os.path.abspath(__file__)), "tests", f"{module_name}.py"
34+
)
35+
36+
# Simply run the module as a script - each module handles its own test execution
37+
result = subprocess.run(
38+
[sys.executable, module_path], capture_output=True, text=True
39+
)
40+
41+
# Log the output from the test module
42+
if result.stdout:
43+
for line in result.stdout.strip().split("\n"):
44+
logger.info(line)
45+
46+
if result.stderr:
47+
for line in result.stderr.strip().split("\n"):
48+
logger.error(line)
49+
50+
return result.returncode == 0
51+
52+
53+
def run_tests() -> List[Tuple[str, bool]]:
54+
"""Run all tests and return results."""
55+
results = []
56+
57+
for module_name in TEST_MODULES:
58+
try:
59+
logger.info(f"\n{'=' * 50}")
60+
logger.info(f"Running test: {module_name}")
61+
logger.info(f"{'-' * 50}")
62+
63+
success = run_test_module(module_name)
64+
results.append((module_name, success))
65+
66+
status = "✅ PASSED" if success else "❌ FAILED"
67+
logger.info(f"Test {module_name}: {status}")
68+
69+
except Exception as e:
70+
logger.error(f"Error loading or running test {module_name}: {str(e)}")
71+
import traceback
72+
73+
logger.error(traceback.format_exc())
74+
results.append((module_name, False))
75+
76+
return results
77+
78+
79+
def print_summary(results: List[Tuple[str, bool]]) -> None:
80+
"""Print a summary of test results."""
81+
logger.info(f"\n{'=' * 50}")
82+
logger.info("TEST SUMMARY")
83+
logger.info(f"{'-' * 50}")
84+
85+
passed = sum(1 for _, success in results if success)
86+
total = len(results)
87+
88+
for module_name, success in results:
89+
status = "✅ PASSED" if success else "❌ FAILED"
90+
logger.info(f"{status} - {module_name}")
91+
92+
logger.info(f"{'-' * 50}")
93+
logger.info(f"Total: {total} | Passed: {passed} | Failed: {total - passed}")
94+
logger.info(f"{'=' * 50}")
95+
96+
97+
if __name__ == "__main__":
98+
# Check if required environment variables are set
99+
required_vars = [
100+
"DATABRICKS_SERVER_HOSTNAME",
101+
"DATABRICKS_HTTP_PATH",
102+
"DATABRICKS_TOKEN",
103+
]
104+
missing_vars = [var for var in required_vars if not os.environ.get(var)]
105+
106+
if missing_vars:
107+
logger.error(
108+
f"Missing required environment variables: {', '.join(missing_vars)}"
47109
)
48-
49-
logger.info(f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}")
50-
logger.info(f"backend type: {type(connection.session.backend)}")
51-
52-
# Close the connection
53-
logger.info("Closing the SEA session...")
54-
connection.close()
55-
logger.info("Successfully closed SEA session")
56-
57-
except Exception as e:
58-
logger.error(f"Error testing SEA session: {str(e)}")
59-
import traceback
60-
logger.error(traceback.format_exc())
110+
logger.error("Please set these variables before running the tests.")
61111
sys.exit(1)
62-
63-
logger.info("SEA session test completed successfully")
64112

65-
if __name__ == "__main__":
66-
test_sea_session()
113+
# Run all tests
114+
results = run_tests()
115+
116+
# Print summary
117+
print_summary(results)
118+
119+
# Exit with appropriate status code
120+
all_passed = all(success for _, success in results)
121+
sys.exit(0 if all_passed else 1)

examples/experimental/tests/__init__.py

Whitespace-only changes.
Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
1+
"""
2+
Test for SEA asynchronous query execution functionality.
3+
"""
4+
import os
5+
import sys
6+
import logging
7+
import time
8+
from databricks.sql.client import Connection
9+
from databricks.sql.backend.types import CommandState
10+
11+
logging.basicConfig(level=logging.INFO)
12+
logger = logging.getLogger(__name__)
13+
14+
15+
def test_sea_async_query_with_cloud_fetch():
16+
"""
17+
Test executing a query asynchronously using the SEA backend with cloud fetch enabled.
18+
19+
This function connects to a Databricks SQL endpoint using the SEA backend,
20+
executes a simple query asynchronously with cloud fetch enabled, and verifies that execution completes successfully.
21+
"""
22+
server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
23+
http_path = os.environ.get("DATABRICKS_HTTP_PATH")
24+
access_token = os.environ.get("DATABRICKS_TOKEN")
25+
catalog = os.environ.get("DATABRICKS_CATALOG")
26+
27+
if not all([server_hostname, http_path, access_token]):
28+
logger.error("Missing required environment variables.")
29+
logger.error(
30+
"Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
31+
)
32+
return False
33+
34+
try:
35+
# Create connection with cloud fetch enabled
36+
logger.info(
37+
"Creating connection for asynchronous query execution with cloud fetch enabled"
38+
)
39+
connection = Connection(
40+
server_hostname=server_hostname,
41+
http_path=http_path,
42+
access_token=access_token,
43+
catalog=catalog,
44+
schema="default",
45+
use_sea=True,
46+
user_agent_entry="SEA-Test-Client",
47+
use_cloud_fetch=True,
48+
)
49+
50+
logger.info(
51+
f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
52+
)
53+
54+
# Execute a simple query asynchronously
55+
cursor = connection.cursor()
56+
logger.info(
57+
"Executing asynchronous query with cloud fetch: SELECT 1 as test_value"
58+
)
59+
cursor.execute_async("SELECT 1 as test_value")
60+
logger.info(
61+
"Asynchronous query submitted successfully with cloud fetch enabled"
62+
)
63+
64+
# Check query state
65+
logger.info("Checking query state...")
66+
while cursor.is_query_pending():
67+
logger.info("Query is still pending, waiting...")
68+
time.sleep(1)
69+
70+
logger.info("Query is no longer pending, getting results...")
71+
cursor.get_async_execution_result()
72+
logger.info(
73+
"Successfully retrieved asynchronous query results with cloud fetch enabled"
74+
)
75+
76+
# Close resources
77+
cursor.close()
78+
connection.close()
79+
logger.info("Successfully closed SEA session")
80+
81+
return True
82+
83+
except Exception as e:
84+
logger.error(
85+
f"Error during SEA asynchronous query execution test with cloud fetch: {str(e)}"
86+
)
87+
import traceback
88+
89+
logger.error(traceback.format_exc())
90+
return False
91+
92+
93+
def test_sea_async_query_without_cloud_fetch():
94+
"""
95+
Test executing a query asynchronously using the SEA backend with cloud fetch disabled.
96+
97+
This function connects to a Databricks SQL endpoint using the SEA backend,
98+
executes a simple query asynchronously with cloud fetch disabled, and verifies that execution completes successfully.
99+
"""
100+
server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
101+
http_path = os.environ.get("DATABRICKS_HTTP_PATH")
102+
access_token = os.environ.get("DATABRICKS_TOKEN")
103+
catalog = os.environ.get("DATABRICKS_CATALOG")
104+
105+
if not all([server_hostname, http_path, access_token]):
106+
logger.error("Missing required environment variables.")
107+
logger.error(
108+
"Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
109+
)
110+
return False
111+
112+
try:
113+
# Create connection with cloud fetch disabled
114+
logger.info(
115+
"Creating connection for asynchronous query execution with cloud fetch disabled"
116+
)
117+
connection = Connection(
118+
server_hostname=server_hostname,
119+
http_path=http_path,
120+
access_token=access_token,
121+
catalog=catalog,
122+
schema="default",
123+
use_sea=True,
124+
user_agent_entry="SEA-Test-Client",
125+
use_cloud_fetch=False,
126+
enable_query_result_lz4_compression=False,
127+
)
128+
129+
logger.info(
130+
f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
131+
)
132+
133+
# Execute a simple query asynchronously
134+
cursor = connection.cursor()
135+
logger.info(
136+
"Executing asynchronous query without cloud fetch: SELECT 1 as test_value"
137+
)
138+
cursor.execute_async("SELECT 1 as test_value")
139+
logger.info(
140+
"Asynchronous query submitted successfully with cloud fetch disabled"
141+
)
142+
143+
# Check query state
144+
logger.info("Checking query state...")
145+
while cursor.is_query_pending():
146+
logger.info("Query is still pending, waiting...")
147+
time.sleep(1)
148+
149+
logger.info("Query is no longer pending, getting results...")
150+
cursor.get_async_execution_result()
151+
logger.info(
152+
"Successfully retrieved asynchronous query results with cloud fetch disabled"
153+
)
154+
155+
# Close resources
156+
cursor.close()
157+
connection.close()
158+
logger.info("Successfully closed SEA session")
159+
160+
return True
161+
162+
except Exception as e:
163+
logger.error(
164+
f"Error during SEA asynchronous query execution test without cloud fetch: {str(e)}"
165+
)
166+
import traceback
167+
168+
logger.error(traceback.format_exc())
169+
return False
170+
171+
172+
def test_sea_async_query_exec():
173+
"""
174+
Run both asynchronous query tests and return overall success.
175+
"""
176+
with_cloud_fetch_success = test_sea_async_query_with_cloud_fetch()
177+
logger.info(
178+
f"Asynchronous query with cloud fetch: {'✅ PASSED' if with_cloud_fetch_success else '❌ FAILED'}"
179+
)
180+
181+
without_cloud_fetch_success = test_sea_async_query_without_cloud_fetch()
182+
logger.info(
183+
f"Asynchronous query without cloud fetch: {'✅ PASSED' if without_cloud_fetch_success else '❌ FAILED'}"
184+
)
185+
186+
return with_cloud_fetch_success and without_cloud_fetch_success
187+
188+
189+
if __name__ == "__main__":
190+
success = test_sea_async_query_exec()
191+
sys.exit(0 if success else 1)

0 commit comments

Comments
 (0)