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

application_sdk/activities/__init__.py

@@ -15,6 +15,7 @@ Example:
 
 import os
 from abc import ABC
+from datetime import datetime, timedelta
 from typing import Any, Dict, Generic, Optional, TypeVar
 
 from pydantic import BaseModel
@@ -62,6 +63,7 @@ class ActivitiesState(BaseModel, Generic[HandlerType]):
     model_config = {"arbitrary_types_allowed": True}
     handler: Optional[HandlerType] = None
     workflow_args: Optional[Dict[str, Any]] = None
+    last_updated_timestamp: Optional[datetime] = None
 
 
 ActivitiesStateType = TypeVar("ActivitiesStateType", bound=ActivitiesState)
@@ -113,12 +115,15 @@ class ActivitiesInterface(ABC, Generic[ActivitiesStateType]):
         Note:
             The workflow ID is automatically retrieved from the current activity context.
             If no state exists for the current workflow, a new one will be created.
+            This method also updates the last_updated_timestamp to enable time-based
+            state refresh functionality.
         """
         workflow_id = get_workflow_id()
         if not self._state.get(workflow_id):
             self._state[workflow_id] = ActivitiesState()
 
         self._state[workflow_id].workflow_args = workflow_args
+        self._state[workflow_id].last_updated_timestamp = datetime.now()
 
     async def _get_state(self, workflow_args: Dict[str, Any]) -> ActivitiesStateType:
         """Retrieve the state for the current workflow.
@@ -142,6 +147,15 @@ class ActivitiesInterface(ABC, Generic[ActivitiesStateType]):
             workflow_id = get_workflow_id()
             if workflow_id not in self._state:
                 await self._set_state(workflow_args)
+
+            else:
+                current_timestamp = datetime.now()
+                # if difference of current_timestamp and last_updated_timestamp is greater than 15 minutes, then again _set_state
+                last_updated = self._state[workflow_id].last_updated_timestamp
+                if last_updated and current_timestamp - last_updated > timedelta(
+                    minutes=15
+                ):
+                    await self._set_state(workflow_args)
             return self._state[workflow_id]
         except OrchestratorError as e:
             logger.error(

application_sdk/activities/metadata_extraction/sql.py

@@ -1,4 +1,5 @@
 import os
+from datetime import datetime
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -60,6 +61,7 @@ class BaseSQLMetadataExtractionActivitiesState(ActivitiesState):
     sql_client: Optional[BaseSQLClient] = None
     handler: Optional[BaseSQLHandler] = None
     transformer: Optional[TransformerInterface] = None
+    last_updated_timestamp: Optional[datetime] = None
 
 
 class BaseSQLMetadataExtractionActivities(ActivitiesInterface):
@@ -149,13 +151,30 @@ class BaseSQLMetadataExtractionActivities(ActivitiesInterface):
 
         Args:
             workflow_args (Dict[str, Any]): Arguments passed to the workflow.
+
+        Note:
+            This method creates and configures the new SQL client before closing
+            the old one to ensure state is never left with a closed client if
+            initialization fails. The timestamp is only updated after the new
+            client is successfully created and assigned.
         """
         workflow_id = get_workflow_id()
         if not self._state.get(workflow_id):
             self._state[workflow_id] = BaseSQLMetadataExtractionActivitiesState()
 
-        await super()._set_state(workflow_args)
+        existing_state = self._state[workflow_id]
+
+        # Update workflow_args early, but preserve old timestamp until new client is ready
+        # This ensures that if initialization fails, the state can still be refreshed
+        existing_state.workflow_args = workflow_args
+
+        # Store reference to old client for cleanup after new client is ready
+        old_sql_client = None
+        if existing_state and existing_state.sql_client is not None:
+            old_sql_client = existing_state.sql_client
 
+        # Create and configure new client BEFORE closing old one
+        # This ensures state is never left with a closed client if initialization fails
         sql_client = self.sql_client_class()
 
         # Load credentials BEFORE creating handler to avoid race condition
@@ -165,10 +184,29 @@ class BaseSQLMetadataExtractionActivities(ActivitiesInterface):
             )
             await sql_client.load(credentials)
 
-        # Assign sql_client and handler to state AFTER credentials are loaded
+        # Only after new client is successfully created and configured,
+        # close old client and assign new one to state
+        if old_sql_client is not None:
+            try:
+                await old_sql_client.close()
+                logger.debug(
+                    f"Closed existing SQL client for workflow {workflow_id} during state refresh"
+                )
+            except Exception as e:
+                logger.warning(
+                    f"Failed to close existing SQL client for workflow {workflow_id}: {e}",
+                    exc_info=True,
+                )
+                # Continue even if close fails - new client is already ready
+
+        # Assign sql_client and handler to state AFTER new client is ready
         self._state[workflow_id].sql_client = sql_client
         handler = self.handler_class(sql_client)
         self._state[workflow_id].handler = handler
+        # Update timestamp only after successful client creation and assignment
+        # This ensures that if initialization fails, the old timestamp remains
+        # and the state can be refreshed again immediately
+        self._state[workflow_id].last_updated_timestamp = datetime.now()
 
         # Create transformer with required parameters from ApplicationConstants
         transformer_params = {

application_sdk/transformers/query/__init__.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Optional, Tuple, Type
 
 import daft
 import yaml
+from daft.functions import to_struct, when
 from pyatlan.model.enums import AtlanConnectorType
 
 from application_sdk.observability.logger_adaptor import get_logger
@@ -227,7 +228,7 @@ class QueryBasedTransformer(TransformerInterface):
         # Only create a struct if we have fields
         if struct_fields:
             # Create the struct first
-            struct = daft.struct(*struct_fields)
+            struct = to_struct(*struct_fields)
 
             # If we have non-null checks, apply them
             if non_null_fields:
@@ -236,8 +237,8 @@ class QueryBasedTransformer(TransformerInterface):
                 for check in non_null_fields[1:]:
                     any_non_null = any_non_null | check
 
-                # Use if_else on the any_non_null Expression
-                return any_non_null.if_else(struct, None).alias(prefix)
+                # Use when().otherwise() for conditional expression (replaces if_else in daft 0.7+)
+                return when(any_non_null, struct).otherwise(None).alias(prefix)
 
             return struct.alias(prefix)
 

application_sdk/version.py

@@ -2,4 +2,4 @@
 Version information for the application_sdk package.
 """
 
-__version__ = "2.1.0"
+__version__ = "2.1.1"

{atlan_application_sdk-2.1.0.dist-info → atlan_application_sdk-2.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atlan-application-sdk
-Version: 2.1.0
+Version: 2.1.1
 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
@@ -31,7 +31,7 @@ Requires-Dist: pydantic<2.13.0,>=2.10.6
 Requires-Dist: python-dotenv<1.3.0,>=1.1.0
 Requires-Dist: uvloop<0.23.0,>=0.21.0; sys_platform != 'win32'
 Provides-Extra: daft
-Requires-Dist: daft<0.8.0,>=0.4.12; extra == 'daft'
+Requires-Dist: daft<0.8.0,>=0.7.1; extra == 'daft'
 Provides-Extra: distributed-lock
 Requires-Dist: redis[hiredis]<7.2.0,>=5.2.0; extra == 'distributed-lock'
 Provides-Extra: iam-auth

{atlan_application_sdk-2.1.0.dist-info → atlan_application_sdk-2.1.1.dist-info}/RECORD

@@ -1,8 +1,8 @@
 application_sdk/__init__.py,sha256=2e2mvmLJ5dxmJGPELtb33xwP-j6JMdoIuqKycEn7hjg,151
 application_sdk/constants.py,sha256=TvdmKQShVWBNQZdVF2y-fxuE31FmeraTnqQ9jT_n5XY,11567
-application_sdk/version.py,sha256=b8rehaDy9Q_sCU-sBPVls3ZFH8_ZNacFO5aQrIaO8ts,84
+application_sdk/version.py,sha256=sNbvXviG7NgxM58lOHKhbZfERat5qAJNr3UZy_toVQs,84
 application_sdk/worker.py,sha256=DLMocpHvvwpdAopyXhxwM7ftaNlKvZMQfkgy1MFyiik,7561
-application_sdk/activities/__init__.py,sha256=6SiefuOPUDGfN3z6oPY4RkQLiUmkHpoDy5xadzpDzAw,11588
+application_sdk/activities/__init__.py,sha256=i7iY6aL1VFg185n2rLLvD_sI2BA9zJ33jL5rD_sY__U,12350
 application_sdk/activities/lock_management.py,sha256=6Wdf3jMKitoarHQP91PIJOoGFz4aaOLS_40c7n1yAOA,3902
 application_sdk/activities/.cursor/BUGBOT.md,sha256=FNykX5aMkdOhzgpiGqstOnSp9JN63iR2XP3onU4AGh8,15843
 application_sdk/activities/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -12,7 +12,7 @@ application_sdk/activities/common/utils.py,sha256=ngyFmiZnMCAQtyu6vGeAlkzwNkM29M
 application_sdk/activities/metadata_extraction/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 application_sdk/activities/metadata_extraction/base.py,sha256=ENFojpxqKdN_eVSL4iet3cGfylPOfcl1jnflfo4zhs8,3920
 application_sdk/activities/metadata_extraction/rest.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-application_sdk/activities/metadata_extraction/sql.py,sha256=IkI1ZhOKyoSwosRT-g8c8IDBuFBq7mwyHLpDvwYO_B4,25451
+application_sdk/activities/metadata_extraction/sql.py,sha256=CmE77EsgbOuDL5AKaRCnq1jApJnDWNVxx-RZ49cJwus,27415
 application_sdk/activities/query_extraction/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 application_sdk/activities/query_extraction/sql.py,sha256=Gsa79R8CYY0uyt3rA2nLMfQs8-C4_zg1pJ_yYSF2cZw,21193
 application_sdk/application/__init__.py,sha256=vcrQsqlfmGvKcCZuOtHHaNRqHSGdXlEDftkb8Tv_shI,9867
@@ -138,7 +138,7 @@ application_sdk/transformers/atlas/__init__.py,sha256=fw3D8bBtt61SseAfYut3JZddpX
 application_sdk/transformers/atlas/sql.py,sha256=rkQXNZ7oebts5oF5E_Bw8NpcHHKScU0TmKciH_1l_k4,50419
 application_sdk/transformers/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 application_sdk/transformers/common/utils.py,sha256=4ISMIQ0Gzghmi31p51FOFm5KLF7XF-fmH9PVT7i0DFE,4899
-application_sdk/transformers/query/__init__.py,sha256=yG1dGP3NhUizwkCgyFAzsr9SV9uWYZKjXoCWPrsIxVw,17358
+application_sdk/transformers/query/__init__.py,sha256=4uVCU-NfDe08PlffjWQ5p4smQa7c518IL2rDgIk6694,17446
 application_sdk/transformers/query/templates/column.yaml,sha256=EXLYwGXN7LKT-v51n2EZnY99o6vHucyFaVSpM-sUSXw,7679
 application_sdk/transformers/query/templates/database.yaml,sha256=SD1hJg5LI7gsBHQL5mW341sa51EkhcsIDDFlIOi9zdk,1374
 application_sdk/transformers/query/templates/extras-procedure.yaml,sha256=XhAfVY4zm99K8fcgkYA1XPLv4ks-SA6SzMO3SMtQ60s,2298
@@ -152,8 +152,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-2.1.0.dist-info/METADATA,sha256=XTa1eREHTJv_zkHQhKdCtxe-txAkf9pmIqSEch_zbUA,5806
-atlan_application_sdk-2.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-atlan_application_sdk-2.1.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-atlan_application_sdk-2.1.0.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
-atlan_application_sdk-2.1.0.dist-info/RECORD,,
+atlan_application_sdk-2.1.1.dist-info/METADATA,sha256=Vc2uG2FMhuNXyZFXmGMmvc_LRpCBaNTcQEHpSV8NpOE,5805
+atlan_application_sdk-2.1.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+atlan_application_sdk-2.1.1.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+atlan_application_sdk-2.1.1.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
+atlan_application_sdk-2.1.1.dist-info/RECORD,,