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