nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/schema/sqlite_store.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+from datetime import datetime
+import json
+import logging
+from pathlib import Path
+import sqlite3
+from typing import List, Optional, Tuple
+
+from nl2sql_adapter_sdk.schema import (
+    SchemaContract,
+    SchemaMetadata,
+    SchemaSnapshot,
+    TableContract,
+    TableMetadata,
+)
+
+from .protocol import generate_schema_fingerprint
+
+logger = logging.getLogger(__name__)
+
+
+class SqliteSchemaStore:
+    """SQLite-backed schema store with versioning and per-table access."""
+
+    def __init__(self, path: Path, max_versions: int = 3):
+        self._path = path
+        self._max_versions = max_versions
+        self._connection = self._connect()
+        self._initialize_schema()
+
+    def _connect(self) -> sqlite3.Connection:
+        if str(self._path) != ":memory:":
+            self._path.parent.mkdir(parents=True, exist_ok=True)
+        connection = sqlite3.connect(str(self._path), check_same_thread=False)
+        connection.execute("PRAGMA journal_mode=WAL;")
+        return connection
+
+    def _initialize_schema(self) -> None:
+        cursor = self._connection.cursor()
+        cursor.execute(
+            """
+            CREATE TABLE IF NOT EXISTS schema_snapshots (
+                datasource_id TEXT NOT NULL,
+                schema_version TEXT NOT NULL,
+                fingerprint TEXT NOT NULL,
+                contract_json TEXT NOT NULL,
+                metadata_json TEXT NOT NULL,
+                created_at INTEGER NOT NULL,
+                PRIMARY KEY (datasource_id, schema_version)
+            );
+            """
+        )
+        cursor.execute(
+            """
+            CREATE INDEX IF NOT EXISTS idx_schema_snapshots_fingerprint
+            ON schema_snapshots (datasource_id, fingerprint);
+            """
+        )
+        cursor.execute(
+            """
+            CREATE INDEX IF NOT EXISTS idx_schema_snapshots_created_at
+            ON schema_snapshots (datasource_id, created_at);
+            """
+        )
+        self._connection.commit()
+
+    def register_snapshot(self, snapshot: SchemaSnapshot) -> Tuple[str, List[str]]:
+        fingerprint = generate_schema_fingerprint(snapshot.contract)
+        existing_version = self._get_version_by_fingerprint(
+            snapshot.contract.datasource_id, fingerprint
+        )
+        if existing_version:
+            logger.info(
+                "Schema for %s already exists with version %s",
+                snapshot.contract.datasource_id,
+                existing_version,
+            )
+            return existing_version, []
+
+        now = datetime.utcnow()
+
+        ts = now.strftime("%Y%m%d%H%M%S")
+        schema_version = f"{ts}_{fingerprint[:8]}"
+        created_at = int(now.timestamp())
+
+        contract_json = json.dumps(snapshot.contract.model_dump(mode="json"))
+        metadata_json = json.dumps(snapshot.metadata.model_dump(mode="json"))
+
+        with self._connection:
+            self._connection.execute(
+                """
+                INSERT INTO schema_snapshots (
+                    datasource_id,
+                    schema_version,
+                    fingerprint,
+                    contract_json,
+                    metadata_json,
+                    created_at
+                ) VALUES (?, ?, ?, ?, ?, ?);
+                """,
+                (
+                    snapshot.contract.datasource_id,
+                    schema_version,
+                    fingerprint,
+                    contract_json,
+                    metadata_json,
+                    created_at,
+                ),
+            )
+
+        evicted_versions = self._evict_old_versions(snapshot.contract.datasource_id)
+        return schema_version, evicted_versions
+
+    def get_snapshot(
+        self, datasource_id: str, schema_version: str
+    ) -> Optional[SchemaSnapshot]:
+        row = self._connection.execute(
+            """
+            SELECT contract_json, metadata_json
+            FROM schema_snapshots
+            WHERE datasource_id = ? AND schema_version = ?;
+            """,
+            (datasource_id, schema_version),
+        ).fetchone()
+        if not row:
+            return None
+        contract = SchemaContract.model_validate(json.loads(row[0]))
+        metadata = SchemaMetadata.model_validate(json.loads(row[1]))
+        return SchemaSnapshot(contract=contract, metadata=metadata)
+
+    def get_latest_snapshot(self, datasource_id: str) -> Optional[SchemaSnapshot]:
+        latest_version = self.get_latest_version(datasource_id)
+        if not latest_version:
+            return None
+        return self.get_snapshot(datasource_id, latest_version)
+
+    def get_latest_version(self, datasource_id: str) -> Optional[str]:
+        with self._connection:
+            row = self._connection.execute(
+                """
+                SELECT schema_version
+                FROM schema_snapshots
+                WHERE datasource_id = ?
+                ORDER BY created_at DESC
+                LIMIT 1;
+                """,
+                (datasource_id,),
+            ).fetchone()
+
+            return row[0] if row else None
+
+    def list_versions(self, datasource_id: str) -> List[str]:
+        rows = self._connection.execute(
+            """
+            SELECT schema_version
+            FROM schema_snapshots
+            WHERE datasource_id = ?
+            ORDER BY created_at ASC;
+            """,
+            (datasource_id,),
+        ).fetchall()
+        return [row[0] for row in rows]
+
+    def get_table_contract(
+        self,
+        datasource_id: str,
+        schema_version: str,
+        table_key: str,
+    ) -> Optional[TableContract]:
+        snapshot = self.get_snapshot(datasource_id, schema_version)
+        if not snapshot:
+            return None
+        return snapshot.contract.tables.get(table_key)
+
+    def get_table_metadata(
+        self,
+        datasource_id: str,
+        schema_version: str,
+        table_key: str,
+    ) -> Optional[TableMetadata]:
+        snapshot = self.get_snapshot(datasource_id, schema_version)
+        if not snapshot:
+            return None
+        return snapshot.metadata.tables.get(table_key)
+
+    def _get_version_by_fingerprint(
+        self, datasource_id: str, fingerprint: str
+    ) -> Optional[str]:
+        row = self._connection.execute(
+            """
+            SELECT schema_version
+            FROM schema_snapshots
+            WHERE datasource_id = ? AND fingerprint = ?
+            ORDER BY created_at DESC
+            LIMIT 1;
+            """,
+            (datasource_id, fingerprint),
+        ).fetchone()
+        return row[0] if row else None
+
+    def _evict_old_versions(self, datasource_id: str) -> List[str]:
+        rows = self._connection.execute(
+            """
+            SELECT schema_version
+            FROM schema_snapshots
+            WHERE datasource_id = ?
+            ORDER BY created_at ASC;
+            """,
+            (datasource_id,),
+        ).fetchall()
+        versions = [row[0] for row in rows]
+        if len(versions) <= self._max_versions:
+            return []
+
+        evicted_versions = versions[: len(versions) - self._max_versions]
+        with self._connection:
+            self._connection.executemany(
+                """
+                DELETE FROM schema_snapshots
+                WHERE datasource_id = ? AND schema_version = ?;
+                """,
+                [(datasource_id, version) for version in evicted_versions],
+            )
+
+        for version in evicted_versions:
+            logger.info(
+                "Evicted old schema version for %s: %s",
+                datasource_id,
+                version,
+            )
+
+        return evicted_versions

nl2sql/schema/store.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+from .in_memory_store import InMemorySchemaStore
+from .protocol import SchemaStore
+from .sqlite_store import SqliteSchemaStore
+
+
+SCHEMA_STORE_BACKENDS = {
+    "memory": InMemorySchemaStore,
+    "sqlite": SqliteSchemaStore,
+}
+
+
+def build_schema_store(
+    backend: str,
+    max_versions: int,
+    path: Optional[Path] = None,
+) -> SchemaStore:
+    backend_key = (backend or "sqlite").lower()
+    if backend_key == "memory":
+        return InMemorySchemaStore(max_versions=max_versions)
+    if backend_key == "sqlite":
+        if path is None:
+            raise ValueError("schema_store_path is required for sqlite backend.")
+        return SqliteSchemaStore(path=path, max_versions=max_versions)
+    raise ValueError(f"Unsupported schema store backend: {backend}")

nl2sql/secrets/__init__.py

@@ -0,0 +1,14 @@
+from .manager import SecretManager
+from .interfaces import SecretProvider
+from .models import SecretProviderConfig
+
+from pathlib import Path
+import os
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+secret_manager = SecretManager()
+
+__all__ = ["SecretManager", "SecretProvider", "SecretProviderConfig", "secret_manager"]

nl2sql/secrets/factory.py

@@ -0,0 +1,85 @@
+from typing import Optional, Any, Dict, Type
+import logging
+from .interfaces import SecretProvider
+from nl2sql.configs.secrets import SecretProviderConfig, AwsSecretConfig, AzureSecretConfig, HashiCorpSecretConfig
+
+logger = logging.getLogger(__name__)
+
+class SecretProviderFactory:
+    """Factory for creating SecretProvider instances.
+    
+    Centralizes the logic for:
+    1. Dynamic imports (lazy loading of heavy dependencies).
+    2. Dependency checking (Boto3, Azure SDK, etc.).
+    3. Instantiation with or without explicit config.
+    """
+
+    @staticmethod
+    def create(config: SecretProviderConfig) -> Optional[SecretProvider]:
+        """Creates a provider instance from configuration.
+        
+        Args:
+            config: A fully resolved configuration object (no placeholders).
+        
+        Returns:
+            The instantiated SecretProvider, or None if dependencies are missing.
+        """
+        try:
+            if isinstance(config, AwsSecretConfig):
+                return SecretProviderFactory._create_aws(config)
+            elif isinstance(config, AzureSecretConfig):
+                return SecretProviderFactory._create_azure(config)
+            elif isinstance(config, HashiCorpSecretConfig):
+                return SecretProviderFactory._create_hashi(config)
+            elif config.type == "env":
+                # Env provider is usually bootstrapped, but we support creating it explicitly too
+                from .providers.env import EnvironmentSecretProvider
+                return EnvironmentSecretProvider()
+            else:
+                logger.warning(f"Unknown provider type: {config.type}")
+                return None
+        except ImportError as e:
+            logger.warning(f"Skipping provider '{config.id}' ({config.type}): {e}")
+            return None
+        except Exception as e:
+            logger.error(f"Failed to create provider '{config.id}': {e}")
+            return None
+
+    @staticmethod
+    def _create_aws(config: AwsSecretConfig) -> SecretProvider:
+        try:
+            from .providers.aws import AwsSecretProvider
+        except ImportError:
+             raise ImportError("Missing dependency 'boto3'. Install 'nl2sql-engine[aws]'.")
+        
+        return AwsSecretProvider(
+            region_name=config.region_name,
+            profile_name=config.profile_name
+        )
+
+    @staticmethod
+    def _create_azure(config: AzureSecretConfig) -> SecretProvider:
+        try:
+            from .providers.azure import AzureSecretProvider
+        except ImportError:
+            raise ImportError("Missing dependencies. Install 'nl2sql-engine[azure]'.")
+
+        return AzureSecretProvider(
+            vault_url=config.vault_url,
+            client_id=config.client_id,
+            client_secret=config.client_secret,
+            tenant_id=config.tenant_id
+        )
+
+    @staticmethod
+    def _create_hashi(config: HashiCorpSecretConfig) -> SecretProvider:
+        try:
+            from .providers.hashi import HashiCorpSecretProvider
+        except ImportError:
+            raise ImportError("Missing dependency 'hvac'. Install 'nl2sql-engine[hashicorp]'.")
+            
+        return HashiCorpSecretProvider(
+            url=config.url,
+            token=config.token,
+            mount_point=config.mount_point
+        )

nl2sql/secrets/interfaces.py

@@ -0,0 +1,16 @@
+from typing import Optional, Protocol
+
+class SecretProvider(Protocol):
+    """Protocol for fetching secrets from secure storage."""
+    
+    def get_secret(self, key: str) -> Optional[str]:
+        """
+        Retrieve a secret by its key.
+        
+        Args:
+            key (str): The identifier for the secret (e.g., 'DB_PASSWORD').
+            
+        Returns:
+            Optional[str]: The secret value, or None if not found.
+        """
+        ...

nl2sql/secrets/manager.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import os
+import re
+from typing import Optional, Dict, Type, List, Any
+from .interfaces import SecretProvider
+from .providers.env import EnvironmentSecretProvider
+from .models import SecretProviderConfig
+import logging
+
+logger = logging.getLogger(__name__)
+
+class SecretManager:
+    """Manages secret resolution using registered providers."""
+    
+    def __init__(self):
+        self._providers: Dict[str, SecretProvider] = {
+            "env": EnvironmentSecretProvider()
+        }
+        self._default_provider = "env"
+
+    def register_provider(self, provider_id: str, provider: SecretProvider) -> None:
+        self._providers[provider_id] = provider
+
+    def configure(self, configs: List[SecretProviderConfig]) -> None:
+        """Configures providers from a list of configuration objects.
+        
+        This method implements a Two-Phase Loading strategy:
+        1. Bootstrap: It assumes the 'env' provider is already active.
+        2. Resolution: It resolves configuration values (like client_secret) that may contain
+           secret references (e.g., "${env:VAR}") using the existing providers.
+        3. Registration: It instantiates and registers the new providers.
+        
+        Args:
+            configs: A list of SecretProviderConfig objects.
+        """
+        from .factory import SecretProviderFactory
+
+        for config in configs:
+            try:
+                if config.type == "env":
+                    continue
+                
+                updates = {}
+                for key, value in config.model_dump(exclude={"id", "type"}).items():
+                    if isinstance(value, str) and value.startswith("${") and value.endswith("}"):
+                        updates[key] = self.resolve(value)
+                
+                resolved_config = config.model_copy(update=updates)
+                provider = SecretProviderFactory.create(resolved_config)
+                    
+                if provider:
+                    self.register_provider(config.id, provider)
+                    logger.info(f"Registered secret provider '{config.id}' (type: {config.type})")
+                    
+            except Exception as e:
+                logger.error(f"Failed to configure secret provider '{config.id}': {e}")
+                
+    def resolve(self, secret_ref: str) -> str:
+        """Resolves a secret reference string.
+
+        Format: ${provider_id:key}
+        
+        Args:
+            secret_ref: The reference string to resolve.
+
+        Returns:
+            str: The resolved secret value.
+
+        Raises:
+            ValueError: If the format is invalid, provider is unknown, or secret
+                is not found.
+        """
+        cleaned_ref = secret_ref.replace('${', '').replace('}', '')
+        
+        parts = cleaned_ref.split(':', 1)
+        if len(parts) != 2:
+            raise ValueError(f"Invalid secret format '{secret_ref}'. Expected '${{provider_id:key}}'.")
+            
+        provider_id, key = parts
+        provider = self._providers.get(provider_id)
+        if not provider:
+            raise ValueError(f"Unknown secret provider ID: '{provider_id}'")
+            
+        val = provider.get_secret(key)
+        if val is not None:
+            return val
+        
+        raise ValueError(f"Secret not found: {secret_ref}")
+
+    def resolve_object(self, obj: Any) -> Any:
+        """Recursively resolves secret references in a generic object.
+
+        Traverses the object structure (Pydantic models, dicts, lists) and resolves
+        any string values matching the secret pattern "${...}". Handles SecretStr
+        by unwrapping, resolving, and re-wrapping.
+
+        Args:
+            obj: The object to resolve secrets in. Can be a Pydantic model,
+                dictionary, list, string, or SecretStr.
+
+        Returns:
+            Any: A new object with all secret references resolved to their
+                actual values.
+        """
+        from pydantic import BaseModel, SecretStr
+
+        if isinstance(obj, str):
+            if obj.startswith("${") and obj.endswith("}"):
+                return self.resolve(obj)
+            return obj
+            
+        if isinstance(obj, SecretStr):
+            secret_val = obj.get_secret_value()
+            if secret_val and secret_val.startswith("${") and secret_val.endswith("}"):
+                resolved_val = self.resolve(secret_val)
+                return SecretStr(resolved_val)
+            return obj
+        
+        if isinstance(obj, BaseModel):
+            updates = {}
+            for field_name in type(obj).model_fields.keys():
+                val = getattr(obj, field_name)
+                resolved = self.resolve_object(val)
+                if resolved != val:
+                    updates[field_name] = resolved
+            
+            if updates:
+                return obj.model_copy(update=updates)
+            return obj
+            
+        if isinstance(obj, list):
+            return [self.resolve_object(item) for item in obj]
+            
+        if isinstance(obj, dict):
+             return {k: self.resolve_object(v) for k, v in obj.items()}
+             
+        return obj
+

nl2sql/secrets/models.py

@@ -0,0 +1,56 @@
+from pydantic import BaseModel, Field, SecretStr
+from typing import Literal, Optional, Union
+
+
+class BaseSecretConfig(BaseModel):
+    """Base configuration for all secret providers."""
+    id: str = Field(..., description="Unique identifier for this provider instance. Used as the scheme in secret references (e.g. ${id:key}).")
+    type: str
+
+class AwsSecretConfig(BaseSecretConfig):
+    """Configuration for AWS Secrets Manager.
+    
+    Attributes:
+        type: Must be 'aws'.
+        region_name: AWS Region (e.g. us-east-1). Defaults to env var if None.
+        profile_name: AWS Profile. Defaults to standard boto3 lookup if None.
+    """
+    type: Literal["aws"] = "aws"
+    region_name: Optional[str] = Field(None, description="AWS Region. If None, uses AWS_DEFAULT_REGION env var.")
+    profile_name: Optional[str] = Field(None, description="AWS Profile. If None, uses default profile.")
+
+class AzureSecretConfig(BaseSecretConfig):
+    """Configuration for Azure Key Vault.
+    
+    Attributes:
+        type: Must be 'azure'.
+        vault_url: The URL of the Key Vault.
+        client_id: Service Principal ID.
+        client_secret: Service Principal Secret.
+        tenant_id: Azure Tenant ID.
+    """
+    type: Literal["azure"] = "azure"
+    vault_url: str = Field(..., description="URL of the Key Vault.")
+    
+    client_id: Optional[str] = Field(None, description="Azure Client ID (Service Principal).")
+    client_secret: Optional[SecretStr] = Field(None, description="Azure Client Secret.")
+    tenant_id: Optional[str] = Field(None, description="Azure Tenant ID.")
+
+class HashiCorpSecretConfig(BaseSecretConfig):
+    """Configuration for HashiCorp Vault."""
+    type: Literal["hashi"] = "hashi"
+    url: str = Field(..., description="URL of the HashiCorp Vault server.")
+    token: Optional[SecretStr] = Field(None, description="Vault Token.")
+    mount_point: str = Field("secret", description="Secrets engine mount point.")
+
+class EnvSecretConfig(BaseSecretConfig):
+    """Configuration for Environment Variable Provider (Explicit)."""
+    type: Literal["env"] = "env"
+
+# Polymorphic Union
+SecretProviderConfig = Union[
+    AwsSecretConfig,
+    AzureSecretConfig,
+    HashiCorpSecretConfig,
+    EnvSecretConfig
+]

nl2sql/secrets/providers/aws.py

@@ -0,0 +1,30 @@
+from typing import Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+class AwsSecretProvider:
+    """Fetches secrets from AWS Secrets Manager."""
+    
+    def __init__(self, region_name: Optional[str] = None, profile_name: Optional[str] = None):
+        try:
+            import boto3
+            session = boto3.Session(profile_name=profile_name, region_name=region_name)
+            self.client = session.client('secretsmanager')
+        except ImportError:
+            self.client = None
+            logger.warning("AWS Secret Provider initialized but 'boto3' is missing. Please install 'nl2sql-engine[aws]' to use AWS secrets.")
+
+    def get_secret(self, key: str) -> Optional[str]:
+        if not self.client:
+            raise ImportError("Cannot fetch AWS secret: 'boto3' is not installed.")
+            
+        try:
+            # key is the SecretId
+            response = self.client.get_secret_value(SecretId=key)
+            if 'SecretString' in response:
+                return response['SecretString']
+            return None # Binary secrets not currently supported for connection strings
+        except Exception as e:
+            logger.error(f"Failed to fetch secret '{key}' from AWS: {e}")
+            return None

nl2sql/secrets/providers/azure.py

@@ -0,0 +1,49 @@
+from typing import Optional
+import os
+import logging
+
+logger = logging.getLogger(__name__)
+
+class AzureSecretProvider:
+    """Fetches secrets from Azure Key Vault."""
+    
+    def __init__(self, vault_url: Optional[str] = None, client_id: Optional[str] = None, client_secret: Optional[str] = None, tenant_id: Optional[str] = None):
+        try:
+            from azure.identity import DefaultAzureCredential, ClientSecretCredential
+            from azure.keyvault.secrets import SecretClient
+            
+            self.vault_url = vault_url or os.environ.get("AZURE_KEYVAULT_URL")
+            if not self.vault_url:
+                raise ValueError("Vault URL is required. specific via config or 'AZURE_KEYVAULT_URL'.")
+                
+            if client_id and client_secret and tenant_id:
+                credential = ClientSecretCredential(
+                    tenant_id=tenant_id,
+                    client_id=client_id,
+                    client_secret=client_secret
+                )
+            else:
+                credential = DefaultAzureCredential()
+
+            self.client = SecretClient(vault_url=self.vault_url, credential=credential)
+            self._available = True
+        except ImportError:
+            self.client = None
+            self._available = False
+            logger.warning("Azure Secret Provider initialized but dependencies missing. Install 'nl2sql-engine[azure]'.")
+        except Exception as e:
+            self.client = None
+            self._available = False
+            logger.warning(f"Azure Secret Provider initialization failed: {e}")
+
+    def get_secret(self, key: str) -> Optional[str]:
+        if not self._available or not self.client:
+             raise ImportError("Azure Secret Provider is not available.")
+             
+        try:
+            # key is the Secret Name
+            secret = self.client.get_secret(key)
+            return secret.value
+        except Exception as e:
+            logger.error(f"Failed to fetch secret '{key}' from Azure Key Vault: {e}")
+            return None

nl2sql/secrets/providers/env.py

@@ -0,0 +1,8 @@
+import os
+from typing import Optional
+
+class EnvironmentSecretProvider:
+    """Fetches secrets from environment variables."""
+    
+    def get_secret(self, key: str) -> Optional[str]:
+        return os.environ.get(key)