nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/public_api.py

@@ -0,0 +1,199 @@
+"""
+Public API for NL2SQL Core Package
+
+This module provides a clean, stable public interface to the NL2SQL core functionality.
+It defines the official API boundaries and ensures backward compatibility.
+"""
+
+from __future__ import annotations
+
+import pathlib
+from typing import Optional, Union
+from dataclasses import dataclass
+
+from nl2sql.context import NL2SQLContext
+from nl2sql.api.query_api import QueryAPI, QueryResult
+from nl2sql.api.datasource_api import DatasourceAPI
+from nl2sql.api.llm_api import LLM_API
+from nl2sql.api.indexing_api import IndexingAPI
+from nl2sql.api.auth_api import AuthAPI
+from nl2sql.api.settings_api import SettingsAPI
+from nl2sql.api.result_api import ResultAPI
+from nl2sql.api.policy_api import PolicyAPI
+from nl2sql.api.benchmark_api import BenchmarkAPI
+
+
+class NL2SQL:
+    """
+    Public API for NL2SQL Core Package
+
+    This class provides a clean, stable interface to the NL2SQL engine functionality.
+    It abstracts away the internal implementation details and provides a consistent
+    API for external consumers.
+    """
+
+    def __init__(
+        self,
+        ds_config_path: Optional[Union[str, pathlib.Path]] = None,
+        secrets_config_path: Optional[Union[str, pathlib.Path]] = None,
+        llm_config_path: Optional[Union[str, pathlib.Path]] = None,
+        vector_store_path: Optional[Union[str, pathlib.Path]] = None,
+        policies_config_path: Optional[Union[str, pathlib.Path]] = None,
+    ):
+        """
+        Initialize the NL2SQL engine with optional configuration paths.
+
+        Args:
+            ds_config_path: Path to datasource configuration file
+            secrets_config_path: Path to secrets configuration file
+            llm_config_path: Path to LLM configuration file
+            vector_store_path: Path to vector store directory
+            policies_config_path: Path to policies configuration file
+        """
+        ds_config_path = pathlib.Path(ds_config_path) if ds_config_path else None
+        secrets_config_path = pathlib.Path(secrets_config_path) if secrets_config_path else None
+        llm_config_path = pathlib.Path(llm_config_path) if llm_config_path else None
+        vector_store_path = pathlib.Path(vector_store_path) if vector_store_path else None
+        policies_config_path = pathlib.Path(policies_config_path) if policies_config_path else None
+
+
+        self._ctx = NL2SQLContext(
+            ds_config_path=ds_config_path,
+            secrets_config_path=secrets_config_path,
+            llm_config_path=llm_config_path,
+            vector_store_path=vector_store_path,
+            policies_config_path=policies_config_path,
+        )
+
+        # Initialize modular APIs
+        self.query = QueryAPI(self._ctx)
+        self.datasource = DatasourceAPI(self._ctx)
+        self.llm = LLM_API(self._ctx)
+        self.indexing = IndexingAPI(self._ctx)
+        self.auth = AuthAPI(self._ctx)
+        self.settings = SettingsAPI(self._ctx)
+        self.results = ResultAPI(self._ctx)
+        self.policy = PolicyAPI(self._ctx)
+        self.benchmark = BenchmarkAPI(self._ctx)
+
+    @property
+    def context(self) -> NL2SQLContext:
+        """Access to the underlying context (internal use only)."""
+        return self._ctx
+
+    # Convenience methods that delegate to the modular APIs
+    def run_query(
+        self,
+        natural_language: str,
+        datasource_id: Optional[str] = None,
+        execute: bool = True,
+        user_context=None,
+    ):
+        """
+        Execute a natural language query against the database.
+        """
+        return self.query.run_query(
+            natural_language=natural_language,
+            datasource_id=datasource_id,
+            execute=execute,
+            user_context=user_context
+        )
+
+    def add_datasource(self, config):
+        """
+        Programmatically add a datasource to the engine.
+        """
+        return self.datasource.add_datasource(config)
+
+    def add_datasource_from_config(self, config_path: Union[str, pathlib.Path]):
+        """
+        Add datasources from a configuration file.
+        """
+        return self.datasource.add_datasource_from_config(config_path)
+
+    def list_datasources(self) -> list:
+        """
+        List all registered datasource IDs.
+        """
+        return self.datasource.list_datasources()
+    
+    def get_datasource_capabilities(self, datasource_id: str) -> dict:
+        """
+        Get capabilities of a specific datasource.
+        """
+        return self.datasource.get_capabilities(datasource_id)
+
+
+    def configure_llm(self, config):
+        """
+        Programmatically configure an LLM.
+        """
+        return self.llm.configure_llm(config)
+
+    def configure_llm_from_config(self, config_path: Union[str, pathlib.Path]):
+        """
+        Configure LLMs from a configuration file.
+        """
+        return self.llm.configure_llm_from_config(config_path)
+
+
+    def list_llms(self) -> dict:
+        """
+        List all configured LLMs.
+        """
+        return self.llm.list_llms()
+    
+    def get_llm(self, llm_name: str) -> dict:
+        """
+        Get details of a specific LLM.
+        """
+        return self.llm.get_llm(llm_name)
+    
+    def index_datasource(self, datasource_id: str):
+        """
+        Index schema for a specific datasource.
+        """
+        return self.indexing.index_datasource(datasource_id)
+
+    def index_all_datasources(self):
+        """
+        Index schema for all registered datasources.
+        """
+        return self.indexing.index_all_datasources()
+
+    def clear_index(self):
+        """
+        Clear the vector store index.
+        """
+        return self.indexing.clear_index()
+
+    # Auth API convenience methods
+    def check_permissions(self, user_context, datasource_id, table):
+        """
+        Check if a user has permission to access a specific resource.
+        """
+        return self.auth.check_permissions(user_context, datasource_id, table)
+
+    def get_allowed_resources(self, user_context):
+        """
+        Get resources a user has access to.
+        """
+        return self.auth.get_allowed_resources(user_context)
+
+    def get_current_settings(self):
+        """
+        Get the current application settings.
+        """
+        return self.settings.get_current_settings()
+
+    def get_setting(self, key):
+        """
+        Get a specific setting value.
+        """
+        return self.settings.get_setting(key)
+
+    def validate_configuration(self):
+        """
+        Validate the current configuration.
+        """
+        return self.settings.validate_configuration()

nl2sql/schema/__init__.py

@@ -0,0 +1,37 @@
+"""Core schema models and stores."""
+
+from nl2sql_adapter_sdk.schema import (
+    TableRef,
+    ColumnStatistics,
+    ColumnMetadata,
+    ColumnContract,
+    ForeignKeyContract,
+    TableContract,
+    TableMetadata,
+    SchemaContract,
+    SchemaMetadata,
+    SchemaSnapshot,
+)
+from .protocol import SchemaStore
+from .in_memory_store import SchemaContractStore, SchemaMetadataStore, InMemorySchemaStore
+from .sqlite_store import SqliteSchemaStore
+from .store import build_schema_store
+
+__all__ = [
+    "TableRef",
+    "ColumnStatistics",
+    "ColumnMetadata",
+    "ColumnContract",
+    "ForeignKeyContract",
+    "TableContract",
+    "TableMetadata",
+    "SchemaContract",
+    "SchemaMetadata",
+    "SchemaSnapshot",
+    "SchemaContractStore",
+    "SchemaMetadataStore",
+    "SchemaStore",
+    "InMemorySchemaStore",
+    "SqliteSchemaStore",
+    "build_schema_store",
+]

nl2sql/schema/in_memory_store.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+from collections import OrderedDict, defaultdict
+from datetime import datetime
+import logging
+from typing import Dict, 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 SchemaContractStore:
+    """Registry for schema contracts by datasource/version."""
+
+    def __init__(self, max_versions: int = 3):
+        self._registry: Dict[str, OrderedDict[str, SchemaContract]] = defaultdict(
+            OrderedDict
+        )
+        self._fingerprint_index: Dict[str, Dict[str, str]] = defaultdict(dict)
+        self._max_versions = max_versions
+
+    def register(self, schema: SchemaContract) -> tuple[str, List[str]]:
+        fingerprint = generate_schema_fingerprint(schema)
+        version = self._check_schema_exists(schema.datasource_id, fingerprint)
+        if version:
+            logger.info(
+                "Schema for %s already exists with version %s",
+                schema.datasource_id,
+                version,
+            )
+            return version, []
+
+        ts = datetime.utcnow().strftime("%Y%m%d%H%M%S")
+        schema_version = f"{ts}_{fingerprint[:8]}"
+        logger.info(
+            "Adding new schema for %s with version %s",
+            schema.datasource_id,
+            schema_version,
+        )
+
+        self._registry[schema.datasource_id][schema_version] = schema
+        self._fingerprint_index[schema.datasource_id][fingerprint] = schema_version
+
+        evicted_versions = self._evict_old_versions(schema.datasource_id)
+        return schema_version, evicted_versions
+
+    def get_all_versions(self, datasource_id: str) -> List[str]:
+        if datasource_id not in self._registry:
+            return []
+        return list(self._registry[datasource_id].keys())
+
+    def _check_schema_exists(self, datasource_id: str, fingerprint: str) -> Optional[str]:
+        if datasource_id not in self._fingerprint_index:
+            return None
+        return self._fingerprint_index[datasource_id].get(fingerprint)
+
+    def get(self, datasource_id: str, schema_version: str) -> Optional[SchemaContract]:
+        return self._registry.get(datasource_id, {}).get(schema_version)
+
+    def get_latest(self, datasource_id: str) -> Optional[SchemaContract]:
+        versions = self._registry.get(datasource_id)
+        if not versions:
+            return None
+        return next(reversed(versions.values()))
+
+    def get_latest_version(self, datasource_id: str) -> Optional[str]:
+        versions = self._registry.get(datasource_id)
+        if not versions:
+            return None
+        return next(reversed(versions.keys()))
+
+    def _evict_old_versions(self, datasource_id: str) -> List[str]:
+        versions = self._registry[datasource_id]
+        fp_index = self._fingerprint_index[datasource_id]
+        evicted_versions: List[str] = []
+
+        while len(versions) > self._max_versions:
+            evicted_version, evicted_schema = versions.popitem(last=False)
+            evicted_versions.append(evicted_version)
+
+            evicted_fp = generate_schema_fingerprint(evicted_schema)
+            fp_index.pop(evicted_fp, None)
+
+            logger.info(
+                "Evicted old schema version for %s: %s",
+                datasource_id,
+                evicted_version,
+            )
+
+        return evicted_versions
+
+
+class SchemaMetadataStore:
+    def __init__(self):
+        self._store: Dict[str, Dict[str, SchemaMetadata]] = defaultdict(dict)
+
+    def get(self, datasource_id: str, schema_version: str) -> Optional[SchemaMetadata]:
+        return self._store.get(datasource_id, {}).get(schema_version)
+
+    def register(self, schema_version: str, schema: SchemaMetadata):
+        self._store[schema.datasource_id][schema_version] = schema
+
+    def delete(self, datasource_id: str, schema_version: str):
+        self._store.get(datasource_id, {}).pop(schema_version, None)
+
+
+class InMemorySchemaStore:
+    """In-memory schema store with versioning and per-table access."""
+
+    def __init__(self, max_versions: int = 3):
+        self._contracts = SchemaContractStore(max_versions=max_versions)
+        self._metadata = SchemaMetadataStore()
+
+    def register_snapshot(self, snapshot: SchemaSnapshot) -> Tuple[str, List[str]]:
+        schema_version, evicted_versions = self._contracts.register(snapshot.contract)
+        self._metadata.register(schema_version, snapshot.metadata)
+
+        for evicted_version in evicted_versions:
+            self._metadata.delete(snapshot.contract.datasource_id, evicted_version)
+
+        return schema_version, evicted_versions
+
+    def get_snapshot(
+        self, datasource_id: str, schema_version: str
+    ) -> Optional[SchemaSnapshot]:
+        contract = self._contracts.get(datasource_id, schema_version)
+        metadata = self._metadata.get(datasource_id, schema_version)
+        if not contract or not metadata:
+            return None
+        return SchemaSnapshot(contract=contract, metadata=metadata)
+
+    def get_latest_snapshot(self, datasource_id: str) -> Optional[SchemaSnapshot]:
+        latest_version = self._contracts.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]:
+        return self._contracts.get_latest_version(datasource_id)
+
+    def list_versions(self, datasource_id: str) -> List[str]:
+        return self._contracts.get_all_versions(datasource_id)
+
+    def get_table_contract(
+        self,
+        datasource_id: str,
+        schema_version: str,
+        table_key: str,
+    ) -> Optional[TableContract]:
+        contract = self._contracts.get(datasource_id, schema_version)
+        if not contract:
+            return None
+        return contract.tables.get(table_key)
+
+    def get_table_metadata(
+        self,
+        datasource_id: str,
+        schema_version: str,
+        table_key: str,
+    ) -> Optional[TableMetadata]:
+        metadata = self._metadata.get(datasource_id, schema_version)
+        if not metadata:
+            return None
+        return metadata.tables.get(table_key)

nl2sql/schema/protocol.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import List, Optional, Protocol, Tuple
+
+from nl2sql_adapter_sdk.schema import (
+    SchemaContract,
+    SchemaMetadata,
+    SchemaSnapshot,
+    TableContract,
+    TableMetadata,
+)
+
+
+def generate_schema_fingerprint(schema: SchemaContract) -> str:
+    payload = {
+        "datasource_id": schema.datasource_id,
+        "engine_type": schema.engine_type,
+        "tables": {
+            table_key: {
+                "columns": [
+                    {
+                        "name": c.name,
+                        "type": c.data_type,
+                        "nullable": c.is_nullable,
+                        "pk": c.is_primary_key,
+                    }
+                    for c in sorted(table.columns.values(), key=lambda c: c.name)
+                ],
+                "fks": [
+                    {
+                        "cols": sorted(fk.constrained_columns),
+                        "ref_table": fk.referred_table,
+                        "ref_cols": sorted(fk.referred_columns),
+                    }
+                    for fk in sorted(
+                        table.foreign_keys,
+                        key=lambda fk: (
+                            fk.referred_table,
+                            sorted(fk.constrained_columns),
+                        ),
+                    )
+                ],
+            }
+            for table_key, table in sorted(schema.tables.items())
+        },
+    }
+
+    raw = json.dumps(payload, sort_keys=True, default=str)
+    return hashlib.sha256(raw.encode()).hexdigest()
+
+
+class SchemaStore(Protocol):
+    """Unified interface for schema snapshot storage backends."""
+
+    def register_snapshot(self, snapshot: SchemaSnapshot) -> Tuple[str, List[str]]:
+        ...
+
+    def get_snapshot(
+        self, datasource_id: str, schema_version: str
+    ) -> Optional[SchemaSnapshot]:
+        ...
+
+    def get_latest_snapshot(self, datasource_id: str) -> Optional[SchemaSnapshot]:
+        ...
+
+    def get_latest_version(self, datasource_id: str) -> Optional[str]:
+        ...
+
+    def list_versions(self, datasource_id: str) -> List[str]:
+        ...
+
+    def get_table_contract(
+        self,
+        datasource_id: str,
+        schema_version: str,
+        table_key: str,
+    ) -> Optional[TableContract]:
+        ...
+
+    def get_table_metadata(
+        self,
+        datasource_id: str,
+        schema_version: str,
+        table_key: str,
+    ) -> Optional[TableMetadata]:
+        ...