atlan-application-sdk 0.1.1rc34__py3-none-any.whl → 0.1.1rc35__py3-none-any.whl

application_sdk/services/secretstore.py

@@ -0,0 +1,344 @@
+"""Unified secret store service for the application."""
+
+import collections.abc
+import copy
+import json
+import uuid
+from typing import Any, Dict
+
+from dapr.clients import DaprClient
+
+from application_sdk.common.dapr_utils import is_component_registered
+from application_sdk.common.error_codes import CommonError
+from application_sdk.constants import (
+    DEPLOYMENT_NAME,
+    DEPLOYMENT_SECRET_PATH,
+    DEPLOYMENT_SECRET_STORE_NAME,
+    LOCAL_ENVIRONMENT,
+    SECRET_STORE_NAME,
+)
+from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services.statestore import StateStore, StateType
+
+logger = get_logger(__name__)
+
+
+class SecretStore:
+    """Unified secret store service for handling secret management."""
+
+    @classmethod
+    async def get_credentials(cls, credential_guid: str) -> Dict[str, Any]:
+        """Resolve credentials based on credential source with automatic secret substitution.
+
+        This method retrieves credential configuration from the state store and resolves
+        any secret references by fetching actual values from the secret store.
+
+        Args:
+            credential_guid (str): The unique GUID of the credential configuration to resolve.
+
+        Returns:
+            Dict[str, Any]: Complete credential data with secrets resolved.
+
+        Raises:
+            CommonError: If credential resolution fails due to missing configuration,
+                        secret store errors, or invalid credential format.
+
+        Examples:
+            >>> # Resolve database credentials
+            >>> creds = await SecretStore.get_credentials("db-cred-abc123")
+            >>> print(f"Connecting to {creds['host']}:{creds['port']}")
+            >>> # Password is automatically resolved from secret store
+
+            >>> # Handle resolution errors
+            >>> try:
+            ...     creds = await SecretStore.get_credentials("invalid-guid")
+            >>> except CommonError as e:
+            ...     logger.error(f"Failed to resolve credentials: {e}")
+        """
+
+        async def _get_credentials_async(credential_guid: str) -> Dict[str, Any]:
+            """Async helper function to perform async I/O operations."""
+            credential_config = await StateStore.get_state(
+                credential_guid, StateType.CREDENTIALS
+            )
+
+            # Fetch secret data from secret store
+            secret_key = credential_config.get("secret-path", credential_guid)
+            secret_data = SecretStore.get_secret(secret_key=secret_key)
+
+            # Resolve credentials
+            credential_source = credential_config.get("credentialSource", "direct")
+            if credential_source == "direct":
+                credential_config.update(secret_data)
+                return credential_config
+            else:
+                return cls.resolve_credentials(credential_config, secret_data)
+
+        try:
+            # Run async operations directly
+            return await _get_credentials_async(credential_guid)
+        except Exception as e:
+            logger.error(f"Error resolving credentials: {str(e)}")
+            raise CommonError(
+                CommonError.CREDENTIALS_RESOLUTION_ERROR,
+                f"Failed to resolve credentials: {str(e)}",
+            )
+
+    @classmethod
+    def resolve_credentials(
+        cls, credential_config: Dict[str, Any], secret_data: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Resolve credentials by substituting secret references with actual values.
+
+        This method processes credential configuration and replaces any reference
+        values with corresponding secrets from the secret data.
+
+        Args:
+            credential_config (Dict[str, Any]): Base credential configuration with potential references.
+            secret_data (Dict[str, Any]): Secret data containing actual secret values.
+
+        Returns:
+            Dict[str, Any]: Credential configuration with all secret references resolved.
+
+        Examples:
+            >>> # Basic secret resolution
+            >>> config = {"host": "db.example.com", "password": "db_password_key"}
+            >>> secrets = {"db_password_key": "actual_secret_password"}
+            >>> resolved = SecretStore.resolve_credentials(config, secrets)
+            >>> print(resolved)  # {"host": "db.example.com", "password": "actual_secret_password"}
+
+            >>> # Resolution with nested 'extra' fields
+            >>> config = {
+            ...     "host": "db.example.com",
+            ...     "extra": {"ssl_cert": "cert_key"}
+            ... }
+            >>> secrets = {"cert_key": "-----BEGIN CERTIFICATE-----..."}
+            >>> resolved = SecretStore.resolve_credentials(config, secrets)
+        """
+        credentials = copy.deepcopy(credential_config)
+
+        # Replace values with secret values
+        for key, value in list(credentials.items()):
+            if isinstance(value, str) and value in secret_data:
+                credentials[key] = secret_data[value]
+
+        # Apply the same substitution to the 'extra' dictionary if it exists
+        if "extra" in credentials and isinstance(credentials["extra"], dict):
+            for key, value in list(credentials["extra"].items()):
+                if isinstance(value, str):
+                    if value in secret_data:
+                        credentials["extra"][key] = secret_data[value]
+                    elif value in secret_data.get("extra", {}):
+                        credentials["extra"][key] = secret_data["extra"][value]
+
+        return credentials
+
+    @classmethod
+    def get_deployment_secret(cls) -> Dict[str, Any]:
+        """Get deployment configuration from the deployment secret store.
+
+        Validates that the deployment secret store component is registered
+        before attempting to fetch secrets to prevent errors. This method
+        is commonly used to retrieve environment-specific configuration.
+
+        Returns:
+            Dict[str, Any]: Deployment configuration data, or empty dict if
+                          component is unavailable or fetch fails.
+
+        Examples:
+            >>> # Get deployment configuration
+            >>> config = SecretStore.get_deployment_secret()
+            >>> if config:
+            ...     print(f"Environment: {config.get('environment')}")
+            ...     print(f"Region: {config.get('region')}")
+            >>> else:
+            ...     print("No deployment configuration available")
+
+            >>> # Use in application initialization
+            >>> deployment_config = SecretStore.get_deployment_secret()
+            >>> if deployment_config.get('debug_mode'):
+            ...     logging.getLogger().setLevel(logging.DEBUG)
+        """
+        if not is_component_registered(DEPLOYMENT_SECRET_STORE_NAME):
+            logger.warning(
+                f"Deployment secret store component '{DEPLOYMENT_SECRET_STORE_NAME}' is not registered"
+            )
+            return {}
+
+        try:
+            return cls.get_secret(DEPLOYMENT_SECRET_PATH, DEPLOYMENT_SECRET_STORE_NAME)
+        except Exception as e:
+            logger.error(f"Failed to fetch deployment config: {e}")
+            return {}
+
+    @classmethod
+    def get_secret(
+        cls, secret_key: str, component_name: str = SECRET_STORE_NAME
+    ) -> Dict[str, Any]:
+        """Get secret from the Dapr secret store component.
+
+        Retrieves secret data from the specified Dapr component and processes
+        it into a standardized dictionary format.
+
+        Args:
+            secret_key (str): Key of the secret to fetch from the secret store.
+            component_name (str, optional): Name of the Dapr component to fetch from.
+                Defaults to SECRET_STORE_NAME.
+
+        Returns:
+            Dict[str, Any]: Processed secret data as a dictionary.
+
+        Raises:
+            Exception: If the secret cannot be retrieved from the component.
+
+        Note:
+            In local development environment, returns empty dict to avoid
+            secret store dependency.
+
+        Examples:
+            >>> # Get database credentials
+            >>> db_secret = SecretStore.get_secret("database-credentials")
+            >>> print(f"Host: {db_secret.get('host')}")
+
+            >>> # Get from specific component
+            >>> api_secret = SecretStore.get_secret(
+            ...     "api-keys",
+            ...     component_name="external-secrets"
+            ... )
+        """
+        if DEPLOYMENT_NAME == LOCAL_ENVIRONMENT:
+            return {}
+
+        try:
+            with DaprClient() as client:
+                dapr_secret_object = client.get_secret(
+                    store_name=component_name, key=secret_key
+                )
+                return cls._process_secret_data(dapr_secret_object.secret)
+        except Exception as e:
+            logger.error(
+                f"Failed to fetch secret using component {component_name}: {str(e)}"
+            )
+            raise
+
+    @classmethod
+    def _process_secret_data(cls, secret_data: Any) -> Dict[str, Any]:
+        """Process raw secret data into a standardized dictionary format.
+
+        Args:
+            secret_data: Raw secret data from various sources.
+
+        Returns:
+            Dict[str, Any]: Processed secret data as a dictionary.
+        """
+        # Convert ScalarMapContainer to dict if needed
+        if isinstance(secret_data, collections.abc.Mapping):
+            secret_data = dict(secret_data)
+
+        # If the dict has a single key and its value is a JSON string, parse it
+        if len(secret_data) == 1 and isinstance(next(iter(secret_data.values())), str):
+            try:
+                parsed = json.loads(next(iter(secret_data.values())))
+                if isinstance(parsed, dict):
+                    secret_data = parsed
+            except Exception as e:
+                logger.error(f"Failed to parse secret data: {e}")
+                pass
+
+        return secret_data
+
+    @classmethod
+    def apply_secret_values(
+        cls, source_data: Dict[str, Any], secret_data: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Apply secret values to source data by substituting references.
+
+        This function replaces values in the source data with actual secret values
+        when the source value exists as a key in the secret data. It supports
+        nested structures and preserves the original data structure.
+
+        Args:
+            source_data (Dict[str, Any]): Original data with potential references to secrets.
+            secret_data (Dict[str, Any]): Secret data containing actual secret values.
+
+        Returns:
+            Dict[str, Any]: Deep copy of source data with secret references resolved.
+
+        Examples:
+            >>> # Simple secret substitution
+            >>> source = {
+            ...     "database_url": "postgresql://user:${db_password}@localhost/app",
+            ...     "api_key": "api_key_ref"
+            ... }
+            >>> secrets = {
+            ...     "api_key_ref": "sk-1234567890abcdef",
+            ...     "db_password": "secure_db_password"
+            ... }
+            >>> resolved = SecretStore.apply_secret_values(source, secrets)
+
+            >>> # With nested extra fields
+            >>> source = {
+            ...     "host": "api.example.com",
+            ...     "extra": {"token": "auth_token_ref"}
+            ... }
+            >>> secrets = {"auth_token_ref": "bearer_token_123"}
+            >>> resolved = SecretStore.apply_secret_values(source, secrets)
+        """
+        result_data = copy.deepcopy(source_data)
+
+        # Replace values with secret values
+        for key, value in list(result_data.items()):
+            if isinstance(value, str) and value in secret_data:
+                result_data[key] = secret_data[value]
+
+        # Apply the same substitution to the 'extra' dictionary if it exists
+        if "extra" in result_data and isinstance(result_data["extra"], dict):
+            for key, value in list(result_data["extra"].items()):
+                if isinstance(value, str) and value in secret_data:
+                    result_data["extra"][key] = secret_data[value]
+
+        return result_data
+
+    @classmethod
+    async def save_secret(cls, config: Dict[str, Any]) -> str:
+        """Store credentials in the state store (development environment only).
+
+        This method is designed for development and testing purposes only.
+        In production environments, secrets should be managed through proper
+        secret management systems.
+
+        Args:
+            config (Dict[str, Any]): The credential configuration to store.
+
+        Returns:
+            str: The generated credential GUID that can be used to retrieve the credentials.
+
+        Raises:
+            ValueError: If called in production environment (non-local deployment).
+            Exception: If there's an error storing the credentials.
+
+        Examples:
+            >>> # Development environment only
+            >>> config = {
+            ...     "host": "localhost",
+            ...     "port": 5432,
+            ...     "username": "dev_user",
+            ...     "password": "dev_password",
+            ...     "database": "app_dev"
+            ... }
+            >>> guid = await SecretStore.save_secret(config)
+            >>> print(f"Stored credentials with GUID: {guid}")
+
+            >>> # Later retrieve these credentials
+            >>> retrieved = await SecretStore.get_credentials(guid)
+        """
+        if DEPLOYMENT_NAME == LOCAL_ENVIRONMENT:
+            # NOTE: (development) temporary solution to store the credentials in the state store.
+            # In production, dapr doesn't support creating secrets.
+            credential_guid = str(uuid.uuid4())
+            await StateStore.save_state_object(
+                id=credential_guid, value=config, type=StateType.CREDENTIALS
+            )
+            return credential_guid
+        else:
+            raise ValueError("Storing credentials is not supported in production.")

application_sdk/services/statestore.py

@@ -0,0 +1,267 @@
+"""Unified state store service for the application."""
+
+import json
+import os
+from enum import Enum
+from typing import Any, Dict
+
+from temporalio import activity
+
+from application_sdk.activities.common.utils import get_object_store_prefix
+from application_sdk.constants import (
+    APPLICATION_NAME,
+    STATE_STORE_PATH_TEMPLATE,
+    TEMPORARY_PATH,
+    UPSTREAM_OBJECT_STORE_NAME,
+)
+from application_sdk.observability.logger_adaptor import get_logger
+from application_sdk.services.objectstore import ObjectStore
+
+logger = get_logger(__name__)
+activity.logger = logger
+
+
+class StateType(Enum):
+    WORKFLOWS = "workflows"
+    CREDENTIALS = "credentials"
+
+    @classmethod
+    def is_member(cls, type: str) -> bool:
+        """Check if a string value is a valid StateType member.
+
+        Args:
+            type (str): The string value to check.
+
+        Returns:
+            bool: True if the value is a valid StateType, False otherwise.
+
+        Examples:
+            >>> StateType.is_member("workflows")
+            True
+            >>> StateType.is_member("invalid")
+            False
+        """
+        return type in cls._value2member_map_
+
+
+def build_state_store_path(id: str, state_type: StateType) -> str:
+    """Build the state file path for the given id and type.
+
+    Args:
+        id (str): The unique identifier for the state.
+        state_type (StateType): The type of state (WORKFLOWS or CREDENTIALS).
+
+    Returns:
+        str: The constructed state file path.
+
+    Examples:
+        >>> from application_sdk.services.statestore import build_state_store_path, StateType
+
+        >>> # Workflow state path
+        >>> path = build_state_store_path("workflow-123", StateType.WORKFLOWS)
+        >>> print(path)
+        './local/tmp/persistent-artifacts/apps/appName/workflows/workflow-123/config.json'
+
+        >>> # Credential state path
+        >>> cred_path = build_state_store_path("db-cred-456", StateType.CREDENTIALS)
+        >>> print(cred_path)
+        './local/tmp/persistent-artifacts/apps/appName/credentials/db-cred-456/config.json'
+    """
+    return os.path.join(
+        TEMPORARY_PATH,
+        STATE_STORE_PATH_TEMPLATE.format(
+            application_name=APPLICATION_NAME, state_type=state_type.value, id=id
+        ),
+    )
+
+
+class StateStore:
+    """Unified state store service for handling state management."""
+
+    @classmethod
+    async def get_state(cls, id: str, type: StateType) -> Dict[str, Any]:
+        """Get state from the store.
+
+        Args:
+            id (str): The unique identifier to retrieve the state for.
+            type (StateType): The type of state to retrieve (WORKFLOWS or CREDENTIALS).
+
+        Returns:
+            Dict[str, Any]: The retrieved state data. Returns empty dict if no state found.
+
+        Raises:
+            IOError: If there's an error with the object store operations.
+            Exception: If there's an unexpected error during state retrieval.
+
+        Examples:
+            >>> from application_sdk.services.statestore import StateStore, StateType
+
+            >>> # Get workflow state
+            >>> state = await StateStore.get_state("workflow-123", StateType.WORKFLOWS)
+            >>> print(f"Current status: {state.get('status', 'unknown')}")
+
+            >>> # Get credential configuration
+            >>> creds = await StateStore.get_state("db-cred-456", StateType.CREDENTIALS)
+            >>> print(f"Database: {creds.get('database')}")
+        """
+
+        state_file_path = build_state_store_path(id, type)
+        state = {}
+
+        try:
+            logger.info(f"Trying to download state object for {id} with type {type}")
+            await ObjectStore.download_file(
+                source=get_object_store_prefix(state_file_path),
+                destination=state_file_path,
+                store_name=UPSTREAM_OBJECT_STORE_NAME,
+            )
+
+            with open(state_file_path, "r") as file:
+                state = json.load(file)
+
+            logger.info(f"State object downloaded for {id} with type {type}")
+        except Exception as e:
+            # local error message is "file not found", while in object store it is "object not found"
+            if "not found" in str(e).lower():
+                logger.info(
+                    f"No state found for {type.value} with id '{id}', returning empty dict"
+                )
+            else:
+                logger.error(f"Failed to extract state: {str(e)}")
+                raise
+
+        return state
+
+    @classmethod
+    async def save_state(cls, key: str, value: Any, id: str, type: StateType) -> None:
+        """Save a single state value to the store.
+
+        This method updates a specific key within the state object, merging with existing state.
+
+        Args:
+            key (str): The key to store the state value under.
+            value (Any): The value to store (can be any JSON-serializable type).
+            id (str): The unique identifier for the state object.
+            type (StateType): The type of state (WORKFLOWS or CREDENTIALS).
+
+        Raises:
+            Exception: If there's an error with the object store operations.
+
+        Examples:
+            >>> from application_sdk.services.statestore import StateStore, StateType
+
+            >>> # Update workflow progress
+            >>> await StateStore.save_state(
+            ...     key="progress",
+            ...     value=75,
+            ...     id="workflow-123",
+            ...     type=StateType.WORKFLOWS
+            ... )
+
+            >>> # Update workflow status with dict
+            >>> await StateStore.save_state(
+            ...     key="execution_info",
+            ...     value={"started_at": "2024-01-15T10:00:00Z", "worker_id": "worker-1"},
+            ...     id="workflow-123",
+            ...     type=StateType.WORKFLOWS
+            ... )
+        """
+        try:
+            # get the current state from object store
+            current_state = await cls.get_state(id, type)
+            state_file_path = build_state_store_path(id, type)
+
+            # update the state with the new value
+            current_state[key] = value
+
+            os.makedirs(os.path.dirname(state_file_path), exist_ok=True)
+
+            # save the state to a local file
+            with open(state_file_path, "w") as file:
+                json.dump(current_state, file)
+
+            # save the state to the object store
+            await ObjectStore.upload_file(
+                source=state_file_path,
+                destination=get_object_store_prefix(state_file_path),
+                store_name=UPSTREAM_OBJECT_STORE_NAME,
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to store state: {str(e)}")
+            raise e
+
+    @classmethod
+    async def save_state_object(
+        cls, id: str, value: Dict[str, Any], type: StateType
+    ) -> Dict[str, Any]:
+        """Save the entire state object to the object store.
+
+        This method merges the provided value with existing state and saves the complete object.
+
+        Args:
+            id (str): The unique identifier for the state object.
+            value (Dict[str, Any]): The state data to save/merge.
+            type (StateType): The type of state (WORKFLOWS or CREDENTIALS).
+
+        Returns:
+            Dict[str, Any]: The complete updated state after merge.
+
+        Raises:
+            Exception: If there's an error with the object store operations.
+
+        Examples:
+            >>> from application_sdk.services.statestore import StateStore, StateType
+
+            >>> # Save complete workflow state
+            >>> workflow_state = {
+            ...     "status": "running",
+            ...     "current_step": "data_processing",
+            ...     "progress": 50,
+            ...     "config": {"batch_size": 1000}
+            ... }
+            >>> updated = await StateStore.save_state_object(
+            ...     id="workflow-123",
+            ...     value=workflow_state,
+            ...     type=StateType.WORKFLOWS
+            ... )
+            >>> print(f"Final state has {len(updated)} keys")
+
+            >>> # Save credential configuration
+            >>> cred_config = {
+            ...     "credential_type": "database",
+            ...     "host": "db.example.com",
+            ...     "port": 5432
+            ... }
+            >>> await StateStore.save_state_object(
+            ...     id="db-cred-456",
+            ...     value=cred_config,
+            ...     type=StateType.CREDENTIALS
+            ... )
+        """
+        try:
+            logger.info(f"Saving state object for {id} with type {type}")
+            # get the current state from object store
+            current_state = await cls.get_state(id, type)
+            state_file_path = build_state_store_path(id, type)
+
+            # update the state with the new value
+            current_state.update(value)
+
+            os.makedirs(os.path.dirname(state_file_path), exist_ok=True)
+
+            # save the state to a local file
+            with open(state_file_path, "w") as file:
+                json.dump(current_state, file)
+
+            # save the state to the object store
+            await ObjectStore.upload_file(
+                source=state_file_path,
+                destination=get_object_store_prefix(state_file_path),
+                store_name=UPSTREAM_OBJECT_STORE_NAME,
+            )
+            logger.info(f"State object saved for {id} with type {type}")
+            return current_state
+        except Exception as e:
+            logger.error(f"Failed to store state: {str(e)}")
+            raise e

application_sdk/version.py

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

application_sdk/worker.py

@@ -22,7 +22,7 @@ from application_sdk.events.models import (
     WorkerStartEventData,
 )
 from application_sdk.observability.logger_adaptor import get_logger
-from application_sdk.outputs.eventstore import EventStore
+from application_sdk.services.eventstore import EventStore
 
 logger = get_logger(__name__)
 

{atlan_application_sdk-0.1.1rc34.dist-info → atlan_application_sdk-0.1.1rc35.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atlan-application-sdk
-Version: 0.1.1rc34
+Version: 0.1.1rc35
 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