atlan-application-sdk 0.1.1rc40__py3-none-any.whl → 0.1.1rc42__py3-none-any.whl

application_sdk/handlers/sql.py

@@ -56,9 +56,13 @@ class BaseSQLHandler(HandlerInterface):
     schema_alias_key: str = SQLConstants.SCHEMA_ALIAS_KEY.value
     database_result_key: str = SQLConstants.DATABASE_RESULT_KEY.value
     schema_result_key: str = SQLConstants.SCHEMA_RESULT_KEY.value
+    multidb: bool = False
 
-    def __init__(self, sql_client: BaseSQLClient | None = None):
+    def __init__(
+        self, sql_client: BaseSQLClient | None = None, multidb: Optional[bool] = False
+    ):
         self.sql_client = sql_client
+        self.multidb = multidb
 
     async def load(self, credentials: Dict[str, Any]) -> None:
         """
@@ -294,35 +298,26 @@ class BaseSQLHandler(HandlerInterface):
                         return False, f"{db}.{sch} schema"
         return True, ""
 
-    async def tables_check(
-        self,
-        payload: Dict[str, Any],
-    ) -> Dict[str, Any]:
+    async def tables_check(self, payload: Dict[str, Any]) -> Dict[str, Any]:
         """
         Method to check the count of tables
         """
         logger.info("Starting tables check")
-        query = prepare_query(
-            query=self.tables_check_sql,
-            workflow_args=payload,
-            temp_table_regex_sql=self.extract_temp_table_regex_table_sql,
-        )
-        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()
-        try:
-            result = 0
-            for row in sql_input.to_dict(orient="records"):
-                result += row["count"]
+
+        def _sum_counts_from_records(records_iter) -> int:
+            total = 0
+            for row in records_iter:
+                total += row["count"]
+            return total
+
+        def _build_success(total: int) -> Dict[str, Any]:
             return {
                 "success": True,
-                "successMessage": f"Tables check successful. Table count: {result}",
+                "successMessage": f"Tables check successful. Table count: {total}",
                 "failureMessage": "",
             }
-        except Exception as exc:
+
+        def _build_failure(exc: Exception) -> Dict[str, Any]:
             logger.error("Error during tables check", exc_info=True)
             return {
                 "success": False,
@@ -331,6 +326,52 @@ class BaseSQLHandler(HandlerInterface):
                 "error": str(exc),
             }
 
+        if self.multidb:
+            try:
+                from application_sdk.activities.metadata_extraction.sql import (
+                    BaseSQLMetadataExtractionActivities,
+                )
+
+                # Use the base query executor in multidb mode to get concatenated df
+                activities = BaseSQLMetadataExtractionActivities()
+                activities.multidb = True
+                concatenated_df = await activities.query_executor(
+                    sql_engine=self.sql_client.engine if self.sql_client else None,
+                    sql_query=self.tables_check_sql,
+                    workflow_args=payload,
+                    output_suffix="raw/table",
+                    typename="table",
+                    write_to_file=False,
+                    concatenate=True,
+                    return_dataframe=True,
+                    sql_client=self.sql_client,
+                )
+
+                if concatenated_df is None:
+                    return _build_success(0)
+
+                total = int(concatenated_df["count"].sum())  # type: ignore[index]
+                return _build_success(total)
+            except Exception as exc:
+                return _build_failure(exc)
+        else:
+            query = prepare_query(
+                query=self.tables_check_sql,
+                workflow_args=payload,
+                temp_table_regex_sql=self.extract_temp_table_regex_table_sql,
+            )
+            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()
+            try:
+                total = _sum_counts_from_records(sql_input.to_dict(orient="records"))
+                return _build_success(total)
+            except Exception as exc:
+                return _build_failure(exc)
+
     async def check_client_version(self) -> Dict[str, Any]:
         """
         Check if the client version meets the minimum required version.

application_sdk/inputs/__init__.py

@@ -1,7 +1,15 @@
+import os
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, AsyncIterator, Iterator, Union
-
+from typing import TYPE_CHECKING, AsyncIterator, Iterator, List, Union
+
+from application_sdk.activities.common.utils import (
+    find_local_files_by_extension,
+    get_object_store_prefix,
+)
+from application_sdk.common.error_codes import IOError
+from application_sdk.constants import TEMPORARY_PATH
 from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services.objectstore import ObjectStore
 
 logger = get_logger(__name__)
 
@@ -15,6 +23,94 @@ class Input(ABC):
     Abstract base class for input data sources.
     """
 
+    async def download_files(self) -> List[str]:
+        """Download files from object store if not available locally.
+
+        Flow:
+        1. Check if files exist locally at self.path
+        2. If not, try to download from object store
+        3. Filter by self.file_names if provided
+        4. Return list of file paths for logging purposes
+
+        Returns:
+            List[str]: List of file paths
+
+        Raises:
+            AttributeError: When the input class doesn't support file operations or _extension
+            IOError: When no files found locally or in object store
+        """
+        # Step 1: Check if files exist locally
+        local_files = find_local_files_by_extension(
+            self.path, self._EXTENSION, self.file_names
+        )
+        if local_files:
+            logger.info(
+                f"Found {len(local_files)} {self._EXTENSION} files locally at: {self.path}"
+            )
+            return local_files
+
+        # Step 2: Try to download from object store
+        logger.info(
+            f"No local {self._EXTENSION} files found at {self.path}, checking object store..."
+        )
+
+        try:
+            # Determine what to download based on path type and filters
+            downloaded_paths = []
+
+            if self.path.endswith(self._EXTENSION):
+                # Single file case (file_names validation already ensures this is valid)
+                source_path = get_object_store_prefix(self.path)
+                destination_path = os.path.join(TEMPORARY_PATH, source_path)
+                await ObjectStore.download_file(
+                    source=source_path,
+                    destination=destination_path,
+                )
+                downloaded_paths.append(destination_path)
+
+            elif self.file_names:
+                # Directory with specific files - download each file individually
+                for file_name in self.file_names:
+                    file_path = os.path.join(self.path, file_name)
+                    source_path = get_object_store_prefix(file_path)
+                    destination_path = os.path.join(TEMPORARY_PATH, source_path)
+                    await ObjectStore.download_file(
+                        source=source_path,
+                        destination=destination_path,
+                    )
+                    downloaded_paths.append(destination_path)
+            else:
+                # Download entire directory
+                source_path = get_object_store_prefix(self.path)
+                destination_path = os.path.join(TEMPORARY_PATH, source_path)
+                await ObjectStore.download_prefix(
+                    source=source_path,
+                    destination=destination_path,
+                )
+                # Find the actual files in the downloaded directory
+                found_files = find_local_files_by_extension(
+                    destination_path, self._EXTENSION, getattr(self, "file_names", None)
+                )
+                downloaded_paths.extend(found_files)
+
+            # Check results
+            if downloaded_paths:
+                logger.info(
+                    f"Successfully downloaded {len(downloaded_paths)} {self._EXTENSION} files from object store"
+                )
+                return downloaded_paths
+            else:
+                raise IOError(
+                    f"{IOError.OBJECT_STORE_READ_ERROR}: Downloaded from object store but no {self._EXTENSION} files found"
+                )
+
+        except Exception as e:
+            logger.error(f"Failed to download from object store: {str(e)}")
+            raise IOError(
+                f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: No {self._EXTENSION} files found locally at '{self.path}' and failed to download from object store. "
+                f"Error: {str(e)}"
+            )
+
     @abstractmethod
     async def get_batched_dataframe(
         self,

application_sdk/inputs/json.py

@@ -1,11 +1,7 @@
-import os
 from typing import TYPE_CHECKING, AsyncIterator, Iterator, List, Optional, Union
 
-from application_sdk.activities.common.utils import get_object_store_prefix
-from application_sdk.common.error_codes import IOError
 from application_sdk.inputs import Input
 from application_sdk.observability.logger_adaptor import get_logger
-from application_sdk.services.objectstore import ObjectStore
 
 if TYPE_CHECKING:
     import daft
@@ -15,56 +11,43 @@ logger = get_logger(__name__)
 
 
 class JsonInput(Input):
-    path: str
-    chunk_size: Optional[int]
-    file_names: Optional[List[str]]
-    download_file_prefix: Optional[str]
+    """
+    JSON Input class to read data from JSON files using daft and pandas.
+    Supports reading both single files and directories containing multiple JSON files.
+    """
+
+    _EXTENSION = ".json"
 
     def __init__(
         self,
         path: str,
         file_names: Optional[List[str]] = None,
-        download_file_prefix: Optional[str] = None,
-        chunk_size: Optional[int] = None,
+        chunk_size: int = 100000,
     ):
         """Initialize the JsonInput class.
 
         Args:
-            path (str): The path to the input directory.
-            file_names (Optional[List[str]]): The list of files to read.
-            download_file_prefix (Optional[str]): The prefix path in object store.
-            chunk_size (Optional[int]): The chunk size to read the data. If None, uses config value.
+            path (str): Path to JSON file or directory containing JSON files.
+                It accepts both types of paths:
+                local path or object store path
+                Wildcards are not supported.
+            file_names (Optional[List[str]]): List of specific file names to read. Defaults to None.
+            chunk_size (int): Number of rows per batch. Defaults to 100000.
+
+        Raises:
+            ValueError: When path is not provided or when single file path is combined with file_names
         """
+
+        # Validate that single file path and file_names are not both specified
+        if path.endswith(self._EXTENSION) and file_names:
+            raise ValueError(
+                f"Cannot specify both a single file path ('{path}') and file_names filter. "
+                f"Either provide a directory path with file_names, or specify the exact file path without file_names."
+            )
+
         self.path = path
-        # If chunk_size is provided, use it; otherwise default to 100,000 rows per batch
-        self.chunk_size = chunk_size if chunk_size is not None else 100000
+        self.chunk_size = chunk_size
         self.file_names = file_names
-        self.download_file_prefix = download_file_prefix
-
-    async def download_files(self):
-        """Download the files from the object store to the local path"""
-        if not self.file_names:
-            logger.debug("No files to download")
-            return
-
-        for file_name in self.file_names or []:
-            try:
-                if self.download_file_prefix is not None and not os.path.exists(
-                    os.path.join(self.path, file_name)
-                ):
-                    destination_file_path = os.path.join(self.path, file_name)
-                    await ObjectStore.download_file(
-                        source=get_object_store_prefix(destination_file_path),
-                        destination=destination_file_path,
-                    )
-            except IOError as e:
-                logger.error(
-                    f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error downloading file {file_name}: {str(e)}",
-                    error_code=IOError.OBJECT_STORE_DOWNLOAD_ERROR.code,
-                )
-                raise IOError(
-                    f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error downloading file {file_name}: {str(e)}"
-                )
 
     async def get_batched_dataframe(
         self,
@@ -76,22 +59,20 @@ class JsonInput(Input):
         try:
             import pandas as pd
 
-            await self.download_files()
+            # Ensure files are available (local or downloaded)
+            json_files = await self.download_files()
+            logger.info(f"Reading {len(json_files)} JSON files in batches")
 
-            for file_name in self.file_names or []:
-                file_path = os.path.join(self.path, file_name)
+            for json_file in json_files:
                 json_reader_obj = pd.read_json(
-                    file_path,
+                    json_file,
                     chunksize=self.chunk_size,
                     lines=True,
                 )
                 for chunk in json_reader_obj:
                     yield chunk
-        except IOError as e:
-            logger.error(
-                f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error reading batched data from JSON: {str(e)}",
-                error_code=IOError.OBJECT_STORE_DOWNLOAD_ERROR.code,
-            )
+        except Exception as e:
+            logger.error(f"Error reading batched data from JSON: {str(e)}")
             raise
 
     async def get_dataframe(self) -> "pd.DataFrame":
@@ -102,21 +83,17 @@ class JsonInput(Input):
         try:
             import pandas as pd
 
-            dataframes = []
-            await self.download_files()
-            for file_name in self.file_names or []:
-                dataframes.append(
-                    pd.read_json(
-                        os.path.join(self.path, file_name),
-                        lines=True,
-                    )
-                )
-            return pd.concat(dataframes, ignore_index=True)
-        except IOError as e:
-            logger.error(
-                f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error reading data from JSON: {str(e)}",
-                error_code=IOError.OBJECT_STORE_DOWNLOAD_ERROR.code,
+            # Ensure files are available (local or downloaded)
+            json_files = await self.download_files()
+            logger.info(f"Reading {len(json_files)} JSON files as pandas dataframe")
+
+            return pd.concat(
+                (pd.read_json(json_file, lines=True) for json_file in json_files),
+                ignore_index=True,
             )
+
+        except Exception as e:
+            logger.error(f"Error reading data from JSON: {str(e)}")
             raise
 
     async def get_batched_daft_dataframe(
@@ -129,18 +106,15 @@ class JsonInput(Input):
         try:
             import daft
 
-            await self.download_files()
-            for file_name in self.file_names or []:
-                json_reader_obj = daft.read_json(
-                    path=os.path.join(self.path, file_name),
-                    _chunk_size=self.chunk_size,
-                )
-                yield json_reader_obj
-        except IOError as e:
-            logger.error(
-                f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error reading batched data from JSON: {str(e)}",
-                error_code=IOError.OBJECT_STORE_DOWNLOAD_ERROR.code,
-            )
+            # Ensure files are available (local or downloaded)
+            json_files = await self.download_files()
+            logger.info(f"Reading {len(json_files)} JSON files as daft batches")
+
+            # Yield each discovered file as separate batch with chunking
+            for json_file in json_files:
+                yield daft.read_json(json_file, _chunk_size=self.chunk_size)
+        except Exception as e:
+            logger.error(f"Error reading batched data from JSON using daft: {str(e)}")
             raise
 
     async def get_daft_dataframe(self) -> "daft.DataFrame":  # noqa: F821
@@ -151,14 +125,12 @@ class JsonInput(Input):
         try:
             import daft
 
-            await self.download_files()
-            if not self.file_names or len(self.file_names) == 0:
-                raise ValueError("No files to read")
-            directory = os.path.join(self.path, self.file_names[0].split("/")[0])
-            return daft.read_json(path=f"{directory}/*.json")
-        except IOError as e:
-            logger.error(
-                f"{IOError.OBJECT_STORE_DOWNLOAD_ERROR}: Error reading data from JSON using daft: {str(e)}",
-                error_code=IOError.OBJECT_STORE_DOWNLOAD_ERROR.code,
-            )
+            # Ensure files are available (local or downloaded)
+            json_files = await self.download_files()
+            logger.info(f"Reading {len(json_files)} JSON files with daft")
+
+            # Use the discovered/downloaded files directly
+            return daft.read_json(json_files)
+        except Exception as e:
+            logger.error(f"Error reading data from JSON using daft: {str(e)}")
             raise