genesis-flow 1.0.0__py3-none-any.whl

mlflow/azure/connection_factory.py

@@ -0,0 +1,340 @@
+"""Database connection factory with Azure authentication support."""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional, Dict, Any
+from urllib.parse import urlparse, urlencode, parse_qs
+
+from sqlalchemy import create_engine, event, Engine, text
+from sqlalchemy.engine import make_url
+from sqlalchemy.pool import QueuePool
+
+from mlflow.azure.auth_handler import AzureAuthHandler
+from mlflow.azure.config import AzureAuthConfig
+from mlflow.azure.exceptions import ConnectionError, ConfigurationError
+from mlflow.azure.utils import sanitize_connection_string_for_logging
+
+logger = logging.getLogger(__name__)
+
+
+class ConnectionFactory:
+    """Factory for creating database connections with Azure authentication."""
+
+    def __init__(self, config: Optional[AzureAuthConfig] = None):
+        """Initialize the connection factory.
+
+        Args:
+            config: Azure authentication configuration
+        """
+        self.config = config or AzureAuthConfig()
+        self.auth_handler = (
+            AzureAuthHandler(self.config) if self.config.should_use_azure_auth else None
+        )
+
+        logger.info(
+            "Initialized connection factory: azure_auth_enabled=%s, auth_method=%s",
+            self.config.should_use_azure_auth,
+            self.config.auth_method.value if self.config.should_use_azure_auth else "sql_auth",
+        )
+
+    def create_engine(
+        self,
+        database_uri: str,
+        **engine_kwargs: Any,
+    ) -> Engine:
+        """Create a SQLAlchemy engine with Azure authentication support.
+
+        Args:
+            database_uri: Database connection URI
+            **engine_kwargs: Additional engine configuration
+
+        Returns:
+            Configured SQLAlchemy engine
+
+        Raises:
+            ConnectionError: If connection configuration fails
+            ConfigurationError: If the URI is invalid
+        """
+        logger.info(
+            "Creating database engine: connection_info=%s, azure_auth=%s",
+            sanitize_connection_string_for_logging(database_uri),
+            self.config.should_use_azure_auth,
+        )
+
+        try:
+            # Parse and validate the URI
+            parsed_uri = self._parse_and_validate_uri(database_uri)
+
+            # Prepare engine configuration
+            engine_config = self._prepare_engine_config(parsed_uri, **engine_kwargs)
+
+            # Create the engine
+            engine = create_engine(**engine_config)
+
+            # Set up Azure authentication if enabled
+            if self.config.should_use_azure_auth and self.auth_handler:
+                self._setup_azure_auth_events(engine)
+
+            logger.info("Database engine created successfully")
+            return engine
+
+        except Exception as e:
+            logger.error("Failed to create database engine: %s", str(e))
+            raise ConnectionError(f"Failed to create database engine: {e}") from e
+
+    def _parse_and_validate_uri(self, database_uri: str) -> Dict[str, Any]:
+        """Parse and validate database URI.
+
+        Args:
+            database_uri: Database connection URI
+
+        Returns:
+            Parsed URI components
+
+        Raises:
+            ConfigurationError: If URI is invalid
+        """
+        try:
+            # Handle azure-postgres:// scheme by converting to postgresql://
+            if database_uri.startswith("azure-postgres://"):
+                database_uri = database_uri.replace("azure-postgres://", "postgresql://", 1)
+
+            url = make_url(database_uri)
+
+            if url.drivername not in ["postgresql", "postgresql+psycopg2"]:
+                raise ConfigurationError(f"Unsupported database driver: {url.drivername}")
+
+            return {
+                "drivername": url.drivername,
+                "username": url.username,
+                "password": url.password,
+                "host": url.host,
+                "port": url.port or 5432,
+                "database": url.database,
+                "query": dict(url.query),
+            }
+
+        except Exception as e:
+            raise ConfigurationError(f"Invalid database URI: {e}") from e
+
+    def _prepare_engine_config(
+        self,
+        parsed_uri: Dict[str, Any],
+        **engine_kwargs: Any,
+    ) -> Dict[str, Any]:
+        """Prepare SQLAlchemy engine configuration.
+
+        Args:
+            parsed_uri: Parsed database URI components
+            **engine_kwargs: Additional engine configuration
+
+        Returns:
+            Engine configuration dictionary
+        """
+        # Build connection URL
+        url_parts = {
+            "drivername": parsed_uri["drivername"],
+            "username": parsed_uri["username"],
+            "host": parsed_uri["host"],
+            "port": parsed_uri["port"],
+            "database": parsed_uri["database"],
+        }
+
+        # Handle password for non-Azure auth
+        if not self.config.should_use_azure_auth:
+            url_parts["password"] = parsed_uri["password"]
+
+        # Prepare query parameters - remove custom auth parameters
+        query_params = parsed_uri["query"].copy()
+        
+        # Remove custom parameters that psycopg2 doesn't understand
+        query_params.pop("auth_method", None)
+
+        # Ensure SSL for Azure connections
+        if self.config.should_use_azure_auth:
+            query_params.setdefault("sslmode", "require")
+
+        url_parts["query"] = query_params
+
+        # Create URL
+        from sqlalchemy.engine.url import URL
+
+        url = URL.create(**url_parts)
+
+        # Default engine configuration
+        config = {
+            "url": url,
+            "echo": engine_kwargs.pop("echo", False),
+            "echo_pool": engine_kwargs.pop("echo_pool", False),
+            "poolclass": engine_kwargs.pop("poolclass", QueuePool),
+            "pool_size": engine_kwargs.pop("pool_size", self.config.pool_size),
+            "max_overflow": engine_kwargs.pop("max_overflow", self.config.max_overflow),
+            "pool_recycle": engine_kwargs.pop("pool_recycle", self.config.pool_recycle),
+            "pool_pre_ping": engine_kwargs.pop("pool_pre_ping", self.config.pool_pre_ping),
+            "connect_args": engine_kwargs.pop("connect_args", {}),
+        }
+
+        # Add connection timeout
+        config["connect_args"].setdefault("connect_timeout", self.config.connection_timeout)
+
+        # Add any remaining engine kwargs
+        config.update(engine_kwargs)
+
+        return config
+
+    def _setup_azure_auth_events(self, engine: Engine) -> None:
+        """Set up SQLAlchemy events for Azure authentication.
+
+        Args:
+            engine: SQLAlchemy engine to configure
+        """
+        if not self.auth_handler:
+            return
+
+        @event.listens_for(engine, "do_connect")
+        def do_connect(dialect, conn_rec, cargs, cparams):
+            """Inject Azure token before connecting."""
+            try:
+                # Get fresh token
+                token = self.auth_handler.get_access_token()
+
+                # Inject token as password
+                cparams["password"] = token
+
+                logger.debug("Injected Azure token for database connection")
+
+            except Exception as e:
+                logger.error("Failed to inject Azure token: %s", str(e))
+                raise ConnectionError(f"Failed to inject Azure token: {e}") from e
+
+        @event.listens_for(engine, "connect")
+        def connect(dbapi_connection, connection_record):
+            """Handle post-connection setup."""
+            logger.debug("Database connection established")
+
+        @event.listens_for(engine, "checkout")
+        def checkout(dbapi_connection, connection_record, connection_proxy):
+            """Handle connection checkout from pool."""
+            # Refresh token if needed before using connection
+            if self.auth_handler:
+                try:
+                    self.auth_handler.refresh_token_if_needed()
+                except Exception as e:
+                    logger.warning("Token refresh failed during checkout: %s", str(e))
+
+    def test_connection(self, database_uri: str) -> bool:
+        """Test database connection.
+
+        Args:
+            database_uri: Database connection URI
+
+        Returns:
+            True if connection successful, False otherwise
+        """
+        try:
+            engine = self.create_engine(database_uri)
+
+            with engine.connect() as conn:
+                # Simple test query
+                result = conn.execute(text("SELECT 1"))
+                result.fetchone()
+
+            logger.info("Database connection test successful")
+            return True
+
+        except Exception as e:
+            logger.error("Database connection test failed: %s", str(e))
+            return False
+
+    def create_connection_url(
+        self,
+        host: str,
+        database: str,
+        username: str,
+        port: int = 5432,
+        password: Optional[str] = None,
+        **params: Any,
+    ) -> str:
+        """Create a PostgreSQL connection URL.
+
+        Args:
+            host: Database host
+            database: Database name
+            username: Database username
+            port: Database port
+            password: Database password (optional for Azure auth)
+            **params: Additional connection parameters
+
+        Returns:
+            PostgreSQL connection URL
+        """
+        # Ensure SSL for Azure connections
+        if self.config.should_use_azure_auth:
+            params.setdefault("sslmode", "require")
+
+        # Build URL components
+        url_parts = ["postgresql://"]
+
+        # Add credentials
+        if username:
+            url_parts.append(username)
+            if password and not self.config.should_use_azure_auth:
+                url_parts.append(f":{password}")
+            url_parts.append("@")
+
+        # Add host and port
+        url_parts.append(f"{host}:{port}")
+
+        # Add database
+        if database:
+            url_parts.append(f"/{database}")
+
+        # Add query parameters
+        if params:
+            query_string = urlencode(params)
+            url_parts.append(f"?{query_string}")
+
+        url = "".join(url_parts)
+        logger.debug("Created connection URL: %s", sanitize_connection_string_for_logging(url))
+
+        return url
+
+
+class PoolManager:
+    """Manages connection pools with token refresh capabilities."""
+
+    def __init__(self, connection_factory: ConnectionFactory):
+        self.connection_factory = connection_factory
+        self._engines: Dict[str, Engine] = {}
+
+    def get_engine(self, connection_string: str, **kwargs) -> Engine:
+        """Get or create engine for connection string."""
+        # Use connection string as key (could be enhanced with hashing)
+        key = connection_string
+
+        if key not in self._engines:
+            self._engines[key] = self.connection_factory.create_engine(connection_string, **kwargs)
+
+        return self._engines[key]
+
+    def refresh_connections(self):
+        """Refresh all connection pools to get new tokens."""
+        for engine in self._engines.values():
+            try:
+                # Dispose current connections to force new ones with fresh tokens
+                engine.dispose()
+                logger.info("Refreshed connection pool")
+            except Exception as e:
+                logger.error("Failed to refresh connection pool: %s", str(e))
+
+    def close_all(self):
+        """Close all connection pools."""
+        for engine in self._engines.values():
+            try:
+                engine.dispose()
+            except Exception as e:
+                logger.error("Failed to dispose engine: %s", str(e))
+
+        self._engines.clear()
+        logger.info("Closed all connection pools")

mlflow/azure/exceptions.py

@@ -0,0 +1,27 @@
+"""Custom exceptions for Azure authentication."""
+
+from __future__ import annotations
+
+
+class AzureAuthError(Exception):
+    """Base exception for Azure authentication errors."""
+
+    pass
+
+
+class TokenAcquisitionError(AzureAuthError):
+    """Raised when token acquisition fails."""
+
+    pass
+
+
+class ConnectionError(AzureAuthError):
+    """Raised when database connection fails."""
+
+    pass
+
+
+class ConfigurationError(AzureAuthError):
+    """Raised when configuration is invalid."""
+
+    pass

mlflow/azure/stores.py

@@ -0,0 +1,327 @@
+"""Azure-authenticated PostgreSQL stores for MLflow."""
+
+from __future__ import annotations
+
+import os
+import logging
+from typing import Optional, Any, Dict
+
+from sqlalchemy import Engine
+
+from mlflow.azure.config import AzureAuthConfig
+from mlflow.azure.connection_factory import ConnectionFactory
+from mlflow.azure.exceptions import ConnectionError, ConfigurationError
+
+logger = logging.getLogger(__name__)
+
+
+class AzurePostgreSQLStore:
+    """MLflow tracking store with Azure Managed Identity authentication for PostgreSQL."""
+
+    def __init__(
+        self,
+        db_uri: str,
+        default_artifact_root: Optional[str] = None,
+        config: Optional[AzureAuthConfig] = None,
+    ):
+        """Initialize the Azure PostgreSQL tracking store.
+
+        Args:
+            db_uri: Database connection URI
+            default_artifact_root: Default artifact root path
+            config: Azure authentication configuration
+        """
+        # Import MLflow classes only when needed to avoid circular imports
+        from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
+
+        self.config = config or AzureAuthConfig()
+        self.connection_factory = ConnectionFactory(self.config)
+
+        logger.info(
+            "Initializing Azure PostgreSQL tracking store: azure_auth_enabled=%s, auth_method=%s",
+            self.config.should_use_azure_auth,
+            self.config.auth_method.value,
+        )
+
+        # Create engine with Azure authentication
+        try:
+            engine = self._create_engine(db_uri)
+        except Exception as e:
+            logger.error("Failed to create database engine: %s", str(e))
+            raise ConnectionError(f"Failed to initialize tracking store: {e}") from e
+
+        # Initialize internal SQLAlchemy store
+        try:
+            # Convert azure-postgres:// scheme to postgresql:// for MLflow compatibility
+            # MLflow only recognizes standard database schemes
+            standard_uri = db_uri.replace("azure-postgres://", "postgresql://", 1)
+            
+            # Remove auth_method parameter from URI for MLflow
+            from urllib.parse import urlparse, parse_qs, urlencode, urlunparse
+            parsed = urlparse(standard_uri)
+            query_params = parse_qs(parsed.query)
+            clean_query_params = {k: v for k, v in query_params.items() if k != "auth_method"}
+            clean_query = urlencode(clean_query_params, doseq=True)
+            clean_parsed = parsed._replace(query=clean_query)
+            clean_uri = urlunparse(clean_parsed)
+
+            # Add our Azure-enabled engine to the engine map BEFORE initializing SqlAlchemyStore
+            # This prevents MLflow from creating its own engine and ensures migrations work correctly
+            SqlAlchemyStore._db_uri_sql_alchemy_engine_map[clean_uri] = engine
+
+            # Initialize with the clean PostgreSQL URI so MLflow's validation passes
+            self._store = SqlAlchemyStore(clean_uri, default_artifact_root)
+
+            logger.info("Azure PostgreSQL tracking store initialized successfully")
+
+        except Exception as e:
+            logger.error("Failed to initialize SQLAlchemy store: %s", str(e))
+            raise
+
+    def __getattr__(self, name):
+        """Delegate all method calls to the internal SQLAlchemy store."""
+        return getattr(self._store, name)
+
+    def _create_engine(self, db_uri: str) -> Engine:
+        """Create database engine with Azure authentication.
+
+        Args:
+            db_uri: Database connection URI
+
+        Returns:
+            Configured SQLAlchemy engine
+        """
+        # Use our connection factory to create the engine
+        return self.connection_factory.create_engine(
+            db_uri,
+            echo=self.config.enable_debug_logging,
+        )
+
+    @classmethod
+    def from_config(
+        cls,
+        db_uri: Optional[str] = None,
+        default_artifact_root: Optional[str] = None,
+        config: Optional[AzureAuthConfig] = None,
+    ) -> "AzurePostgreSQLStore":
+        """Create store from configuration.
+
+        Args:
+            db_uri: Database URI (defaults to MLFLOW_BACKEND_STORE_URI)
+            default_artifact_root: Artifact root (defaults to MLFLOW_DEFAULT_ARTIFACT_ROOT)
+            config: Azure authentication configuration
+
+        Returns:
+            Configured store instance
+        """
+        # Get URI from environment if not provided
+        if not db_uri:
+            db_uri = os.getenv("MLFLOW_BACKEND_STORE_URI")
+            if not db_uri:
+                raise ConfigurationError(
+                    "Database URI must be provided or set in MLFLOW_BACKEND_STORE_URI"
+                )
+
+        # Get artifact root from environment if not provided
+        if not default_artifact_root:
+            default_artifact_root = os.getenv("MLFLOW_DEFAULT_ARTIFACT_ROOT")
+
+        # Create config if not provided
+        if not config:
+            config = AzureAuthConfig()
+
+        return cls(db_uri, default_artifact_root, config)
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit with cleanup."""
+        try:
+            if hasattr(self, "engine") and self.engine:
+                self.engine.dispose()
+                logger.debug("Database engine disposed")
+        except Exception as e:
+            logger.warning("Error during cleanup: %s", str(e))
+
+
+def create_store(store_uri: str, artifact_uri: Optional[str] = None):
+    """Factory function for creating tracking store with Azure authentication.
+
+    This function is called by MLflow's plugin system and automatically detects
+    if Azure authentication should be used based on environment variables.
+
+    Args:
+        store_uri: Database connection URI
+        artifact_uri: Artifact storage URI
+
+    Returns:
+        Azure-enabled tracking store if Azure auth is enabled, standard store otherwise
+    """
+    # Check if Azure authentication should be used
+    config = AzureAuthConfig()
+    
+    # Auto-detect Azure PostgreSQL and enable auth if credentials are available
+    is_azure_postgres = (
+        store_uri.startswith("azure-postgres://") or
+        ".postgres.database.azure.com" in store_uri or
+        "auth_method=managed_identity" in store_uri
+    )
+    
+    # For Azure PostgreSQL, enable auth if we have any Azure credentials
+    if is_azure_postgres and not config.auth_enabled:
+        # Enable auth automatically if we have Azure credentials
+        if config.client_id or config.client_secret or os.getenv("AZURE_CLIENT_ID"):
+            config.auth_enabled = True
+            if not config.auth_method or config.auth_method == config.auth_method.SQL_AUTH:
+                # Default to service principal if we have credentials, managed identity otherwise  
+                config.auth_method = (
+                    config.auth_method.SERVICE_PRINCIPAL 
+                    if config.client_secret or os.getenv("AZURE_CLIENT_SECRET") 
+                    else config.auth_method.MANAGED_IDENTITY
+                )
+    
+    if config.should_use_azure_auth:
+        logger.info("Creating Azure-enabled tracking store via plugin: auth_method=%s", config.auth_method.value)
+        
+        try:
+            # Import MLflow store
+            from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
+
+            # Ensure we have a postgresql:// scheme
+            standard_uri = store_uri.replace("azure-postgres://", "postgresql://", 1)
+            
+            # Clean the URI
+            from urllib.parse import urlparse, parse_qs, urlencode, urlunparse
+            parsed = urlparse(standard_uri)
+            query_params = parse_qs(parsed.query)
+            clean_query_params = {k: v for k, v in query_params.items() if k != "auth_method"}
+            clean_query = urlencode(clean_query_params, doseq=True)
+            clean_parsed = parsed._replace(query=clean_query)
+            clean_uri = urlunparse(clean_parsed)
+
+            # Check if we already have this engine cached
+            if clean_uri not in SqlAlchemyStore._db_uri_sql_alchemy_engine_map:
+                # Create our Azure-enabled engine and cache it
+                connection_factory = ConnectionFactory(config)
+                engine = connection_factory.create_engine(standard_uri)
+                SqlAlchemyStore._db_uri_sql_alchemy_engine_map[clean_uri] = engine
+                logger.debug("Cached new Azure-enabled engine")
+            else:
+                logger.debug("Reusing cached Azure-enabled engine")
+
+            # Create the standard MLflow store - it will use our cached engine
+            # Ensure artifact_uri is not None or empty to avoid FileNotFoundError
+            logger.debug(f"Raw artifact_uri: {repr(artifact_uri)}")
+            safe_artifact_uri = artifact_uri if (artifact_uri and artifact_uri.strip()) else "./mlflow-artifacts"
+            logger.debug(f"Safe artifact_uri: {repr(safe_artifact_uri)}")
+            return SqlAlchemyStore(clean_uri, safe_artifact_uri)
+
+        except Exception as e:
+            logger.error("Failed to create Azure-enabled store via plugin: %s", str(e))
+            raise
+    else:
+        # Azure auth not enabled, fall back to standard MLflow behavior
+        logger.info("Azure authentication not enabled, using standard tracking store")
+        from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
+        # Ensure artifact_uri is not None or empty to avoid FileNotFoundError
+        logger.debug(f"Fallback raw artifact_uri: {repr(artifact_uri)}")
+        safe_artifact_uri = artifact_uri if (artifact_uri and artifact_uri.strip()) else "./mlflow-artifacts"
+        logger.debug(f"Fallback safe artifact_uri: {repr(safe_artifact_uri)}")
+        return SqlAlchemyStore(store_uri, safe_artifact_uri)
+
+
+# For backward compatibility and testing
+def get_store(store_uri: str, artifact_uri: Optional[str] = None) -> AzurePostgreSQLStore:
+    """Get a configured Azure PostgreSQL tracking store.
+
+    Args:
+        store_uri: Database connection URI
+        artifact_uri: Artifact storage URI
+
+    Returns:
+        Configured store instance
+    """
+    return create_store(store_uri, artifact_uri)
+
+
+# Utility functions for manual store creation
+def get_azure_tracking_store(
+    db_uri: str, artifact_uri: Optional[str] = None, config: Optional[AzureAuthConfig] = None
+) -> AzurePostgreSQLStore:
+    """
+    Create a tracking store with specific configuration.
+
+    Args:
+        db_uri: Database connection URI
+        artifact_uri: Artifact storage URI
+        config: Authentication configuration
+
+    Returns:
+        Configured AzurePostgreSQLStore
+    """
+    if config:
+        # Temporarily override global config
+        import os
+
+        original_env = {}
+
+        try:
+            # Save original environment
+            for key in [
+                "MLFLOW_AZURE_AUTH_ENABLED",
+                "MLFLOW_AZURE_AUTH_METHOD",
+                "AZURE_CLIENT_ID",
+                "AZURE_TENANT_ID",
+            ]:
+                if key in os.environ:
+                    original_env[key] = os.environ[key]
+
+            # Set new environment
+            os.environ["MLFLOW_AZURE_AUTH_ENABLED"] = str(config.auth_enabled).lower()
+            os.environ["MLFLOW_AZURE_AUTH_METHOD"] = config.auth_method.value
+            if config.client_id:
+                os.environ["AZURE_CLIENT_ID"] = config.client_id
+            if config.tenant_id:
+                os.environ["AZURE_TENANT_ID"] = config.tenant_id
+
+            # Create store with new config
+            store = AzurePostgreSQLStore(db_uri, artifact_uri)
+
+        finally:
+            # Restore original environment
+            for key, value in original_env.items():
+                os.environ[key] = value
+
+            # Remove keys that weren't originally set
+            for key in [
+                "MLFLOW_AZURE_AUTH_ENABLED",
+                "MLFLOW_AZURE_AUTH_METHOD",
+                "AZURE_CLIENT_ID",
+                "AZURE_TENANT_ID",
+            ]:
+                if key not in original_env and key in os.environ:
+                    del os.environ[key]
+
+        return store
+
+    else:
+        return AzurePostgreSQLStore(db_uri, artifact_uri)
+
+
+def test_azure_connection(db_uri: str, config: Optional[AzureAuthConfig] = None) -> bool:
+    """
+    Test Azure database connection.
+
+    Args:
+        db_uri: Database connection URI
+        config: Authentication configuration
+
+    Returns:
+        True if connection successful, False otherwise
+    """
+    test_config = config or AzureAuthConfig()
+    connection_factory = ConnectionFactory(test_config)
+
+    return connection_factory.test_connection(db_uri)