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