atlan-application-sdk 1.1.1__py3-none-any.whl → 2.1.0__py3-none-any.whl

application_sdk/activities/query_extraction/sql.py

@@ -16,9 +16,8 @@ from application_sdk.clients.sql import BaseSQLClient
 from application_sdk.constants import UPSTREAM_OBJECT_STORE_NAME
 from application_sdk.handlers import HandlerInterface
 from application_sdk.handlers.sql import BaseSQLHandler
-from application_sdk.inputs.sql_query import SQLQueryInput
+from application_sdk.io.parquet import ParquetFileWriter
 from application_sdk.observability.logger_adaptor import get_logger
-from application_sdk.outputs.parquet import ParquetOutput
 from application_sdk.services.objectstore import ObjectStore
 from application_sdk.services.secretstore import SecretStore
 from application_sdk.transformers import TransformerInterface
@@ -202,21 +201,23 @@ class SQLQueryExtractionActivities(ActivitiesInterface):
 
         try:
             state = await self._get_state(workflow_args)
-            sql_input = SQLQueryInput(
-                engine=state.sql_client.engine,
-                query=self.get_formatted_query(self.fetch_queries_sql, workflow_args),
-                chunk_size=None,
+            sql_client = state.sql_client
+            if not sql_client:
+                logger.error("SQL client not initialized")
+                raise ValueError("SQL client not initialized")
+
+            formatted_query = self.get_formatted_query(
+                self.fetch_queries_sql, workflow_args
             )
-            sql_input = await sql_input.get_dataframe()
+            sql_results = await sql_client.get_results(formatted_query)
 
-            raw_output = ParquetOutput(
-                output_path=workflow_args["output_path"],
-                output_suffix="raw/query",
+            raw_output = ParquetFileWriter(
+                path=os.path.join(workflow_args["output_path"], "raw/query"),
                 chunk_size=workflow_args["miner_args"].get("chunk_size", 100000),
                 start_marker=workflow_args["start_marker"],
                 end_marker=workflow_args["end_marker"],
             )
-            await raw_output.write_dataframe(sql_input)
+            await raw_output.write(sql_results)
             logger.info(
                 f"Query fetch completed, {raw_output.total_record_count} records processed",
             )

application_sdk/application/__init__.py

@@ -5,8 +5,8 @@ from application_sdk.activities import ActivitiesInterface
 from application_sdk.clients.base import BaseClient
 from application_sdk.clients.utils import get_workflow_client
 from application_sdk.constants import ENABLE_MCP
-from application_sdk.events.models import EventRegistration
 from application_sdk.handlers.base import BaseHandler
+from application_sdk.interceptors.models import EventRegistration
 from application_sdk.observability.logger_adaptor import get_logger
 from application_sdk.server import ServerInterface
 from application_sdk.server.fastapi import APIServer, HttpWorkflowTrigger

application_sdk/clients/sql.py

@@ -6,8 +6,19 @@ database operations, supporting batch processing and server-side cursors.
 """
 
 import asyncio
+import concurrent
 from concurrent.futures import ThreadPoolExecutor
-from typing import Any, Dict, List, Optional
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncIterator,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Union,
+    cast,
+)
 from urllib.parse import quote_plus
 
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
@@ -27,6 +38,11 @@ from application_sdk.observability.logger_adaptor import get_logger
 logger = get_logger(__name__)
 activity.logger = logger
 
+if TYPE_CHECKING:
+    import daft
+    import pandas as pd
+    from sqlalchemy.orm import Session
+
 
 class BaseSQLClient(ClientInterface):
     """SQL client for database operations.
@@ -53,6 +69,7 @@ class BaseSQLClient(ClientInterface):
         self,
         use_server_side_cursor: bool = USE_SERVER_SIDE_CURSOR,
         credentials: Dict[str, Any] = {},
+        chunk_size: int = 5000,
     ):
         """
         Initialize the SQL client.
@@ -64,6 +81,7 @@ class BaseSQLClient(ClientInterface):
         """
         self.use_server_side_cursor = use_server_side_cursor
         self.credentials = credentials
+        self.chunk_size = chunk_size
 
     async def load(self, credentials: Dict[str, Any]) -> None:
         """Load credentials and prepare engine for lazy connections.
@@ -383,6 +401,154 @@ class BaseSQLClient(ClientInterface):
 
         logger.info("Query execution completed")
 
+    def _execute_pandas_query(
+        self, conn, query, chunksize: Optional[int]
+    ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
+        """Helper function to execute SQL query using pandas.
+           The function is responsible for using import_optional_dependency method of the pandas library to import sqlalchemy
+           This function helps pandas in determining weather to use the sqlalchemy connection object and constructs like text()
+           or use the underlying database connection object. This has been done to make sure connectors like the Redshift connector,
+           which do not support the sqlalchemy connection object, can be made compatible with the application-sdk.
+
+        Args:
+            conn: Database connection object.
+
+        Returns:
+            Union["pd.DataFrame", Iterator["pd.DataFrame"]]: Query results as DataFrame
+                or iterator of DataFrames if chunked.
+        """
+        import pandas as pd
+        from pandas.compat._optional import import_optional_dependency
+        from sqlalchemy import text
+
+        if import_optional_dependency("sqlalchemy", errors="ignore"):
+            return pd.read_sql_query(text(query), conn, chunksize=chunksize)
+        else:
+            dbapi_conn = getattr(conn, "connection", None)
+            return pd.read_sql_query(query, dbapi_conn, chunksize=chunksize)
+
+    def _read_sql_query(
+        self, session: "Session", query: str, chunksize: Optional[int]
+    ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
+        """Execute SQL query using the provided session.
+
+        Args:
+            session: SQLAlchemy session for database operations.
+
+        Returns:
+            Union["pd.DataFrame", Iterator["pd.DataFrame"]]: Query results as DataFrame
+                or iterator of DataFrames if chunked.
+        """
+        conn = session.connection()
+        return self._execute_pandas_query(conn, query, chunksize=chunksize)
+
+    def _execute_query_daft(
+        self, query: str, chunksize: Optional[int]
+    ) -> Union["daft.DataFrame", Iterator["daft.DataFrame"]]:
+        """Execute SQL query using the provided engine and daft.
+
+        Returns:
+            Union["daft.DataFrame", Iterator["daft.DataFrame"]]: Query results as DataFrame
+                or iterator of DataFrames if chunked.
+        """
+        # Daft uses ConnectorX to read data from SQL by default for supported connectors
+        # If a connection string is passed, it will use ConnectorX to read data
+        # For unsupported connectors and if directly engine is passed, it will use SQLAlchemy
+        import daft
+
+        if not self.engine:
+            raise ValueError("Engine is not initialized. Call load() first.")
+
+        if isinstance(self.engine, str):
+            return daft.read_sql(query, self.engine, infer_schema_length=chunksize)
+        return daft.read_sql(query, self.engine.connect, infer_schema_length=chunksize)
+
+    def _execute_query(
+        self, query: str, chunksize: Optional[int]
+    ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
+        """Execute SQL query using the provided engine and pandas.
+
+        Returns:
+            Union["pd.DataFrame", Iterator["pd.DataFrame"]]: Query results as DataFrame
+                or iterator of DataFrames if chunked.
+        """
+        if not self.engine:
+            raise ValueError("Engine is not initialized. Call load() first.")
+
+        with self.engine.connect() as conn:
+            return self._execute_pandas_query(conn, query, chunksize)
+
+    async def _execute_async_read_operation(
+        self, query: str, chunksize: Optional[int]
+    ) -> Union["pd.DataFrame", Iterator["pd.DataFrame"]]:
+        """Helper to execute async read operation with either async session or thread executor."""
+        if isinstance(self.engine, str):
+            raise ValueError("Engine should be an SQLAlchemy engine object")
+
+        from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
+
+        async_session = None
+        if self.engine and isinstance(self.engine, AsyncEngine):
+            from sqlalchemy.orm import sessionmaker
+
+            async_session = sessionmaker(
+                self.engine, expire_on_commit=False, class_=AsyncSession
+            )
+
+        if async_session:
+            async with async_session() as session:
+                return await session.run_sync(
+                    self._read_sql_query, query, chunksize=chunksize
+                )
+        else:
+            # Run the blocking operation in a thread pool
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                return await asyncio.get_event_loop().run_in_executor(
+                    executor, self._execute_query, query, chunksize
+                )
+
+    async def get_batched_results(
+        self,
+        query: str,
+    ) -> Union[AsyncIterator["pd.DataFrame"], Iterator["pd.DataFrame"]]:  # type: ignore
+        """Get query results as batched pandas DataFrames asynchronously.
+
+        Returns:
+            AsyncIterator["pd.DataFrame"]: Async iterator yielding batches of query results.
+
+        Raises:
+            ValueError: If engine is a string instead of SQLAlchemy engine.
+            Exception: If there's an error executing the query.
+        """
+        try:
+            # We cast to Iterator because passing chunk_size guarantees an Iterator return
+            result = await self._execute_async_read_operation(query, self.chunk_size)
+            return cast(Iterator["pd.DataFrame"], result)
+        except Exception as e:
+            logger.error(f"Error reading batched data(pandas) from SQL: {str(e)}")
+
+    async def get_results(self, query: str) -> "pd.DataFrame":
+        """Get all query results as a single pandas DataFrame asynchronously.
+
+        Returns:
+            pd.DataFrame: Query results as a DataFrame.
+
+        Raises:
+            ValueError: If engine is a string instead of SQLAlchemy engine.
+            Exception: If there's an error executing the query.
+        """
+        try:
+            result = await self._execute_async_read_operation(query, None)
+            import pandas as pd
+
+            if isinstance(result, pd.DataFrame):
+                return result
+            raise Exception("Unable to get pandas dataframe from SQL query results")
+
+        except Exception as e:
+            logger.error(f"Error reading data(pandas) from SQL: {str(e)}")
+            raise e
+
 
 class AsyncBaseSQLClient(BaseSQLClient):
     """Asynchronous SQL client for database operations.

application_sdk/clients/temporal.py

@@ -26,18 +26,18 @@ from application_sdk.constants import (
     WORKFLOW_PORT,
     WORKFLOW_TLS_ENABLED,
 )
-from application_sdk.events.models import (
-    ApplicationEventNames,
-    Event,
-    EventTypes,
-    WorkerTokenRefreshEventData,
-)
 from application_sdk.interceptors.cleanup import CleanupInterceptor, cleanup
 from application_sdk.interceptors.correlation_context import (
     CorrelationContextInterceptor,
 )
 from application_sdk.interceptors.events import EventInterceptor, publish_event
 from application_sdk.interceptors.lock import RedisLockInterceptor
+from application_sdk.interceptors.models import (
+    ApplicationEventNames,
+    Event,
+    EventTypes,
+    WorkerTokenRefreshEventData,
+)
 from application_sdk.observability.logger_adaptor import get_logger
 from application_sdk.services.eventstore import EventStore
 from application_sdk.services.secretstore import SecretStore

application_sdk/common/types.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class DataframeType(Enum):
+    """Enumeration of dataframe types."""
+
+    pandas = "pandas"
+    daft = "daft"

application_sdk/common/utils.py

@@ -20,7 +20,6 @@ from typing import (
 from application_sdk.activities.common.utils import get_object_store_prefix
 from application_sdk.common.error_codes import CommonError
 from application_sdk.constants import TEMPORARY_PATH
-from application_sdk.inputs.sql_query import SQLQueryInput
 from application_sdk.observability.logger_adaptor import get_logger
 from application_sdk.services.objectstore import ObjectStore
 
@@ -280,13 +279,7 @@ async def get_database_names(
             temp_table_regex_sql=temp_table_regex_sql,
             use_posix_regex=True,
         )
-        # We'll run the query to get all the database names
-        database_sql_input = SQLQueryInput(
-            engine=sql_client.engine,
-            query=prepared_query,  # type: ignore
-            chunk_size=None,
-        )
-        database_dataframe = await database_sql_input.get_dataframe()
+        database_dataframe = await sql_client.get_results(prepared_query)
         database_names = list(database_dataframe["database_name"])
     return database_names
 

application_sdk/constants.py

@@ -244,7 +244,7 @@ TRACES_FILE_NAME = "traces.parquet"
 
 # Dapr Sink Configuration
 ENABLE_OBSERVABILITY_DAPR_SINK = (
-    os.getenv("ATLAN_ENABLE_OBSERVABILITY_DAPR_SINK", "false").lower() == "true"
+    os.getenv("ATLAN_ENABLE_OBSERVABILITY_DAPR_SINK", "true").lower() == "true"
 )
 
 # atlan_client configuration (non ATLAN_ prefix are rooted in pyatlan SDK, to be revisited)

application_sdk/handlers/sql.py

@@ -1,4 +1,5 @@
 import asyncio
+import os
 import re
 from enum import Enum
 from typing import Any, Dict, List, Optional, Set, Tuple
@@ -13,7 +14,6 @@ from application_sdk.common.utils import (
 )
 from application_sdk.constants import SQL_QUERIES_PATH, SQL_SERVER_MIN_VERSION
 from application_sdk.handlers import HandlerInterface
-from application_sdk.inputs.sql_query import SQLQueryInput
 from application_sdk.observability.logger_adaptor import get_logger
 from application_sdk.server.fastapi.models import MetadataType
 
@@ -77,10 +77,7 @@ class BaseSQLHandler(HandlerInterface):
         if self.metadata_sql is None:
             raise ValueError("metadata_sql is not defined")
 
-        sql_input = SQLQueryInput(
-            engine=self.sql_client.engine, query=self.metadata_sql, chunk_size=None
-        )
-        df = await sql_input.get_dataframe()
+        df = await self.sql_client.get_results(self.metadata_sql)
         result: List[Dict[Any, Any]] = []
         try:
             for row in df.to_dict(orient="records"):
@@ -103,12 +100,7 @@ class BaseSQLHandler(HandlerInterface):
         :raises Exception: If the credentials are invalid.
         """
         try:
-            sql_input = SQLQueryInput(
-                engine=self.sql_client.engine,
-                query=self.test_authentication_sql,
-                chunk_size=None,
-            )
-            df = await sql_input.get_dataframe()
+            df = await self.sql_client.get_results(self.test_authentication_sql)
             df.to_dict(orient="records")
             return True
         except Exception as exc:
@@ -335,16 +327,16 @@ class BaseSQLHandler(HandlerInterface):
                 # Use the base query executor in multidb mode to get concatenated df
                 activities = BaseSQLMetadataExtractionActivities()
                 activities.multidb = True
+                base_output_path = payload.get("output_path", "")
                 concatenated_df = await activities.query_executor(
-                    sql_engine=self.sql_client.engine if self.sql_client else None,
+                    sql_client=self.sql_client,
                     sql_query=self.tables_check_sql,
                     workflow_args=payload,
-                    output_suffix="raw/table",
+                    output_path=os.path.join(base_output_path, "raw", "table"),
                     typename="table",
                     write_to_file=False,
                     concatenate=True,
                     return_dataframe=True,
-                    sql_client=self.sql_client,
                 )
 
                 if concatenated_df is None:
@@ -362,12 +354,9 @@ class BaseSQLHandler(HandlerInterface):
             )
             if not query:
                 raise ValueError("tables_check_sql is not defined")
-            sql_input = SQLQueryInput(
-                engine=self.sql_client.engine, query=query, chunk_size=None
-            )
-            sql_input = await sql_input.get_dataframe()
+            sql_results = await self.sql_client.get_results(query)
             try:
-                total = _sum_counts_from_records(sql_input.to_dict(orient="records"))
+                total = _sum_counts_from_records(sql_results.to_dict(orient="records"))
                 return _build_success(total)
             except Exception as exc:
                 return _build_failure(exc)
@@ -404,13 +393,9 @@ class BaseSQLHandler(HandlerInterface):
 
             # If dialect version not available and client_version_sql is defined, use SQL query
             if not client_version and self.client_version_sql:
-                sql_input = await SQLQueryInput(
-                    query=self.client_version_sql,
-                    engine=self.sql_client.engine,
-                    chunk_size=None,
-                ).get_dataframe()
+                sql_results = await self.sql_client.get_results(self.client_version_sql)
                 version_string = next(
-                    iter(sql_input.to_dict(orient="records")[0].values())
+                    iter(sql_results.to_dict(orient="records")[0].values())
                 )
                 version_match = re.search(r"(\d+\.\d+(?:\.\d+)?)", version_string)
                 if version_match:

application_sdk/interceptors/events.py

@@ -12,7 +12,7 @@ from temporalio.worker import (
     WorkflowInterceptorClassInput,
 )
 
-from application_sdk.events.models import (
+from application_sdk.interceptors.models import (
     ApplicationEventNames,
     Event,
     EventMetadata,