clickzetta-dbutils 1.0.8__py3-none-any.whl → 1.0.9__py3-none-any.whl

clickzetta_dbutils/db_utils.py

@@ -50,6 +50,7 @@ class DatabaseConnectionManager:
         self._engine: Optional[Engine] = None
         self._options = {}
         self._query: Dict[str, str] = {}
+        self._host: Optional[str] = None
 
     @classmethod
     def load_connection_configs(cls) -> Dict[str, ConnectionConfig]:
@@ -61,10 +62,18 @@ class DatabaseConnectionManager:
         """
         if not hasattr(DatabaseConnectionManager, '_connection_cache'):
             # Retrieve and decode connection info from environment variable
-            conn_info_str = os.environ.get('connectionInfos', '[]')
+            conn_info_str = ""
+            if name := os.environ.get('EXECUTE_LOG_ID', ''):
+                pipe_path = '/tmp/' + name
+                try:
+                    with open(pipe_path, 'r') as f:
+                        conn_info_str = f.read()
+                except FileNotFoundError:
+                    pass
+            # Fall back to environment variable if pipe is not found
+            conn_info_str = conn_info_str or os.environ.get('connectionInfos', '[]')
             if not conn_info_str:
-                raise DatabaseConnectionError(
-                    "No connection information found in environment variable 'connectionInfos'")
+                raise DatabaseConnectionError("No connection information found in env")
             decoded_info = urllib.parse.unquote(conn_info_str)
             conn_list = json.loads(decoded_info)
 
@@ -180,6 +189,18 @@ class DatabaseConnectionManager:
             self._query.update(query)
         return self
 
+    def use_host(self, host: str) -> 'DatabaseConnectionManager':
+        """
+        Set host for the connection.
+        Args:
+            host (str): Host name
+        Returns:
+            self: For method chaining
+        """
+        if host:
+            self._host = host
+        return self
+
     def build(self, *args, **kwargs) -> Engine:
         """
         Create SQLAlchemy engine based on data source name and optional schema
@@ -194,7 +215,8 @@ class DatabaseConnectionManager:
         ds_type = conn_info.dsType
         options = conn_info.options or {}
         schema = self._schema or conn_info.schema
-        host_parts = conn_info.host.split(':')
+        host = self._host or conn_info.host
+        host_parts = host.split(':')
         connect_args = {}
 
         # Construct connection URL based on data source type
@@ -291,6 +313,7 @@ def get_active_engine(
         options: Optional[Dict[str, str]] = None,
         query: Optional[Dict[str, str]] = None,
         driver: Optional[str] = None,
+        host: Optional[str] = None,
         *args, **kwargs
 ) -> Engine:
     """
@@ -304,6 +327,7 @@ def get_active_engine(
         options (dict, optional): Additional connection options.
         query (dict, optional): Additional query parameters for SQLAlchemy url.
         driver (str, optional): Driver name for the connection.
+        host (str, optional): Host name for the connection.
         *args: Additional arguments for SQLAlchemy engine.
 
     Returns:
@@ -323,6 +347,8 @@ def get_active_engine(
         manager.use_query(query)
     if driver:
         manager.use_driver(driver)
+    if host:
+        manager.use_host(host)
 
     return manager.build(*args, **kwargs)
 
@@ -379,3 +405,8 @@ def get_active_lakehouse_engine(
         driver=driver,
         *args, **kwargs
     )
+
+def clean_connection_cache():
+    """Clear connection cache before each test"""
+    if hasattr(DatabaseConnectionManager, '_connection_cache'):
+        delattr(DatabaseConnectionManager, '_connection_cache')

clickzetta_dbutils/version.py

@@ -1 +1 @@
-__version__ = "1.0.8"
+__version__ = "1.0.9"

{clickzetta_dbutils-1.0.8.dist-info → clickzetta_dbutils-1.0.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: clickzetta-dbutils
-Version: 1.0.8
+Version: 1.0.9
 Summary: clickzetta dbutils
 Author-email: "lin.zhang" <lin.zhang@clickzetta.com>
 Project-URL: documentation, https://www.yunqi.tech/

clickzetta_dbutils-1.0.9.dist-info/RECORD

@@ -0,0 +1,7 @@
+clickzetta_dbutils/__init__.py,sha256=XSofC7BRKAZbgXkil1HFUiHZztzWWyHcbKNi5tpOjCk,350
+clickzetta_dbutils/db_utils.py,sha256=swBzlWMbwNRcEQiehXeOoWZGibaOkm-JO4pRZdAZxgM,13756
+clickzetta_dbutils/version.py,sha256=lFohcdMk8DViWSfrTasSRt9RPcMXfTvmskn8IZQNlws,21
+clickzetta_dbutils-1.0.9.dist-info/METADATA,sha256=wCpqpipf2S5eaHf0RVEI5CzGUGqZXe6k4N4Ya5AyT48,852
+clickzetta_dbutils-1.0.9.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+clickzetta_dbutils-1.0.9.dist-info/top_level.txt,sha256=8o5KqMSg9pxnPNejHjMaqZV2vEDvwvsz2GdChZI0N6I,19
+clickzetta_dbutils-1.0.9.dist-info/RECORD,,

clickzetta_dbutils-1.0.8.dist-info/RECORD

@@ -1,7 +0,0 @@
-clickzetta_dbutils/__init__.py,sha256=XSofC7BRKAZbgXkil1HFUiHZztzWWyHcbKNi5tpOjCk,350
-clickzetta_dbutils/db_utils.py,sha256=AHtyOD-Kk980CtPg6mKIRhW1M7gG9j68-Y7TBJvgF0Q,12681
-clickzetta_dbutils/version.py,sha256=p_ws-9W7HeP_RfNuGscH3UIczaDFgLeGOeDgRBhdzrc,21
-clickzetta_dbutils-1.0.8.dist-info/METADATA,sha256=_sO2JXZFTe5AFT0cPa2Uuqz_v-mlRo0qDay1MZ0IOPs,852
-clickzetta_dbutils-1.0.8.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-clickzetta_dbutils-1.0.8.dist-info/top_level.txt,sha256=8o5KqMSg9pxnPNejHjMaqZV2vEDvwvsz2GdChZI0N6I,19
-clickzetta_dbutils-1.0.8.dist-info/RECORD,,