clickzetta-dbutils 1.0.3__py3-none-any.whl → 1.0.5__py3-none-any.whl

clickzetta_dbutils/__init__.py

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

clickzetta_dbutils/db_utils.py

@@ -2,12 +2,15 @@ import json
 import os
 import urllib.parse
 from dataclasses import dataclass, field
+from logging import getLogger
 from typing import Optional, Dict
 
 from sqlalchemy import create_engine as sa_create_engine
 from sqlalchemy.engine import URL
 from sqlalchemy.engine.base import Engine
 
+_log = getLogger(__name__)
+
 
 class DatabaseConnectionError(Exception):
     """Custom exception for database connection errors."""
@@ -18,14 +21,15 @@ class DatabaseConnectionError(Exception):
 class ConnectionConfig:
     dsName: str
     dsType: int
-    schema: str
+    schema: str = "public"
+    workspaceName: str = "default"
     host: Optional[str] = None
     magicToken: Optional[str] = None
     username: Optional[str] = None
     password: Optional[str] = None
     instanceName: Optional[str] = None
-    workspaceName: Optional[str] = None
     options: Dict[str, str] = field(default_factory=dict)
+    query: Dict[str, str] = field(default_factory=dict)
 
 
 class DatabaseConnectionManager:
@@ -33,16 +37,18 @@ class DatabaseConnectionManager:
     Manages database connections with flexible configuration options.
     """
 
-    def __init__(self):
+    def __init__(self, ds_name: str):
         """
         Initialize a database connection for a specific data source.
         """
+        self._ds_name = ds_name
         self._vcluster: Optional[str] = None
         self._workspace: Optional[str] = None
         self._driver: Optional[str] = None
         self._schema: Optional[str] = None
         self._engine: Optional[Engine] = None
         self._options = {}
+        self._query: Dict[str, str] = {}
 
     @classmethod
     def _load_connection_configs(cls) -> Dict[str, ConnectionConfig]:
@@ -55,6 +61,9 @@ class DatabaseConnectionManager:
         if not hasattr(DatabaseConnectionManager, '_connection_cache'):
             # Retrieve and decode connection info from environment variable
             conn_info_str = os.environ.get('connectionInfos', '[]')
+            if not conn_info_str:
+                raise DatabaseConnectionError(
+                    "No connection information found in environment variable 'connectionInfos'")
             decoded_info = urllib.parse.unquote(conn_info_str)
             conn_list = json.loads(decoded_info)
 
@@ -63,6 +72,10 @@ class DatabaseConnectionManager:
                 info.get('dsName'): ConnectionConfig(**info)
                 for info in conn_list
             }
+
+            resource_names = [f"{conn.dsName} ({conn.workspaceName})" for conn in cls._connection_cache.values()]
+            _log.info(f"Successfully loaded connection configurations: {', '.join(resource_names)}")
+
         return cls._connection_cache
 
     def get_connection_info(self, ds_name: str) -> ConnectionConfig:
@@ -77,8 +90,16 @@ class DatabaseConnectionManager:
 
         config = connections.get(ds_name)
         config.options.update(self._options)
+        if self._query:
+            config.query.update(self._query)
         return config
 
+    def get_connection_infos(self):
+        """
+        Get all connection infos
+        """
+        return self._load_connection_configs()
+
     def use_workspace(self, workspace: str) -> 'DatabaseConnectionManager':
         """
         Set workspace for the connection.
@@ -143,15 +164,27 @@ class DatabaseConnectionManager:
         """
         if options:
             self._options.update(options)
+        return self
+
+    def use_query(self, query: dict) -> 'DatabaseConnectionManager':
+        """
+        Set query for the connection.
+        Args:
+            query (str): Query string
+        Returns:
+            self: For method chaining
+        """
+        if query:
+            self._query.update(query)
+        return self
 
-    def connect(self, ds_name: str, *args, **kwargs) -> Engine:
+    def build(self, *args, **kwargs) -> Engine:
         """
         Create SQLAlchemy engine based on data source name and optional schema
 
-        :param ds_name: Name of the data source
         :return: SQLAlchemy Engine
         """
-        conn_info: ConnectionConfig = self.get_connection_info(ds_name)
+        conn_info: ConnectionConfig = self.get_connection_info(self._ds_name)
 
         if not conn_info.host:
             raise DatabaseConnectionError("Missing connection host for MySQL data source")
@@ -160,15 +193,14 @@ class DatabaseConnectionManager:
         options = conn_info.options or {}
         schema = self._schema or conn_info.schema
         host_parts = conn_info.host.split(':')
-
-
+        connect_args = {}
 
         # Construct connection URL based on data source type
         if ds_type == 5:  # Mysql
             if not conn_info.username or not conn_info.password:
                 raise DatabaseConnectionError("Missing username or password for MySQL data source")
-            # Split host into host and port if provided
 
+            options.update(conn_info.query)
             url = URL.create(
                 drivername=self._driver or 'mysql+mysqlconnector',
                 username=conn_info.username,
@@ -187,8 +219,10 @@ class DatabaseConnectionManager:
                 password=conn_info.password,
                 host=host_parts[0],
                 port=host_parts[1] if len(host_parts) > 1 else None,
-                database=schema
+                database=schema,
+                query=conn_info.query
             )
+            connect_args = {'options': self._convert_options(options)}
         elif ds_type == 1:  # ClickZetta
             if not conn_info.workspaceName or not conn_info.instanceName:
                 raise DatabaseConnectionError("Missing required parameters 'workspace_name', "
@@ -211,16 +245,22 @@ class DatabaseConnectionManager:
             else:
                 raise ValueError("username and password or token must be specified")
 
-
             # Add schema if provided
             if schema:
                 base_url += f"&schema={schema}"
 
+            options.update(conn_info.query)
+
+            # Add query string if provided
+            if options:
+                base_url += f"&{urllib.parse.urlencode(options)}"
+
             url = base_url
         else:
             raise ValueError(f"Unsupported data source type: {ds_type}")
 
-        return sa_create_engine(url, connect_args={'options': self._convert_options(options)}, *args, **kwargs)
+        self._engine = sa_create_engine(url, connect_args=connect_args, *args, **kwargs)
+        return self._engine
 
     @staticmethod
     def _convert_options(options):
@@ -229,8 +269,8 @@ class DatabaseConnectionManager:
         return ' '.join([f'-c {k}={v}' for k, v in options.items()])
 
 
-def get_lakehouse_client(conn):
-    return conn.connection.connection._client
+def get_lakehouse_connection(conn):
+    return conn.connection.connection
 
 
 def get_active_engine(
@@ -245,15 +285,16 @@ def get_active_engine(
     Convenience function to create a database engine.
 
     Args:
-        ds_name (str): Data source name
-        workspace (str, optional): Workspace name
-        schema (str, optional): Schema name
-        vcluster (str, optional): Virtual cluster name
+        ds_name (str): Data source name. Required.
+        vcluster (str, optional): Virtual cluster name for ClickZetta data source. Required for ClickZetta.
+        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.
 
     Returns:
         SQLAlchemy Engine instance
     """
-    manager = DatabaseConnectionManager()
+    manager = DatabaseConnectionManager(ds_name)
 
     if workspace:
         manager.use_workspace(workspace)
@@ -264,4 +305,4 @@ def get_active_engine(
     if options:
         manager.use_options(options)
 
-    return manager.connect(ds_name, *args, **kwargs)
+    return manager.build(*args, **kwargs)

clickzetta_dbutils/version.py

@@ -1 +1 @@
-__version__ = "1.0.3"
+__version__ = "1.0.5"

{clickzetta_dbutils-1.0.3.dist-info → clickzetta_dbutils-1.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: clickzetta-dbutils
-Version: 1.0.3
+Version: 1.0.5
 Summary: clickzetta dbutils
 Author-email: "lin.zhang" <lin.zhang@clickzetta.com>
 Project-URL: documentation, https://www.yunqi.tech/
@@ -10,7 +10,7 @@ Classifier: Development Status :: 3 - Alpha
 Requires-Python: >=3.7
 Requires-Dist: sqlalchemy<2.0.0,>=1.4.0
 Requires-Dist: numpy==1.24.4
-Requires-Dist: psycopg2
+Requires-Dist: psycopg2-binary
 Requires-Dist: mysql-connector-python
 Provides-Extra: dev
 Requires-Dist: pytest==8.2.1; extra == "dev"

clickzetta_dbutils-1.0.5.dist-info/RECORD

@@ -0,0 +1,7 @@
+clickzetta_dbutils/__init__.py,sha256=Q_6kas0RvZ0767qlaA_xGESmXxm0dks1YQ8DqCX8LV0,290
+clickzetta_dbutils/db_utils.py,sha256=TQobsTlxNTzUwIFgWZvVMtydNUPk_7AG82ChWQ4BCQY,10085
+clickzetta_dbutils/version.py,sha256=ZR1VA9cGs0vIK6cWK4YKLfBTnmUCAcDaaP9ARPPYxEs,21
+clickzetta_dbutils-1.0.5.dist-info/METADATA,sha256=2QYdvLVx4oQCI0J727G0ACdKrwuNAlYXjPpUKuSS6mI,938
+clickzetta_dbutils-1.0.5.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+clickzetta_dbutils-1.0.5.dist-info/top_level.txt,sha256=8o5KqMSg9pxnPNejHjMaqZV2vEDvwvsz2GdChZI0N6I,19
+clickzetta_dbutils-1.0.5.dist-info/RECORD,,

clickzetta_dbutils-1.0.3.dist-info/RECORD

@@ -1,7 +0,0 @@
-clickzetta_dbutils/__init__.py,sha256=OevYNnzvgLUw0-KDJKG7loCONkN_7DxtmqZAmzCHGAg,282
-clickzetta_dbutils/db_utils.py,sha256=L170MOzkMH4cfEwYJcHPz7_aQcTi96e9BqkF6CCRfc4,8506
-clickzetta_dbutils/version.py,sha256=-PYTKPEkq2aGJHOdzkgu5oKZIPKSq6MCqGrFyy4QQXc,21
-clickzetta_dbutils-1.0.3.dist-info/METADATA,sha256=Io3EEH8-lVA9DMp33N7x8WDgMPeeWHkzY9Qi4GDwPg4,931
-clickzetta_dbutils-1.0.3.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-clickzetta_dbutils-1.0.3.dist-info/top_level.txt,sha256=8o5KqMSg9pxnPNejHjMaqZV2vEDvwvsz2GdChZI0N6I,19
-clickzetta_dbutils-1.0.3.dist-info/RECORD,,