2121)
2222from .region import Region
2323from .runtime import Runtime
24- from .session import Session
24+ from .connection import Connection
2525
2626apilevel = "2.0"
2727threadsafety = 1
@@ -35,7 +35,7 @@ def connect(
3535 runtime : Runtime = DEFAULT_RUNTIME ,
3636 region : Region = DEFAULT_REGION ,
3737 wait_timeout_seconds : int = DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS ,
38- ) -> Session :
38+ ) -> Connection :
3939 if not token and not api_key :
4040 raise ValueError ("At least one of `token` or `api_key` is required" )
4141 if token and api_key :
@@ -113,11 +113,11 @@ def http_to_ws(uri: str) -> str:
113113 return str (urllib .parse .urlunparse (parsed ))
114114
115115
116- def connect_direct (uri : str , headers : dict [str , str ] = None ) -> Session :
116+ def connect_direct (uri : str , headers : dict [str , str ] = None ) -> Connection :
117117 logging .info ("Connecting to SQL session at %s ..." , uri )
118118 try :
119- connection = websockets .sync .client .connect (uri = uri , additional_headers = headers )
120- session = Session (ws = connection )
119+ ws = websockets .sync .client .connect (uri = uri , additional_headers = headers )
120+ session = Connection (ws )
121121 return session
122122 except Exception as e :
123123 raise InterfaceError ("Failed to connect to SQL session!" , e )
0 commit comments