Skip to content

Commit 05ee4e7

Browse files
varun-edachali-dbxvarun-edachali-dbx
committed
add test scripts
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent e8e8ee7 commit 05ee4e7

File tree

5 files changed

+521
-0
lines changed

5 files changed

+521
-0
lines changed

‎examples/experimental/tests/__init__.py

Whitespace-only changes.

‎examples/experimental/tests/test_sea_async_query.py

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)

‎examples/experimental/tests/test_sea_metadata.py

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
"""
2+
Test for SEA metadata functionality.
3+
"""
4+
import os
5+
import sys
6+
import logging
7+
from databricks.sql.client import Connection
8+
9+
logging.basicConfig(level=logging.INFO)
10+
logger = logging.getLogger(__name__)
11+
12+
13+
def test_sea_metadata():
14+
"""
15+
Test metadata operations using the SEA backend.
16+
17+
This function connects to a Databricks SQL endpoint using the SEA backend,
18+
and executes metadata operations like catalogs(), schemas(), tables(), and columns().
19+
"""
20+
server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
21+
http_path = os.environ.get("DATABRICKS_HTTP_PATH")
22+
access_token = os.environ.get("DATABRICKS_TOKEN")
23+
catalog = os.environ.get("DATABRICKS_CATALOG")
24+
25+
if not all([server_hostname, http_path, access_token]):
26+
logger.error("Missing required environment variables.")
27+
logger.error(
28+
"Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
29+
)
30+
return False
31+
32+
if not catalog:
33+
logger.error(
34+
"DATABRICKS_CATALOG environment variable is required for metadata tests."
35+
)
36+
return False
37+
38+
try:
39+
# Create connection
40+
logger.info("Creating connection for metadata operations")
41+
connection = Connection(
42+
server_hostname=server_hostname,
43+
http_path=http_path,
44+
access_token=access_token,
45+
catalog=catalog,
46+
schema="default",
47+
use_sea=True,
48+
user_agent_entry="SEA-Test-Client",
49+
)
50+
51+
logger.info(
52+
f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
53+
)
54+
55+
# Test catalogs
56+
cursor = connection.cursor()
57+
logger.info("Fetching catalogs...")
58+
cursor.catalogs()
59+
logger.info("Successfully fetched catalogs")
60+
61+
# Test schemas
62+
logger.info(f"Fetching schemas for catalog '{catalog}'...")
63+
cursor.schemas(catalog_name=catalog)
64+
logger.info("Successfully fetched schemas")
65+
66+
# Test tables
67+
logger.info(f"Fetching tables for catalog '{catalog}', schema 'default'...")
68+
cursor.tables(catalog_name=catalog, schema_name="default")
69+
logger.info("Successfully fetched tables")
70+
71+
# Test columns for a specific table
72+
# Using a common table that should exist in most environments
73+
logger.info(
74+
f"Fetching columns for catalog '{catalog}', schema 'default', table 'customer'..."
75+
)
76+
cursor.columns(
77+
catalog_name=catalog, schema_name="default", table_name="customer"
78+
)
79+
logger.info("Successfully fetched columns")
80+
81+
# Close resources
82+
cursor.close()
83+
connection.close()
84+
logger.info("Successfully closed SEA session")
85+
86+
return True
87+
88+
except Exception as e:
89+
logger.error(f"Error during SEA metadata test: {str(e)}")
90+
import traceback
91+
92+
logger.error(traceback.format_exc())
93+
return False
94+
95+
96+
if __name__ == "__main__":
97+
success = test_sea_metadata()
98+
sys.exit(0 if success else 1)

‎examples/experimental/tests/test_sea_session.py

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
"""
2+
Test for SEA session management functionality.
3+
"""
4+
import os
5+
import sys
6+
import logging
7+
from databricks.sql.client import Connection
8+
9+
logging.basicConfig(level=logging.INFO)
10+
logger = logging.getLogger(__name__)
11+
12+
13+
def test_sea_session():
14+
"""
15+
Test opening and closing a SEA session using the connector.
16+
17+
This function connects to a Databricks SQL endpoint using the SEA backend,
18+
opens a session, and then closes it.
19+
20+
Required environment variables:
21+
- DATABRICKS_SERVER_HOSTNAME: Databricks server hostname
22+
- DATABRICKS_HTTP_PATH: HTTP path for the SQL endpoint
23+
- DATABRICKS_TOKEN: Personal access token for authentication
24+
"""
25+
server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
26+
http_path = os.environ.get("DATABRICKS_HTTP_PATH")
27+
access_token = os.environ.get("DATABRICKS_TOKEN")
28+
catalog = os.environ.get("DATABRICKS_CATALOG")
29+
30+
if not all([server_hostname, http_path, access_token]):
31+
logger.error("Missing required environment variables.")
32+
logger.error(
33+
"Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN."
34+
)
35+
return False
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",
47+
)
48+
49+
logger.info(
50+
f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}"
51+
)
52+
logger.info(f"Backend type: {type(connection.session.backend)}")
53+
54+
# Close the connection
55+
logger.info("Closing the SEA session...")
56+
connection.close()
57+
logger.info("Successfully closed SEA session")
58+
59+
return True
60+
61+
except Exception as e:
62+
logger.error(f"Error testing SEA session: {str(e)}")
63+
import traceback
64+
65+
logger.error(traceback.format_exc())
66+
return False
67+
68+
69+
if __name__ == "__main__":
70+
success = test_sea_session()
71+
sys.exit(0 if success else 1)

0 commit comments

Comments
 (0)