clickzetta-dbutils 1.0.6__py3-none-any.whl → 1.0.8__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.
         """
@@ -52,7 +52,7 @@ class DatabaseConnectionManager:
         self._query: Dict[str, str] = {}
 
     @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.
 
@@ -83,7 +83,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 +95,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':
         """
@@ -255,6 +256,7 @@ class DatabaseConnectionManager:
                     query=query_params
                 )
             elif conn_info.magicToken:
+                # Use magic token for authentication, do not require username and password
                 query_params["magic_token"] = conn_info.magicToken
                 url = URL.create(
                     drivername="clickzetta",
@@ -287,6 +289,8 @@ def get_active_engine(
         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:
     """
@@ -298,6 +302,9 @@ def get_active_engine(
         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.
 
     Returns:
         SQLAlchemy Engine instance
@@ -312,5 +319,63 @@ def get_active_engine(
         manager.use_vcluster(vcluster)
     if options:
         manager.use_options(options)
+    if query:
+        manager.use_query(query)
+    if driver:
+        manager.use_driver(driver)
 
     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
+    )

clickzetta_dbutils/version.py

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: clickzetta-dbutils
-Version: 1.0.6
+Version: 1.0.8
 Summary: clickzetta dbutils
 Author-email: "lin.zhang" <lin.zhang@clickzetta.com>
 Project-URL: documentation, https://www.yunqi.tech/
@@ -16,6 +16,8 @@ Provides-Extra: dev
 Requires-Dist: pytest==8.2.1; extra == "dev"
 Requires-Dist: build; extra == "dev"
 Requires-Dist: pytest-xdist; extra == "dev"
+Provides-Extra: cliclzetta-legecy-sdk
+Requires-Dist: clickzetta-connector; extra == "cliclzetta-legecy-sdk"
+Requires-Dist: clickzetta-sqlalchemy; extra == "cliclzetta-legecy-sdk"
 Provides-Extra: cliclzetta
-Requires-Dist: clickzetta-connector; extra == "cliclzetta"
-Requires-Dist: clickzetta-sqlalchemy; extra == "cliclzetta"
+Requires-Dist: clickzetta-connector-python; extra == "cliclzetta"

clickzetta_dbutils-1.0.8.dist-info/RECORD

@@ -0,0 +1,7 @@
+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,,

clickzetta_dbutils-1.0.6.dist-info/RECORD

@@ -1,7 +0,0 @@
-clickzetta_dbutils/__init__.py,sha256=Q_6kas0RvZ0767qlaA_xGESmXxm0dks1YQ8DqCX8LV0,290
-clickzetta_dbutils/db_utils.py,sha256=1iVaiYIQosukrQBAbgjNXJJW51p7k1ihQKBIcl_MvQM,10166
-clickzetta_dbutils/version.py,sha256=fCDDAyG3nMZcE_hvt1RHdxkiFN3DfNSyU_k-rLUDrpE,21
-clickzetta_dbutils-1.0.6.dist-info/METADATA,sha256=mpfK67NkOGCBN_wu5NcTvMLLTsVGlLxmnNdcX3KTiFw,726
-clickzetta_dbutils-1.0.6.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-clickzetta_dbutils-1.0.6.dist-info/top_level.txt,sha256=8o5KqMSg9pxnPNejHjMaqZV2vEDvwvsz2GdChZI0N6I,19
-clickzetta_dbutils-1.0.6.dist-info/RECORD,,