atlan-application-sdk 0.1.1rc39__py3-none-any.whl → 0.1.1rc41__py3-none-any.whl

application_sdk/clients/sql.py

@@ -7,13 +7,14 @@ database operations, supporting batch processing and server-side cursors.
 
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 from urllib.parse import quote_plus
 
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
 from temporalio import activity
 
 from application_sdk.clients import ClientInterface
+from application_sdk.clients.models import DatabaseConfig
 from application_sdk.common.aws_utils import (
     generate_aws_rds_token_with_iam_role,
     generate_aws_rds_token_with_iam_user,
@@ -48,7 +49,7 @@ class BaseSQLClient(ClientInterface):
     credentials: Dict[str, Any] = {}
     resolved_credentials: Dict[str, Any] = {}
     use_server_side_cursor: bool = USE_SERVER_SIDE_CURSOR
-    DB_CONFIG: Dict[str, Any] = {}
+    DB_CONFIG: Optional[DatabaseConfig] = None
 
     def __init__(
         self,
@@ -71,25 +72,36 @@ class BaseSQLClient(ClientInterface):
         self.sql_alchemy_connect_args = sql_alchemy_connect_args
 
     async def load(self, credentials: Dict[str, Any]) -> None:
-        """Load and establish the database connection.
+        """Load credentials and prepare engine for lazy connections.
+
+        This method now only stores credentials and creates the engine without
+        establishing a persistent connection. Connections are created on-demand.
 
         Args:
             credentials (Dict[str, Any]): Database connection credentials.
 
         Raises:
-            ClientError: If connection fails due to authentication or connection issues
+            ClientError: If credentials are invalid or engine creation fails
         """
         self.credentials = credentials  # Update the instance credentials
         try:
             from sqlalchemy import create_engine
 
+            # Create engine but no persistent connection
             self.engine = create_engine(
                 self.get_sqlalchemy_connection_string(),
                 connect_args=self.sql_alchemy_connect_args,
                 pool_pre_ping=True,
             )
-            self.connection = self.engine.connect()
-        except ClientError as e:
+
+            # Test connection briefly to validate credentials
+            with self.engine.connect() as _:
+                pass  # Connection test successful
+
+            # Don't store persistent connection
+            self.connection = None
+
+        except Exception as e:
             logger.error(
                 f"{ClientError.SQL_CLIENT_AUTH_ERROR}: Error loading SQL client: {str(e)}"
             )
@@ -100,8 +112,10 @@ class BaseSQLClient(ClientInterface):
 
     async def close(self) -> None:
         """Close the database connection."""
-        if self.connection:
-            self.connection.close()
+        if self.engine:
+            self.engine.dispose()
+            self.engine = None
+        self.connection = None  # Should already be None, but ensure cleanup
 
     def get_iam_user_token(self):
         """Get an IAM user token for AWS RDS database authentication.
@@ -249,7 +263,9 @@ class BaseSQLClient(ClientInterface):
         Returns:
             str: The updated URL with the dialect.
         """
-        installed_dialect = self.DB_CONFIG["template"].split("://")[0]
+        if not self.DB_CONFIG:
+            raise ValueError("DB_CONFIG is not configured for this SQL client.")
+        installed_dialect = self.DB_CONFIG.template.split("://")[0]
         url_dialect = sqlalchemy_url.split("://")[0]
         if installed_dialect != url_dialect:
             sqlalchemy_url = sqlalchemy_url.replace(url_dialect, installed_dialect)
@@ -268,6 +284,9 @@ class BaseSQLClient(ClientInterface):
         Raises:
             ValueError: If required connection parameters are missing.
         """
+        if not self.DB_CONFIG:
+            raise ValueError("DB_CONFIG is not configured for this SQL client.")
+
         extra = parse_credentials_extra(self.credentials)
 
         # TODO: Uncomment this when the native deployment is ready
@@ -280,7 +299,7 @@ class BaseSQLClient(ClientInterface):
 
         # Prepare parameters
         param_values = {}
-        for param in self.DB_CONFIG["required"]:
+        for param in self.DB_CONFIG.required:
             if param == "password":
                 param_values[param] = auth_token
             else:
@@ -290,30 +309,28 @@ class BaseSQLClient(ClientInterface):
                 param_values[param] = value
 
         # Fill in base template
-        conn_str = self.DB_CONFIG["template"].format(**param_values)
+        conn_str = self.DB_CONFIG.template.format(**param_values)
 
         # Append defaults if not already in the template
-        if self.DB_CONFIG.get("defaults"):
-            conn_str = self.add_connection_params(conn_str, self.DB_CONFIG["defaults"])
+        if self.DB_CONFIG.defaults:
+            conn_str = self.add_connection_params(conn_str, self.DB_CONFIG.defaults)
 
-        if self.DB_CONFIG.get("parameters"):
-            parameter_keys = self.DB_CONFIG["parameters"]
-            self.DB_CONFIG["parameters"] = {
+        if self.DB_CONFIG.parameters:
+            parameter_keys = self.DB_CONFIG.parameters
+            parameter_values = {
                 key: self.credentials.get(key) or extra.get(key)
                 for key in parameter_keys
             }
-            conn_str = self.add_connection_params(
-                conn_str, self.DB_CONFIG["parameters"]
-            )
+            conn_str = self.add_connection_params(conn_str, parameter_values)
 
         return conn_str
 
     async def run_query(self, query: str, batch_size: int = 100000):
-        """Execute a SQL query and return results in batches.
+        """Execute a SQL query and return results in batches using lazy connections.
 
-        This method executes the provided SQL query and yields results in batches
-        to efficiently manage memory usage for large result sets. It supports both
-        server-side and client-side cursors based on configuration.
+        This method creates a connection on-demand, executes the query in batches,
+        and automatically closes the connection when done. This prevents memory
+        leaks from persistent connections.
 
         Args:
             query (str): SQL query to execute.
@@ -325,44 +342,47 @@ class BaseSQLClient(ClientInterface):
                 a dictionary mapping column names to values.
 
         Raises:
-            ValueError: If database connection is not established.
+            ValueError: If engine is not initialized.
             Exception: If query execution fails.
         """
-        if not self.connection:
-            raise ValueError("Connection is not established")
-        loop = asyncio.get_running_loop()
-
-        if self.use_server_side_cursor:
-            self.connection.execution_options(yield_per=batch_size)
+        if not self.engine:
+            raise ValueError("Engine is not initialized. Call load() first.")
 
+        loop = asyncio.get_running_loop()
         logger.info(f"Running query: {query}")
 
-        with ThreadPoolExecutor() as pool:
-            try:
-                from sqlalchemy import text
+        # Use context manager for automatic connection cleanup
+        with self.engine.connect() as connection:
+            if self.use_server_side_cursor:
+                connection = connection.execution_options(yield_per=batch_size)
 
-                cursor = await loop.run_in_executor(
-                    pool, self.connection.execute, text(query)
-                )
-                if not cursor or not cursor.cursor:
-                    raise ValueError("Cursor is not supported")
-                column_names: List[str] = [
-                    description.name.lower()
-                    for description in cursor.cursor.description
-                ]
+            with ThreadPoolExecutor() as pool:
+                try:
+                    from sqlalchemy import text
 
-                while True:
-                    rows = await loop.run_in_executor(
-                        pool, cursor.fetchmany, batch_size
+                    cursor = await loop.run_in_executor(
+                        pool, connection.execute, text(query)
                     )
-                    if not rows:
-                        break
-
-                    results = [dict(zip(column_names, row)) for row in rows]
-                    yield results
-            except Exception as e:
-                logger.error("Error running query in batch: {error}", error=str(e))
-                raise e
+                    if not cursor or not cursor.cursor:
+                        raise ValueError("Cursor is not supported")
+                    column_names: List[str] = [
+                        description.name.lower()
+                        for description in cursor.cursor.description
+                    ]
+
+                    while True:
+                        rows = await loop.run_in_executor(
+                            pool, cursor.fetchmany, batch_size
+                        )
+                        if not rows:
+                            break
+
+                        results = [dict(zip(column_names, row)) for row in rows]
+                        yield results
+                except Exception as e:
+                    logger.error("Error running query in batch: {error}", error=str(e))
+                    raise e
+            # Connection automatically closed by context manager
 
         logger.info("Query execution completed")
 
@@ -386,22 +406,23 @@ class AsyncBaseSQLClient(BaseSQLClient):
     engine: "AsyncEngine"
 
     async def load(self, credentials: Dict[str, Any]) -> None:
-        """Load and establish an asynchronous database connection.
+        """Load credentials and prepare async engine for lazy connections.
 
-        This method creates an async SQLAlchemy engine and establishes a connection
-        to the database using the provided credentials.
+        This method stores credentials and creates an async engine without establishing
+        a persistent connection. Connections are created on-demand for better memory efficiency.
 
         Args:
             credentials (Dict[str, Any]): Database connection credentials including
                 host, port, username, password, and other connection parameters.
 
         Raises:
-            ValueError: If connection fails due to invalid credentials or connection issues.
+            ValueError: If credentials are invalid or engine creation fails.
         """
         self.credentials = credentials
         try:
             from sqlalchemy.ext.asyncio import create_async_engine
 
+            # Create async engine but no persistent connection
             self.engine = create_async_engine(
                 self.get_sqlalchemy_connection_string(),
                 connect_args=self.sql_alchemy_connect_args,
@@ -409,7 +430,14 @@ class AsyncBaseSQLClient(BaseSQLClient):
             )
             if not self.engine:
                 raise ValueError("Failed to create async engine")
-            self.connection = await self.engine.connect()
+
+            # Test connection briefly to validate credentials
+            async with self.engine.connect() as _:
+                pass  # Connection test successful
+
+            # Don't store persistent connection
+            self.connection = None
+
         except Exception as e:
             logger.error(f"Error establishing database connection: {str(e)}")
             if self.engine:
@@ -417,11 +445,19 @@ class AsyncBaseSQLClient(BaseSQLClient):
                 self.engine = None
             raise ValueError(str(e))
 
+    async def close(self) -> None:
+        """Close the async database connection and dispose of the engine."""
+        if self.engine:
+            await self.engine.dispose()
+            self.engine = None
+        self.connection = None
+
     async def run_query(self, query: str, batch_size: int = 100000):
-        """Execute a SQL query asynchronously and return results in batches.
+        """Execute a SQL query asynchronously and return results in batches using lazy connections.
 
-        This method executes the provided SQL query using an async connection and
-        yields results in batches to manage memory usage for large result sets.
+        This method creates an async connection on-demand, executes the query in batches,
+        and automatically closes the connection when done. This prevents memory leaks
+        from persistent connections.
 
         Args:
             query (str): SQL query to execute.
@@ -433,42 +469,46 @@ class AsyncBaseSQLClient(BaseSQLClient):
                 a dictionary mapping column names to values.
 
         Raises:
+            ValueError: If engine is not initialized.
             Exception: If query execution fails.
         """
-        if not self.connection:
-            raise ValueError("Connection is not established")
+        if not self.engine:
+            raise ValueError("Engine is not initialized. Call load() first.")
 
         logger.info(f"Running query: {query}")
         use_server_side_cursor = self.use_server_side_cursor
 
-        try:
-            from sqlalchemy import text
-
-            if use_server_side_cursor:
-                await self.connection.execution_options(yield_per=batch_size)
-
-            result = (
-                await self.connection.stream(text(query))
-                if use_server_side_cursor
-                else await self.connection.execute(text(query))
-            )
+        # Use async context manager for automatic connection cleanup
+        async with self.engine.connect() as connection:
+            try:
+                from sqlalchemy import text
 
-            column_names = list(result.keys())
+                if use_server_side_cursor:
+                    connection = connection.execution_options(yield_per=batch_size)
 
-            while True:
-                rows = (
-                    await result.fetchmany(batch_size)
+                result = (
+                    await connection.stream(text(query))
                     if use_server_side_cursor
-                    else result.cursor.fetchmany(batch_size)
-                    if result.cursor
-                    else None
+                    else await connection.execute(text(query))
                 )
-                if not rows:
-                    break
-                yield [dict(zip(column_names, row)) for row in rows]
 
-        except Exception as e:
-            logger.error(f"Error executing query: {str(e)}")
-            raise
+                column_names = list(result.keys())
+
+                while True:
+                    rows = (
+                        await result.fetchmany(batch_size)
+                        if use_server_side_cursor
+                        else result.cursor.fetchmany(batch_size)
+                        if result.cursor
+                        else None
+                    )
+                    if not rows:
+                        break
+                    yield [dict(zip(column_names, row)) for row in rows]
+
+            except Exception as e:
+                logger.error(f"Error executing query: {str(e)}")
+                raise
+            # Async connection automatically closed by context manager
 
         logger.info("Query execution completed")

application_sdk/clients/temporal.py

@@ -26,6 +26,7 @@ from application_sdk.constants import (
     WORKFLOW_PORT,
     WORKFLOW_TLS_ENABLED_KEY,
 )
+from application_sdk.interceptors.cleanup import CleanupInterceptor, cleanup
 from application_sdk.interceptors.events import EventInterceptor, publish_event
 from application_sdk.interceptors.lock import RedisLockInterceptor
 from application_sdk.observability.logger_adaptor import get_logger
@@ -359,7 +360,7 @@ class TemporalWorkflowClient(WorkflowClient):
             )
 
         # Start with provided activities and add system activities
-        final_activities = list(activities) + [publish_event]
+        final_activities = list(activities) + [publish_event, cleanup]
 
         # Add lock management activities if needed
         if not IS_LOCKING_DISABLED:
@@ -395,6 +396,7 @@ class TemporalWorkflowClient(WorkflowClient):
             activity_executor=activity_executor,
             interceptors=[
                 EventInterceptor(),
+                CleanupInterceptor(),
                 RedisLockInterceptor(activities_dict),
             ],
         )