removing statestore dependency

Author: nishantmunjal7

application_sdk/activities/metadata_extraction/sql.py

@@ -136,11 +136,19 @@ async def _set_state(self, workflow_args: Dict[str, Any]):
         handler = self.handler_class(sql_client)
         self._state[workflow_id].handler = handler
 
+        # Handle credentials - support both StateStore and direct approaches
         if "credential_guid" in workflow_args:
+            # StateStore approach - load credentials from SecretStore
             credentials = SecretStoreInput.extract_credentials(
                 workflow_args["credential_guid"]
             )
             await sql_client.load(credentials)
+        elif "credentials" in workflow_args:
+            # Direct approach - use credentials directly from workflow_args
+            credentials = workflow_args["credentials"]
+            await sql_client.load(credentials)
+        else:
+            logger.warning("No credentials found in workflow_args")
 
         self._state[workflow_id].sql_client = sql_client
 

application_sdk/clients/temporal.py

@@ -23,6 +23,7 @@
 from application_sdk.clients.workflow import WorkflowClient
 from application_sdk.constants import (
     APPLICATION_NAME,
+    DEPLOYMENT_NAME,
     MAX_CONCURRENT_ACTIVITIES,
     WORKFLOW_AUTH_ENABLED,
     WORKFLOW_HOST,
@@ -265,13 +266,16 @@ def __init__(
     def get_worker_task_queue(self) -> str:
         """Get the worker task queue name.
 
-        The task queue name is derived from the application name and is used
-        to route workflow tasks to appropriate workers.
+        The task queue name is derived from the application name and deployment name
+        and is used to route workflow tasks to appropriate workers.
 
         Returns:
-            str: The task queue name, which is the same as the application name.
+            str: The task queue name in format "app_name-deployment_name".
         """
-        return self.application_name
+        if DEPLOYMENT_NAME:
+            return f"{self.application_name}-{DEPLOYMENT_NAME}"
+        else:
+            return self.application_name
 
     def get_connection_string(self) -> str:
         """Get the Temporal server connection string.
@@ -307,10 +311,13 @@ async def load(self) -> None:
             ConnectionError: If connection to the Temporal server fails.
             ValueError: If authentication is enabled but credentials are missing.
         """
+        # Set TLS based on host - disable for localhost/127.0.0.1, enable for others
+        tls_enabled = self.host not in ["127.0.0.1", "localhost"]
+
         connection_options: Dict[str, Any] = {
             "target_host": self.get_connection_string(),
             "namespace": self.namespace,
-            "tls": True,
+            "tls": tls_enabled,
         }
 
         if self.auth_enabled:
@@ -355,12 +362,24 @@ async def start_workflow(
             WorkflowFailureError: If the workflow fails to start.
             ValueError: If the client is not loaded.
         """
+        # Check if credentials should be stored based on credentialSource
+        should_store_credentials = False
         if "credentials" in workflow_args:
-            # remove credentials from workflow_args and add reference to credentials
-            workflow_args["credential_guid"] = SecretStoreOutput.store_credentials(
-                workflow_args["credentials"]
-            )
-            del workflow_args["credentials"]
+            credential_source = workflow_args["credentials"].get("credentialSource", "")
+            should_store_credentials = credential_source == "direct"
+
+            if should_store_credentials:
+                # Only store credentials if credentialSource is "direct"
+                workflow_args["credential_guid"] = SecretStoreOutput.store_credentials(
+                    workflow_args["credentials"]
+                )
+                del workflow_args["credentials"]
+            else:
+                # For non-direct credential sources, keep credentials in workflow_args
+                # but don't store them in SecretStore
+                logger.info(
+                    f"Skipping credential storage for credentialSource: {credential_source}"
+                )
 
         workflow_id = workflow_args.get("workflow_id")
         output_prefix = workflow_args.get("output_prefix", "/tmp/output")
@@ -376,19 +395,36 @@ async def start_workflow(
                 }
             )
 
+        # Determine configuration approach based on credential storage
+        if should_store_credentials:
+            # StateStore approach - store configuration and pass only workflow_id with flag
             StateStoreOutput.store_configuration(workflow_id, workflow_args)
-
-            logger.info(f"Created workflow config with ID: {workflow_id}")
+            args = [{"workflow_id": workflow_id, "_use_statestore": True}]
+            logger.info(
+                f"Created workflow config with ID: {workflow_id} (StateStore approach)"
+            )
+        else:
+            # Direct approach - pass full configuration with flag
+            workflow_args["_use_statestore"] = False
+            args = [workflow_args]
+            logger.info(f"Created workflow with ID: {workflow_id} (direct approach)")
 
         try:
-            # Pass the full workflow_args to the workflow
+            # Get task_queue from workflow_args or credentials or use default
+            task_queue = (
+                workflow_args.get("task_queue")
+                or workflow_args.get("credentials", {}).get("task_queue")
+                or self.worker_task_queue
+            )
+
+            # Pass the conditional args to the workflow
             if not self.client:
                 raise ValueError("Client is not loaded")
             handle = await self.client.start_workflow(
                 workflow_class,  # type: ignore
-                args=[{"workflow_id": workflow_id}],
+                args=args,
                 id=workflow_id,
-                task_queue=self.worker_task_queue,
+                task_queue=task_queue,
                 cron_schedule=workflow_args.get("cron_schedule", ""),
                 execution_timeout=WORKFLOW_MAX_TIMEOUT_HOURS,
             )

application_sdk/constants.py

@@ -31,7 +31,7 @@
 #: Name of the application, used for identification
 APPLICATION_NAME = os.getenv("ATLAN_APPLICATION_NAME", "default")
 #: Name of the deployment, used to distinguish between different deployments of the same application
-DEPLOYMENT_NAME = os.getenv("ATLAN_DEPLOYMENT_NAME", "")
+DEPLOYMENT_NAME = os.getenv("ATLAN_DEPLOYMENT_NAME", "atlan")
 #: Host address for the application's HTTP server
 APP_HOST = str(os.getenv("ATLAN_APP_HTTP_HOST", "localhost"))
 #: Port number for the application's HTTP server

application_sdk/workflows/__init__.py

@@ -64,23 +64,32 @@ async def run(self, workflow_config: Dict[str, Any]) -> None:
         """Run the workflow with the given configuration.
 
         This method provides the base implementation for workflow execution. It:
-        1. Extracts workflow configuration from the state store
+        1. Extracts workflow configuration from the state store or uses direct configuration
         2. Sets up workflow run ID and retry policy
         3. Executes the preflight check activity
 
         Args:
-            workflow_config (Dict[str, Any]): Includes workflow_id and other parameters
-                workflow_id is used to extract the workflow configuration from the
-                state store.
+            workflow_config (Dict[str, Any]): Either includes workflow_id and _use_statestore flag
+                for StateStore approach, or contains the complete workflow configuration for
+                direct approach.
         """
-        # Get the workflow configuration from the state store
-        workflow_args: Dict[str, Any] = await workflow.execute_activity_method(
-            self.activities_cls.get_workflow_args,
-            workflow_config,  # Pass the whole config containing workflow_id
-            retry_policy=RetryPolicy(maximum_attempts=3, backoff_coefficient=2),
-            start_to_close_timeout=self.default_start_to_close_timeout,
-            heartbeat_timeout=self.default_heartbeat_timeout,
-        )
+        # Determine configuration approach based on _use_statestore flag
+        use_statestore = workflow_config.get("_use_statestore", False)
+
+        if use_statestore:
+            # StateStore approach - retrieve workflow configuration from state store
+            workflow_args: Dict[str, Any] = await workflow.execute_activity_method(
+                self.activities_cls.get_workflow_args,
+                workflow_config,  # Pass the whole config containing workflow_id
+                retry_policy=RetryPolicy(maximum_attempts=3, backoff_coefficient=2),
+                start_to_close_timeout=self.default_start_to_close_timeout,
+                heartbeat_timeout=self.default_heartbeat_timeout,
+            )
+            logger.info("Retrieved workflow configuration from StateStore")
+        else:
+            # Direct approach - use passed configuration directly
+            workflow_args = workflow_config
+            logger.info("Using workflow configuration passed directly")
 
         logger.info("Starting workflow execution")
 

application_sdk/workflows/metadata_extraction/sql.py

@@ -190,16 +190,25 @@ async def run(self, workflow_config: Dict[str, Any]) -> None:
         workflow_success = False
 
         try:
+            # Let the base workflow handle the hybrid approach and preflight check
             await super().run(workflow_config)
 
-            # Get the workflow configuration from the state store
-            workflow_args: Dict[str, Any] = await workflow.execute_activity_method(
-                self.activities_cls.get_workflow_args,
-                workflow_config,  # Pass the whole config containing workflow_id
-                retry_policy=RetryPolicy(maximum_attempts=3, backoff_coefficient=2),
-                start_to_close_timeout=self.default_start_to_close_timeout,
-                heartbeat_timeout=self.default_heartbeat_timeout,
-            )
+            # For the direct approach, use workflow_config as workflow_args
+            # For the StateStore approach, we need to retrieve it again
+            use_statestore = workflow_config.get("_use_statestore", False)
+
+            if use_statestore:
+                # StateStore approach - retrieve workflow configuration from state store
+                workflow_args: Dict[str, Any] = await workflow.execute_activity_method(
+                    self.activities_cls.get_workflow_args,
+                    workflow_config,
+                    retry_policy=RetryPolicy(maximum_attempts=3, backoff_coefficient=2),
+                    start_to_close_timeout=self.default_start_to_close_timeout,
+                    heartbeat_timeout=self.default_heartbeat_timeout,
+                )
+            else:
+                # Direct approach - use passed configuration directly
+                workflow_args = workflow_config
 
             workflow_run_id = workflow.info().run_id
             workflow_args["workflow_run_id"] = workflow_run_id