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

application_sdk/worker.py

@@ -1,10 +1,11 @@
 """Worker module for managing Temporal workers.
 
 This module provides the Worker class for managing Temporal workflow workers,
-including their initialization, configuration, and execution.
+including their initialization, configuration, and execution with graceful shutdown support.
 """
 
 import asyncio
+import signal
 import sys
 import threading
 from concurrent.futures import ThreadPoolExecutor
@@ -14,7 +15,11 @@ from temporalio.types import CallableType, ClassType
 from temporalio.worker import Worker as TemporalWorker
 
 from application_sdk.clients.workflow import WorkflowClient
-from application_sdk.constants import DEPLOYMENT_NAME, MAX_CONCURRENT_ACTIVITIES
+from application_sdk.constants import (
+    DEPLOYMENT_NAME,
+    GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS,
+    MAX_CONCURRENT_ACTIVITIES,
+)
 from application_sdk.interceptors.models import (
     ApplicationEventNames,
     Event,
@@ -37,13 +42,16 @@ if sys.platform not in ("win32", "cygwin"):
         # uvloop is not available, use default asyncio
         logger.warning("uvloop is not available, using default asyncio")
         pass
+elif sys.platform == "win32":
+    # Use WindowsSelectorEventLoopPolicy for Windows platform
+    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
 
 
 class Worker:
     """Worker class for managing Temporal workflow workers.
 
     This class handles the initialization and execution of Temporal workers,
-    including their activities, workflows, and module configurations.
+    including their activities, workflows, module configurations, and graceful shutdown.
 
     Attributes:
         workflow_client: Client for interacting with Temporal.
@@ -53,6 +61,13 @@ class Worker:
         passthrough_modules: List of module names to pass through.
         max_concurrent_activities: Maximum number of concurrent activities.
 
+    Graceful Shutdown:
+        When SIGTERM or SIGINT is received:
+        1. Signal handlers trigger worker.shutdown()
+        2. Worker stops polling for new tasks
+        3. In-flight activities are allowed to complete within graceful_shutdown_timeout
+        4. Worker exits early if all activities complete, or at timeout
+
     Note:
         This class is designed to be thread-safe when running workers in daemon mode.
         However, care should be taken when modifying worker attributes after initialization.
@@ -71,6 +86,78 @@ class Worker:
 
     default_passthrough_modules = ["application_sdk", "pandas", "os", "app"]
 
+    def _setup_signal_handlers(self) -> None:
+        """Set up SIGTERM and SIGINT handlers for graceful shutdown.
+
+        Signal handlers can only be registered from the main thread.
+
+        Platform Notes:
+            - Unix/Linux/macOS: Full signal handling support via asyncio event loop.
+            - Windows: Signal handling is not supported. Workers on Windows will not
+              respond to SIGTERM/SIGINT for graceful shutdown. On Windows, the worker
+              will continue running until the process is forcefully terminated.
+        """
+        # Signal handlers only work on Unix-like systems
+        if sys.platform in ("win32", "cygwin"):
+            logger.warning(
+                "Signal handlers not supported on Windows. "
+                "Graceful shutdown via SIGTERM/SIGINT is not available. "
+                "For production deployments, use Unix-based systems."
+            )
+            return
+
+        # Signal handlers can only be registered from the main thread
+        if threading.current_thread() is not threading.main_thread():
+            logger.debug(
+                "Skipping signal handler registration - not running in main thread"
+            )
+            return
+
+        loop = asyncio.get_running_loop()
+
+        def handle_signal(sig_name: str) -> None:
+            """Handle shutdown signal by triggering worker.shutdown().
+
+            Uses a flag to prevent multiple shutdown tasks from being created
+            if multiple signals are received in quick succession.
+            """
+            if self._shutdown_initiated:
+                logger.debug(f"Received {sig_name}, but shutdown already in progress")
+                return
+
+            self._shutdown_initiated = True
+            logger.info(
+                f"Received {sig_name}, initiating graceful shutdown "
+                f"(timeout: {GRACEFUL_SHUTDOWN_TIMEOUT_SECONDS}s)"
+            )
+            if self.workflow_worker:
+                asyncio.create_task(self._shutdown_worker())
+
+        try:
+            loop.add_signal_handler(signal.SIGTERM, lambda: handle_signal("SIGTERM"))
+            loop.add_signal_handler(signal.SIGINT, lambda: handle_signal("SIGINT"))
+            logger.debug("Registered SIGTERM and SIGINT handlers")
+        except (ValueError, RuntimeError) as e:
+            logger.warning(f"Could not set up signal handlers: {e}")
+
+    async def _shutdown_worker(self) -> None:
+        """Shutdown the worker gracefully.
+
+        Calls worker.shutdown() which:
+        1. Stops polling for new tasks
+        2. Waits for in-flight activities (up to graceful_shutdown_timeout)
+        3. Returns when done or timeout reached
+        """
+        if not self.workflow_worker:
+            return
+
+        try:
+            logger.info("Stopping polling, waiting for in-flight activities...")
+            await self.workflow_worker.shutdown()
+            logger.info("Worker shutdown complete")
+        except Exception as e:
+            logger.error(f"Error during shutdown: {e}")
+
     def __init__(
         self,
         workflow_client: Optional[WorkflowClient] = None,
@@ -105,6 +192,7 @@ class Worker:
         """
         self.workflow_client = workflow_client
         self.workflow_worker: Optional[TemporalWorker] = None
+        self._shutdown_initiated = False
         self.workflow_activities = workflow_activities
         self.workflow_classes = workflow_classes
         self.passthrough_modules = list(
@@ -153,9 +241,11 @@ class Worker:
             RuntimeError: If worker creation fails.
             ConnectionError: If connection to Temporal server fails.
 
-        Note:
-            When running as a daemon, the worker runs in a separate thread and
-            does not block the main thread.
+        Graceful Shutdown:
+            When SIGTERM/SIGINT is received:
+            1. Worker stops polling for new tasks
+            2. In-flight activities complete within graceful_shutdown_timeout
+            3. Worker exits early if activities finish, or at timeout
         """
         if daemon:
             worker_thread = threading.Thread(
@@ -184,11 +274,19 @@ class Worker:
                 max_concurrent_activities=self.max_concurrent_activities,
                 activity_executor=self.activity_executor,
             )
+            self.workflow_worker = worker
 
             logger.info(
                 f"Starting worker with task queue: {self.workflow_client.worker_task_queue}"
             )
+
+            # Set up signal handlers and run worker
+            self._setup_signal_handlers()
+
             await worker.run()
+
+            logger.info("Worker stopped")
+
         except Exception as e:
             logger.error(f"Error starting worker: {e}")
             raise e

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atlan-application-sdk
-Version: 2.1.1
+Version: 2.3.0
 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
@@ -30,6 +30,10 @@ Requires-Dist: pyatlan<8.5.0,>=8.0.2
 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: azure
+Requires-Dist: azure-identity>=1.15.0; extra == 'azure'
+Requires-Dist: azure-storage-blob>=12.19.0; extra == 'azure'
+Requires-Dist: azure-storage-file-datalake>=12.19.0; extra == 'azure'
 Provides-Extra: daft
 Requires-Dist: daft<0.8.0,>=0.7.1; extra == 'daft'
 Provides-Extra: distributed-lock

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

@@ -1,7 +1,7 @@
 application_sdk/__init__.py,sha256=2e2mvmLJ5dxmJGPELtb33xwP-j6JMdoIuqKycEn7hjg,151
-application_sdk/constants.py,sha256=TvdmKQShVWBNQZdVF2y-fxuE31FmeraTnqQ9jT_n5XY,11567
-application_sdk/version.py,sha256=sNbvXviG7NgxM58lOHKhbZfERat5qAJNr3UZy_toVQs,84
-application_sdk/worker.py,sha256=DLMocpHvvwpdAopyXhxwM7ftaNlKvZMQfkgy1MFyiik,7561
+application_sdk/constants.py,sha256=_JprQls6AurPBHMG8TgGkLT7q4Lj8YxMAabYQx8Heks,12544
+application_sdk/version.py,sha256=Lnq43x1REzO0vMqhztEDG59AS0CKBDL70U5_-QOtwwc,84
+application_sdk/worker.py,sha256=5V_zxkx64moHSkMaZTNnbfstoygMBk7SAOSEy6AZZZM,11552
 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
@@ -15,8 +15,8 @@ application_sdk/activities/metadata_extraction/rest.py,sha256=47DEQpj8HBSa-_TImW
 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
-application_sdk/application/metadata_extraction/sql.py,sha256=rOd06Wodr4GyzupCYxVSCsNcuNar1rJM66ej9vocNHw,8138
+application_sdk/application/__init__.py,sha256=sTgyt4MBDvsQJ2yI0TxzB1Wp68WC-5fBkSExUgtW8LY,12559
+application_sdk/application/metadata_extraction/sql.py,sha256=TmHTLdWLxqtSbVyzfNmnQNe6BUePVIlZxaiQHZXXbyI,9812
 application_sdk/clients/__init__.py,sha256=C9T84J7V6ZumcoWJPAxdd3tqSmbyciaGBJn-CaCCny0,1341
 application_sdk/clients/atlan.py,sha256=l6yV39fr1006SJFwkOTNDQlbSFlHCZQaUPfdUlzdVEg,5053
 application_sdk/clients/atlan_auth.py,sha256=_MykgutI-Ill1t8ERgc1a7QrfaxnrtZjD48FAT-ER9M,8642
@@ -24,13 +24,16 @@ application_sdk/clients/base.py,sha256=TIn3pG89eXUc1XSYf4jk66m1vajWp0WxcCQOOltda
 application_sdk/clients/models.py,sha256=iZOTyH6LO64kozdiUPCFCN0NgLhd_Gtv0lH7ZIPdo8w,1800
 application_sdk/clients/redis.py,sha256=IfAD32vLp88BCvsDTaQtxFHxzHlEx4V7TK7h1HwDDBg,15917
 application_sdk/clients/sql.py,sha256=J43FCxLW2YbnH2MlSm5hCTRFOMOEBtHFqi4ZTTul4JQ,26300
-application_sdk/clients/temporal.py,sha256=7ZkQSwSSZTFkBkhwFlqmLBFxlD2-jLS8QWpSiFTX7V4,20024
+application_sdk/clients/temporal.py,sha256=7xxtzHrZdWjZ5pY1UPoRDY8U9IGpaGycANoWLY-r3io,20238
 application_sdk/clients/utils.py,sha256=zLFOJbTr_6TOqnjfVFGY85OtIXZ4FQy_rquzjaydkbY,779
 application_sdk/clients/workflow.py,sha256=6bSqmA3sNCk9oY68dOjBUDZ9DhNKQxPD75qqE0cfldc,6104
 application_sdk/clients/.cursor/BUGBOT.md,sha256=7nEDUqWBEMI_uU6eK1jCSZGeXoQtLQcKwOrDn8AIDWo,10595
+application_sdk/clients/azure/__init__.py,sha256=koHM4dpsLDf6NQfBWh1axIwF8DGFePcc_nu8_4a-D1g,233
+application_sdk/clients/azure/auth.py,sha256=lvdyNt0V099XlSuTapprNx9xmtpr1El5I_2pKtEnDX0,11129
+application_sdk/clients/azure/client.py,sha256=7f2kanMWgF61xM3WXeFnsgrZaG3Pdl9a1Vt8818e8HQ,13538
 application_sdk/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 application_sdk/common/aws_utils.py,sha256=xlSMIQyjvQ-CydEXaxXrnPUygv7AAbCLsxhZ2wtKnzg,11219
-application_sdk/common/error_codes.py,sha256=bxgvugN_0H5b8VXfJw-44mybgX5I9lRJbRdYjtPjqDI,14561
+application_sdk/common/error_codes.py,sha256=BGdOSfhKe4pd9oM9T2N8GASZeIJU2Be1NcThL8LkKBQ,14900
 application_sdk/common/file_converter.py,sha256=ta0PVh7uIEGJg0BTPUJnSjj55ve2iVAOkqwAeg96_-g,3079
 application_sdk/common/types.py,sha256=qkVr3SAR1zn5_0w9hFt18vGtcnaPLKlMJLSBhTSKePU,134
 application_sdk/common/utils.py,sha256=czcWvqoe2PawDvHOahk_AI88Zqth-CM3KzdGmehHQJ4,19286
@@ -78,8 +81,8 @@ application_sdk/observability/utils.py,sha256=-02GAFom8Bg4SNyCTNYySmen2dzvLfTu43
 application_sdk/observability/decorators/observability_decorator.py,sha256=yd6qfrg1MmH5KcZ5Ydzb0RaBzmxx5FrmiI9qwvZx3EU,8963
 application_sdk/server/__init__.py,sha256=KTqE1YPw_3WDVMWatJUuf9OOiobLM2K5SMaBrI62sCo,1568
 application_sdk/server/.cursor/BUGBOT.md,sha256=p_MMoWUW5G1894WfOKYReZKWCuyJT_OJz3rL5g21NbI,16566
-application_sdk/server/fastapi/__init__.py,sha256=BVqf63z1hxEdpJqLU4LXpFTbk5q8dVkjEbWbu_vbW_Y,29578
-application_sdk/server/fastapi/models.py,sha256=h0hMtMg_p5G0Ug2MBkmBcT94W025VKCLMFyh0FciNoQ,7559
+application_sdk/server/fastapi/__init__.py,sha256=Ce1eXW9Y_8icuEL3RawYfoqcJbjpyfEK2nHq9y-o04k,30941
+application_sdk/server/fastapi/models.py,sha256=ZKz0RWISrJKQ50kepPiPsVqZ6li44I_hFIy5GXochLc,9796
 application_sdk/server/fastapi/utils.py,sha256=WoDGDmq6E1kwS2FN5pjIuzygKNTHpA-tg2SQJZgJOOI,1415
 application_sdk/server/fastapi/middleware/logmiddleware.py,sha256=sKKi-ysI9XQDT0uKW3sfw2r3XQq2iJT4_XrNGXZWdjI,2907
 application_sdk/server/fastapi/middleware/metrics.py,sha256=F_EhayzNgEn9KJvITv8VFxwXhBZVaOQkDptyfa9nQK4,1820
@@ -138,7 +141,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=4uVCU-NfDe08PlffjWQ5p4smQa7c518IL2rDgIk6694,17446
+application_sdk/transformers/query/__init__.py,sha256=asp3IC0zLjrwM8wyUxI8FSFeXmBT5IoFZ23IHn67Muw,18711
 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 +155,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.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,,
+atlan_application_sdk-2.3.0.dist-info/METADATA,sha256=F-6xR-DchbERcVzifQhz_M4jjUhcYke0tZhe3vIwHao,6014
+atlan_application_sdk-2.3.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+atlan_application_sdk-2.3.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+atlan_application_sdk-2.3.0.dist-info/licenses/NOTICE,sha256=A-XVVGt3KOYuuMmvSMIFkg534F1vHiCggEBp4Ez3wGk,1041
+atlan_application_sdk-2.3.0.dist-info/RECORD,,