Skip to content

Commit 0d11c28

Browse files
[squashed from prev branch] introduce sea client with session open and close functionality
Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>
1 parent 3c78ed7 commit 0d11c28

File tree

9 files changed

+763
-18
lines changed

9 files changed

+763
-18
lines changed

.github/workflows/code-quality-checks.yml

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,15 @@
11
name: Code Quality Checks
2-
3-
on: [pull_request]
4-
2+
on:
3+
push:
4+
branches:
5+
- main
6+
- sea-migration
7+
- telemetry
8+
pull_request:
9+
branches:
10+
- main
11+
- sea-migration
12+
- telemetry
513
jobs:
614
run-unit-tests:
715
runs-on: ubuntu-latest

.github/workflows/integration.yml

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
name: Integration Tests
2-
32
on:
4-
push:
3+
push:
4+
paths-ignore:
5+
- "**.MD"
6+
- "**.md"
7+
pull_request:
58
branches:
69
- main
7-
pull_request:
10+
- sea-migration
11+
- telemetry
812

913
jobs:
1014
run-e2e-tests:
Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
import os
2+
import sys
3+
import logging
4+
from databricks.sql.client import Connection
5+
6+
logging.basicConfig(level=logging.DEBUG)
7+
logger = logging.getLogger(__name__)
8+
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+
server_hostname = os.environ.get("DATABRICKS_SERVER_HOSTNAME")
22+
http_path = os.environ.get("DATABRICKS_HTTP_PATH")
23+
access_token = os.environ.get("DATABRICKS_TOKEN")
24+
catalog = os.environ.get("DATABRICKS_CATALOG")
25+
26+
if not all([server_hostname, http_path, access_token]):
27+
logger.error("Missing required environment variables.")
28+
logger.error("Please set DATABRICKS_SERVER_HOSTNAME, DATABRICKS_HTTP_PATH, and DATABRICKS_TOKEN.")
29+
sys.exit(1)
30+
31+
logger.info(f"Connecting to {server_hostname}")
32+
logger.info(f"HTTP Path: {http_path}")
33+
if catalog:
34+
logger.info(f"Using catalog: {catalog}")
35+
36+
try:
37+
logger.info("Creating connection with SEA backend...")
38+
connection = Connection(
39+
server_hostname=server_hostname,
40+
http_path=http_path,
41+
access_token=access_token,
42+
catalog=catalog,
43+
schema="default",
44+
use_sea=True,
45+
user_agent_entry="SEA-Test-Client" # add custom user agent
46+
)
47+
48+
logger.info(f"Successfully opened SEA session with ID: {connection.get_session_id_hex()}")
49+
logger.info(f"backend type: {type(connection.session.backend)}")
50+
51+
# Close the connection
52+
logger.info("Closing the SEA session...")
53+
connection.close()
54+
logger.info("Successfully closed SEA session")
55+
56+
except Exception as e:
57+
logger.error(f"Error testing SEA session: {str(e)}")
58+
import traceback
59+
logger.error(traceback.format_exc())
60+
sys.exit(1)
61+
62+
logger.info("SEA session test completed successfully")
63+
64+
if __name__ == "__main__":
65+
test_sea_session()

0 commit comments

Comments
 (0)