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

clickzetta_dbutils/__init__.py

@@ -1,5 +1,5 @@
 from .db_utils import get_active_engine, get_lakehouse_connection, DatabaseConnectionManager, ConnectionConfig, \
-    DatabaseConnectionError
+    DatabaseConnectionError, get_active_lakehouse_engine
 
-__all__ = ["get_active_engine", "get_lakehouse_connection", "DatabaseConnectionManager", "ConnectionConfig",
+__all__ = ["get_active_engine", "get_lakehouse_connection", "get_active_lakehouse_engine", "DatabaseConnectionManager", "ConnectionConfig",
            "DatabaseConnectionError"]

clickzetta_dbutils/db_utils.py

@@ -38,7 +38,7 @@ class DatabaseConnectionManager:
     Manages database connections with flexible configuration options.
     """
 
-    def __init__(self, ds_name: str):
+    def __init__(self, ds_name: Optional[str] = None):
         """
         Initialize a database connection for a specific data source.
         """
@@ -50,9 +50,10 @@ 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]:
+    def load_connection_configs(cls) -> Dict[str, ConnectionConfig]:
         """
         Load and cache connection configurations from environment variables.
 
@@ -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)
 
@@ -83,7 +92,7 @@ class DatabaseConnectionManager:
         """
         Find connection info by data source name
         """
-        connections = self._load_connection_configs()
+        connections = self.load_connection_configs()
 
         # Validate data source exists
         if ds_name not in connections:
@@ -95,11 +104,12 @@ class DatabaseConnectionManager:
             config.query.update(self._query)
         return config
 
-    def get_connection_infos(self):
+    @classmethod
+    def get_connection_infos(cls) -> Dict[str, ConnectionConfig]:
         """
         Get all connection infos
         """
-        return self._load_connection_configs()
+        return DatabaseConnectionManager.load_connection_configs()
 
     def use_workspace(self, workspace: str) -> 'DatabaseConnectionManager':
         """
@@ -179,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
@@ -193,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
@@ -290,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:
     """
@@ -303,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:
@@ -322,5 +347,66 @@ 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)
+
+
+def get_active_lakehouse_engine(
+        vcluster: Optional[str] = None,
+        workspace: Optional[str] = None,
+        schema: Optional[str] = None,
+        options: Optional[Dict[str, str]] = None,
+        query: Optional[Dict[str, str]] = None,
+        driver: Optional[str] = None,
+        *args, **kwargs
+) -> Engine:
+    """
+    Convenience function to create a database engine for lakehouse (ClickZetta) data source.
+    
+    Args:
+        vcluster (str, optional): Virtual cluster name for ClickZetta data source. Required.
+        workspace (str, optional): Workspace name. Default is 'default'.
+        schema (str, optional): Schema name for the connection. Default is 'public'.
+        options (dict, optional): Additional connection options.
+        query (dict, optional): Additional query parameters for SQLAlchemy url.
+        driver (str, optional): Driver name for the connection.
+        *args: Additional arguments for SQLAlchemy engine.
+        **kwargs: Additional keyword arguments for SQLAlchemy engine.
+
+    Returns:
+        SQLAlchemy Engine instance
+
+    Raises:
+        DatabaseConnectionError: If no lakehouse data source is found in the configuration.
+    """
+    # Get all connection configurations
+    conn_infos = DatabaseConnectionManager.get_connection_infos()
+
+    # Find the lakehouse (ClickZetta) data source
+    lakehouse_ds = None
+    for ds_name, conn_info in conn_infos.items():
+        if conn_info.dsType == 1:  # ClickZetta type
+            lakehouse_ds = ds_name
+            break
+
+    if not lakehouse_ds:
+        raise DatabaseConnectionError("No lakehouse (ClickZetta) data source found in configuration")
+
+    # Create engine using the found lakehouse data source
+    return get_active_engine(
+        ds_name=lakehouse_ds,
+        vcluster=vcluster,
+        workspace=workspace,
+        schema=schema,
+        options=options,
+        query=query,
+        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.7"
+__version__ = "1.0.9"

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: clickzetta-dbutils
-Version: 1.0.7
+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.7.dist-info/RECORD

@@ -1,7 +0,0 @@
-clickzetta_dbutils/__init__.py,sha256=Q_6kas0RvZ0767qlaA_xGESmXxm0dks1YQ8DqCX8LV0,290
-clickzetta_dbutils/db_utils.py,sha256=Vv00OcjoBZdC8z1ucEtekEwqPhu95GkgazBdoRBd-bI,10643
-clickzetta_dbutils/version.py,sha256=98rTj2jvgnCZhTIeHpbml7e2xuhIqt_9BJNSWkWEQb8,21
-clickzetta_dbutils-1.0.7.dist-info/METADATA,sha256=IPehhi4jMgaS7VZfqMhSwTkYkAbnLXHKf708o4537jo,852
-clickzetta_dbutils-1.0.7.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-clickzetta_dbutils-1.0.7.dist-info/top_level.txt,sha256=8o5KqMSg9pxnPNejHjMaqZV2vEDvwvsz2GdChZI0N6I,19
-clickzetta_dbutils-1.0.7.dist-info/RECORD,,