atlan-application-sdk 0.1.1rc53__py3-none-any.whl → 0.1.1rc55__py3-none-any.whl

application_sdk/activities/__init__.py

@@ -20,6 +20,7 @@ from typing import Any, Dict, Generic, Optional, TypeVar
 from pydantic import BaseModel
 from temporalio import activity
 
+from application_sdk.activities.common.models import ActivityResult
 from application_sdk.activities.common.utils import (
     auto_heartbeater,
     build_output_path,
@@ -27,6 +28,7 @@ from application_sdk.activities.common.utils import (
     get_workflow_run_id,
 )
 from application_sdk.common.error_codes import OrchestratorError
+from application_sdk.common.file_converter import FileType, convert_data_files
 from application_sdk.constants import TEMPORARY_PATH
 from application_sdk.handlers import HandlerInterface
 from application_sdk.observability.logger_adaptor import get_logger
@@ -268,3 +270,26 @@ class ActivitiesInterface(ABC, Generic[ActivitiesStateType]):
                 exc_info=e,
             )
             raise
+
+    @activity.defn
+    @auto_heartbeater
+    async def convert_files(self, workflow_args: Dict[str, Any]) -> ActivityResult:
+        """
+        Convert the input files to the specified output type.
+        """
+        converted_files = []
+        if workflow_args.get("input_files") and workflow_args.get("output_file_type"):
+            converted_files = await convert_data_files(
+                workflow_args["input_files"],
+                FileType(workflow_args["output_file_type"]),
+            )
+            return ActivityResult(
+                status="success",
+                message=f"Successfully converted files to {workflow_args['output_file_type']}",
+                metadata={"input_files": converted_files},
+            )
+        return ActivityResult(
+            status="warning",
+            message="Unable to get input files or output file type",
+            metadata={"input_files": converted_files},
+        )

application_sdk/activities/common/models.py

@@ -4,7 +4,7 @@ This module contains Pydantic models used to represent various data structures
 needed by activities, such as statistics and configuration.
 """
 
-from typing import List, Optional
+from typing import Any, Dict, List, Optional, TypedDict
 
 from pydantic import BaseModel
 
@@ -36,3 +36,9 @@ class ActivityStatistics(BaseModel):
     chunk_count: int = 0
     partitions: Optional[List[int]] = []
     typename: Optional[str] = None
+
+
+class ActivityResult(TypedDict):
+    status: str
+    message: str
+    metadata: Dict[str, Any]

application_sdk/clients/models.py

@@ -31,7 +31,11 @@ class DatabaseConfig(BaseModel):
     )
     parameters: Optional[List[str]] = Field(
         default=None,
-        description="List of additional connection parameter names that can be dynamically added from credentials",
+        description="List of additional connection parameter names that can be dynamically added from credentials to the connection string. ex: ['ssl_mode'] will be added to the connection string as ?ssl_mode=require",
+    )
+    connect_args: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional connection arguments to be passed to SQLAlchemy. ex: {'sslmode': 'require'}",
     )
 
     class Config:

application_sdk/clients/sql.py

@@ -37,7 +37,6 @@ class BaseSQLClient(ClientInterface):
     Attributes:
         connection: Database connection instance.
         engine: SQLAlchemy engine instance.
-        sql_alchemy_connect_args (Dict[str, Any]): Additional connection arguments.
         credentials (Dict[str, Any]): Database credentials.
         resolved_credentials (Dict[str, Any]): Resolved credentials after reading from secret manager.
         use_server_side_cursor (bool): Whether to use server-side cursors.
@@ -45,7 +44,6 @@ class BaseSQLClient(ClientInterface):
 
     connection = None
     engine = None
-    sql_alchemy_connect_args: Dict[str, Any] = {}
     credentials: Dict[str, Any] = {}
     resolved_credentials: Dict[str, Any] = {}
     use_server_side_cursor: bool = USE_SERVER_SIDE_CURSOR
@@ -55,7 +53,6 @@ class BaseSQLClient(ClientInterface):
         self,
         use_server_side_cursor: bool = USE_SERVER_SIDE_CURSOR,
         credentials: Dict[str, Any] = {},
-        sql_alchemy_connect_args: Dict[str, Any] = {},
     ):
         """
         Initialize the SQL client.
@@ -64,12 +61,9 @@ class BaseSQLClient(ClientInterface):
             use_server_side_cursor (bool, optional): Whether to use server-side cursors.
                 Defaults to USE_SERVER_SIDE_CURSOR.
             credentials (Dict[str, Any], optional): Database credentials. Defaults to {}.
-            sql_alchemy_connect_args (Dict[str, Any], optional): Additional SQLAlchemy
-                connection arguments. Defaults to {}.
         """
         self.use_server_side_cursor = use_server_side_cursor
         self.credentials = credentials
-        self.sql_alchemy_connect_args = sql_alchemy_connect_args
 
     async def load(self, credentials: Dict[str, Any]) -> None:
         """Load credentials and prepare engine for lazy connections.
@@ -83,6 +77,9 @@ class BaseSQLClient(ClientInterface):
         Raises:
             ClientError: If credentials are invalid or engine creation fails
         """
+        if not self.DB_CONFIG:
+            raise ValueError("DB_CONFIG is not configured for this SQL client.")
+
         self.credentials = credentials  # Update the instance credentials
         try:
             from sqlalchemy import create_engine
@@ -90,7 +87,7 @@ class BaseSQLClient(ClientInterface):
             # Create engine but no persistent connection
             self.engine = create_engine(
                 self.get_sqlalchemy_connection_string(),
-                connect_args=self.sql_alchemy_connect_args,
+                connect_args=self.DB_CONFIG.connect_args,
                 pool_pre_ping=True,
             )
 
@@ -397,7 +394,6 @@ class AsyncBaseSQLClient(BaseSQLClient):
     Attributes:
         connection (AsyncConnection): Async database connection instance.
         engine (AsyncEngine): Async SQLAlchemy engine instance.
-        sql_alchemy_connect_args (Dict[str, Any]): Additional connection arguments.
         credentials (Dict[str, Any]): Database credentials.
         use_server_side_cursor (bool): Whether to use server-side cursors.
     """
@@ -419,13 +415,16 @@ class AsyncBaseSQLClient(BaseSQLClient):
             ValueError: If credentials are invalid or engine creation fails.
         """
         self.credentials = credentials
+        if not self.DB_CONFIG:
+            raise ValueError("DB_CONFIG is not configured for this SQL client.")
+
         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,
+                connect_args=self.DB_CONFIG.connect_args,
                 pool_pre_ping=True,
             )
             if not self.engine:

application_sdk/common/file_converter.py

@@ -0,0 +1,99 @@
+from collections import namedtuple
+from enum import Enum
+from typing import List, Optional
+
+import pandas as pd
+
+from application_sdk.observability.logger_adaptor import get_logger
+
+logger = get_logger(__name__)
+
+
+def enum_register():
+    """
+    Helps us register custom function for enum values
+    """
+    registry = {}
+
+    def add(name: str):
+        def inner(fn):
+            registry[name] = fn
+            return fn
+
+        return inner
+
+    Register = namedtuple("Register", ["add", "registry"])
+    return Register(add, registry)
+
+
+file_converter_registry = enum_register()
+
+
+# Edit the enums here to add new file types
+class FileType(Enum):
+    JSON = "json"
+    PARQUET = "parquet"
+
+
+# Edit the enums here to add new file conversions
+class ConvertFile(Enum):
+    JSON_TO_PARQUET = "json_to_parquet"
+    PARQUET_TO_JSON = "parquet_to_json"
+
+
+async def convert_data_files(
+    input_file_paths: List[str], output_file_type: FileType
+) -> List[str]:
+    """
+    Convert the input files to the specified file type
+    Args:
+        input_file_paths: List[str] - List of input file paths
+        output_file_type: FileType - The file type to convert to
+    Returns:
+        List[str] - List of converted file paths
+    """
+    if not input_file_paths:
+        return []
+    input_file_type = input_file_paths[0].split(".")[-1]
+    convert_file = ConvertFile(f"{input_file_type}_to_{output_file_type.value}")
+    converter_func = file_converter_registry.registry.get(convert_file)
+    converted_files = []
+    try:
+        for file in input_file_paths:
+            converted_file = converter_func(file)
+            if converted_file:
+                converted_files.append(converted_file)
+    except KeyError:
+        raise ValueError(f"No converter found for file type: {convert_file}")
+
+    return converted_files
+
+
+# Add the main logic here to convert the files here
+@file_converter_registry.add(ConvertFile.JSON_TO_PARQUET)
+def convert_json_to_parquet(file_path: str) -> Optional[str]:
+    """Convert the downloaded files from json to parquet"""
+    try:
+        logger.info(f"Converting {file_path} to parquet")
+        df = pd.read_json(file_path, orient="records", lines=True)
+        df = df.loc[:, ~df.where(df.astype(bool)).isna().all(axis=0)]
+        parquet_file_path = file_path.replace(".json", ".parquet")
+        df.to_parquet(parquet_file_path)
+        return parquet_file_path
+    except Exception as e:
+        logger.error(f"Error converting {file_path} to parquet: {e}")
+        return None
+
+
+@file_converter_registry.add(ConvertFile.PARQUET_TO_JSON)
+def convert_parquet_to_json(file_path: str) -> Optional[str]:
+    """Convert the downloaded files from parquet to json"""
+    try:
+        logger.info(f"Converting {file_path} to json")
+        df = pd.read_parquet(file_path)
+        json_file_path = file_path.replace(".parquet", ".json")
+        df.to_json(json_file_path, orient="records", lines=True)
+        return json_file_path
+    except Exception as e:
+        logger.error(f"Error converting {file_path} to json: {e}")
+        return None

application_sdk/version.py

@@ -2,4 +2,4 @@
 Version information for the application_sdk package.
 """
 
-__version__ = "0.1.1rc53"
+__version__ = "0.1.1rc55"

{atlan_application_sdk-0.1.1rc53.dist-info → atlan_application_sdk-0.1.1rc55.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atlan-application-sdk
-Version: 0.1.1rc53
+Version: 0.1.1rc55
 Summary: Atlan Application SDK is a Python library for developing applications on the Atlan Platform
 Project-URL: Repository, https://github.com/atlanhq/application-sdk
 Project-URL: Documentation, https://github.com/atlanhq/application-sdk/README.md

{atlan_application_sdk-0.1.1rc53.dist-info → atlan_application_sdk-0.1.1rc55.dist-info}/RECORD

@@ -1,12 +1,12 @@
 application_sdk/__init__.py,sha256=2e2mvmLJ5dxmJGPELtb33xwP-j6JMdoIuqKycEn7hjg,151
 application_sdk/constants.py,sha256=EDGR-3SuCxNV-3x0D4wA9is9vBbVWa3nHvJ8r2w4lYY,10977
-application_sdk/version.py,sha256=0GDfah3Is6w2M-p8lHG9431TleRLSCtZ79vP9LRgIfY,88
+application_sdk/version.py,sha256=GSgUx92mCFDwTnAC5gx_IwXSMV4KwnJrIkRB3TCBNxQ,88
 application_sdk/worker.py,sha256=i5f0AeKI39IfsLO05QkwC6uMz0zDPSJqP7B2byri1VI,7489
-application_sdk/activities/__init__.py,sha256=OhbOrz8ildUCFe4xrjYvceQM36dlDpKTPtw0OnTaPwU,10090
+application_sdk/activities/__init__.py,sha256=L5WXkTwOwGtjWAlXrUJRCKGwyIyp3z8fBv8BZVCRFQI,11175
 application_sdk/activities/lock_management.py,sha256=oX2qPpfEu_xP0MiaCakVGk9ivZDvG4EddVZag1DuHSE,3976
 application_sdk/activities/.cursor/BUGBOT.md,sha256=FNykX5aMkdOhzgpiGqstOnSp9JN63iR2XP3onU4AGh8,15843
 application_sdk/activities/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-application_sdk/activities/common/models.py,sha256=LIZfWvTtgtbAUvvn-rwrPQgD7fP2J0Gxdxr_ITgw-jM,1243
+application_sdk/activities/common/models.py,sha256=43MF_w0EzEQiJvGIqF_FNet4X6MEmwqYd3YAsHdQn08,1362
 application_sdk/activities/common/utils.py,sha256=nSNGkY5eS5pPc8etdPWkXBFTSaConGAD8LDtNqOMHF4,9836
 application_sdk/activities/metadata_extraction/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 application_sdk/activities/metadata_extraction/base.py,sha256=ENFojpxqKdN_eVSL4iet3cGfylPOfcl1jnflfo4zhs8,3920
@@ -20,9 +20,9 @@ application_sdk/clients/__init__.py,sha256=C9T84J7V6ZumcoWJPAxdd3tqSmbyciaGBJn-C
 application_sdk/clients/atlan.py,sha256=l6yV39fr1006SJFwkOTNDQlbSFlHCZQaUPfdUlzdVEg,5053
 application_sdk/clients/atlan_auth.py,sha256=D7FuNqv81ohNXLJtdx1AFw_jU6a3g0Pw6149ia4ucFY,8930
 application_sdk/clients/base.py,sha256=TIn3pG89eXUc1XSYf4jk66m1vajWp0WxcCQOOltdazA,14021
-application_sdk/clients/models.py,sha256=yPgmiqt3I7am2NdFF10BnKiwfgjLceEXrpie0rvGCg8,1509
+application_sdk/clients/models.py,sha256=iZOTyH6LO64kozdiUPCFCN0NgLhd_Gtv0lH7ZIPdo8w,1800
 application_sdk/clients/redis.py,sha256=IfAD32vLp88BCvsDTaQtxFHxzHlEx4V7TK7h1HwDDBg,15917
-application_sdk/clients/sql.py,sha256=r-8rghnATLRxxISchLZaNnMCAZMRLKyPAwPZengzMMY,19846
+application_sdk/clients/sql.py,sha256=lXeVu_dute30IaWWK5gHBhjEs2dXp_e0XkOMsbOsq64,19589
 application_sdk/clients/temporal.py,sha256=jC3U8LmW8G6gg-Qmxk0rcAifIGF0KekwP1UkMGXN7RA,18314
 application_sdk/clients/utils.py,sha256=zLFOJbTr_6TOqnjfVFGY85OtIXZ4FQy_rquzjaydkbY,779
 application_sdk/clients/workflow.py,sha256=6bSqmA3sNCk9oY68dOjBUDZ9DhNKQxPD75qqE0cfldc,6104
@@ -32,6 +32,7 @@ application_sdk/common/aws_utils.py,sha256=xlSMIQyjvQ-CydEXaxXrnPUygv7AAbCLsxhZ2
 application_sdk/common/dapr_utils.py,sha256=0yHqDP6qNb1OT-bX2XRYQPZ5xkGkV13nyRw6GkPlHs8,1136
 application_sdk/common/dataframe_utils.py,sha256=PId9vT6AUoq3tesiTd4sSUvW7RUhPWdAAEBLuOprks4,1262
 application_sdk/common/error_codes.py,sha256=bxgvugN_0H5b8VXfJw-44mybgX5I9lRJbRdYjtPjqDI,14561
+application_sdk/common/file_converter.py,sha256=ta0PVh7uIEGJg0BTPUJnSjj55ve2iVAOkqwAeg96_-g,3079
 application_sdk/common/utils.py,sha256=ImCrlyCj5Mj571CVWfqy5MynVVju9xhn1ItSlJoaebc,19572
 application_sdk/common/.cursor/BUGBOT.md,sha256=OkB5TMAEJFzaBfbNb3g9ZDPW2r1krQE_KEuJbytMPuI,12176
 application_sdk/decorators/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -156,8 +157,8 @@ application_sdk/workflows/metadata_extraction/__init__.py,sha256=jHUe_ZBQ66jx8bg
 application_sdk/workflows/metadata_extraction/sql.py,sha256=6ZaVt84n-8U2ZvR9GR7uIJKv5v8CuyQjhlnoRJvDszc,12435
 application_sdk/workflows/query_extraction/__init__.py,sha256=n066_CX5RpJz6DIxGMkKS3eGSRg03ilaCtsqfJWQb7Q,117
 application_sdk/workflows/query_extraction/sql.py,sha256=kT_JQkLCRZ44ZpaC4QvPL6DxnRIIVh8gYHLqRbMI-hA,4826
-atlan_application_sdk-0.1.1rc53.dist-info/METADATA,sha256=IIkNqjucZ5QB-EUMFTFaucgdxl-U6ouulLKzH3atL_I,5634
-atlan_application_sdk-0.1.1rc53.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-atlan_application_sdk-0.1.1rc53.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-atlan_application_sdk-0.1.1rc53.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
-atlan_application_sdk-0.1.1rc53.dist-info/RECORD,,
+atlan_application_sdk-0.1.1rc55.dist-info/METADATA,sha256=oPFC4krNFgKJYPm5q8u6Cji9hKoZfflnMTTGi8-c7SQ,5634
+atlan_application_sdk-0.1.1rc55.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+atlan_application_sdk-0.1.1rc55.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+atlan_application_sdk-0.1.1rc55.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
+atlan_application_sdk-0.1.1rc55.dist-info/RECORD,,