nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/api/policy_api.py

@@ -0,0 +1,135 @@
+"""
+Policy API for NL2SQL.
+
+Provides public entry points for policy validation and integrity checks.
+"""
+
+from __future__ import annotations
+
+import pathlib
+from dataclasses import dataclass, field
+from typing import List, Optional, Set
+
+from nl2sql.common.settings import settings
+from nl2sql.configs import ConfigManager
+from nl2sql.context import NL2SQLContext
+from nl2sql.datasources import DatasourceRegistry
+from nl2sql.secrets import secret_manager
+
+
+@dataclass
+class PolicyValidationEntry:
+    role: str
+    target: str
+    status: str
+    details: str
+
+
+@dataclass
+class PolicyValidationReport:
+    ok: bool
+    entries: List[PolicyValidationEntry] = field(default_factory=list)
+    errors: List[str] = field(default_factory=list)
+    available_datasources: Set[str] = field(default_factory=set)
+
+
+class PolicyAPI:
+    """
+    API for policy validation and integrity checks.
+    """
+
+    def __init__(self, ctx: Optional[NL2SQLContext] = None):
+        self._ctx = ctx
+
+    def validate_policies(
+        self,
+        policies_path: Optional[pathlib.Path] = None,
+        datasources_path: Optional[pathlib.Path] = None,
+        secrets_path: Optional[pathlib.Path] = None,
+    ) -> PolicyValidationReport:
+        """
+        Validate policy syntax and integrity against defined datasources.
+        """
+        cm = self._ctx.config_manager if self._ctx else ConfigManager()
+
+        policies_path = pathlib.Path(policies_path) if policies_path else pathlib.Path(settings.policies_config_path)
+        datasources_path = pathlib.Path(datasources_path) if datasources_path else pathlib.Path(settings.datasource_config_path)
+        secrets_path = pathlib.Path(secrets_path) if secrets_path else pathlib.Path(settings.secrets_config_path)
+
+        try:
+            policy_cfg = cm.load_policies(policies_path)
+        except Exception as e:
+            return PolicyValidationReport(ok=False, errors=[str(e)])
+
+        try:
+            secret_configs = cm.load_secrets(secrets_path)
+            if secret_configs:
+                secret_manager.configure(secret_configs)
+
+            ds_configs = cm.load_datasources(datasources_path)
+            registry = DatasourceRegistry(secret_manager)
+            registry.register_datasources(ds_configs)
+            available_ds = set(registry.list_ids())
+        except Exception as e:
+            return PolicyValidationReport(ok=False, errors=[str(e)])
+
+        entries: List[PolicyValidationEntry] = []
+        has_errors = False
+
+        for role_id, role_def in policy_cfg.roles.items():
+            for ds in role_def.allowed_datasources:
+                if ds == "*":
+                    entries.append(
+                        PolicyValidationEntry(role=role_id, target="Datasource: *", status="OK", details="Global Access")
+                    )
+                    continue
+                if ds not in available_ds:
+                    entries.append(
+                        PolicyValidationEntry(
+                            role=role_id,
+                            target=f"Datasource: {ds}",
+                            status="MISSING",
+                            details="Datasource not defined in config",
+                        )
+                    )
+                    has_errors = True
+                else:
+                    entries.append(
+                        PolicyValidationEntry(role=role_id, target=f"Datasource: {ds}", status="OK", details="Verified")
+                    )
+
+            for rule in role_def.allowed_tables:
+                if rule == "*":
+                    entries.append(
+                        PolicyValidationEntry(role=role_id, target="Table: *", status="OK", details="Global Access")
+                    )
+                    continue
+
+                parts = rule.split(".")
+                if len(parts) >= 2:
+                    ds_part = parts[0]
+                    if ds_part not in available_ds and ds_part != "*":
+                        entries.append(
+                            PolicyValidationEntry(
+                                role=role_id,
+                                target=f"Table Rule: {rule}",
+                                status="INVALID_DS",
+                                details=f"Datasource '{ds_part}' unknown",
+                            )
+                        )
+                        has_errors = True
+                    else:
+                        entries.append(
+                            PolicyValidationEntry(
+                                role=role_id,
+                                target=f"Table Rule: {rule}",
+                                status="OK",
+                                details="DS Verified",
+                            )
+                        )
+
+        return PolicyValidationReport(
+            ok=not has_errors,
+            entries=entries,
+            available_datasources=available_ds,
+        )

nl2sql/api/query_api.py

@@ -0,0 +1,138 @@
+"""
+Query API for NL2SQL
+
+Provides functionality for executing natural language queries against databases.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, List, Dict, Any
+
+from pydantic import BaseModel, Field
+from nl2sql.context import NL2SQLContext
+from nl2sql.pipeline.runtime import run_with_graph
+from nl2sql.auth import UserContext
+from nl2sql.execution.contracts import ArtifactRef
+
+
+class SubQueryResult(BaseModel):
+    """Represents the result of a sub-query execution."""
+    id: str = Field(default="")
+    intent: str = Field(default="")
+    sql: str = Field(default="")
+    datasource_id: str = Field(default="")
+    schema_version: str = Field(default="")
+
+
+class QueryResult(BaseModel):
+    """Represents the result of a query execution.
+
+    Row data is deliberately not inlined: results live in artifact storage and are
+    addressable through ``artifact_refs``.
+    """
+    sub_queries: List[SubQueryResult] = Field(default_factory=list)
+    final_answer: Optional[Dict[str, Any]] = None
+    errors: List[Dict[str, Any]] = Field(default_factory=list)
+    trace_id: str = Field(default="")
+    reasoning: List[Dict[str, Any]] = Field(default_factory=list)
+    warnings: List[Dict[str, Any]] = Field(default_factory=list)
+    artifact_refs: Dict[str, ArtifactRef] = Field(default_factory=dict)
+
+
+def _field(source: Any, name: str, default: Any = None) -> Any:
+    """Read ``name`` off a mapping or an object.
+
+    LangGraph may hand back either dicts or model instances for nested state.
+    """
+    if source is None:
+        return default
+    if isinstance(source, dict):
+        return source.get(name, default)
+    return getattr(source, name, default)
+
+
+def _enum_value(value: Any) -> str:
+    """Render an enum (or plain value) as its string value."""
+    if value is None:
+        return ""
+    return str(getattr(value, "value", value))
+
+
+def _error_summary(error: Any) -> Dict[str, Any]:
+    """Project a PipelineError onto a client-safe summary (no stack traces)."""
+    return {
+        "node": _field(error, "node", "") or "",
+        "message": _field(error, "message", "") or "",
+        "error_code": _enum_value(_field(error, "error_code")),
+        "severity": _enum_value(_field(error, "severity")),
+    }
+
+
+def _sub_query_results(state: Dict[str, Any]) -> List[SubQueryResult]:
+    """Build the per-sub-query view from ``subgraph_outputs``."""
+    results: List[SubQueryResult] = []
+    for output in (state.get("subgraph_outputs") or {}).values():
+        sub_query = _field(output, "sub_query")
+        if not sub_query:
+            continue
+        results.append(
+            SubQueryResult(
+                id=_field(sub_query, "id", "") or "",
+                intent=_field(sub_query, "intent", "") or "",
+                datasource_id=_field(sub_query, "datasource_id", "") or "",
+                schema_version=_field(sub_query, "schema_version", "") or "",
+                sql=_field(output, "sql_draft", "") or "",
+            )
+        )
+    return results
+
+
+def result_from_state(state: Dict[str, Any]) -> QueryResult:
+    """Build a typed :class:`QueryResult` from a raw pipeline graph state."""
+    state = state or {}
+    return QueryResult(
+        sub_queries=_sub_query_results(state),
+        final_answer=_field(state.get("answer_synthesizer_response"), "final_answer"),
+        errors=[_error_summary(error) for error in (state.get("errors") or [])],
+        trace_id=state.get("trace_id") or "",
+        reasoning=list(state.get("reasoning") or []),
+        warnings=list(state.get("warnings") or []),
+        artifact_refs=state.get("artifact_refs") or {},
+    )
+
+
+class QueryAPI:
+    """
+    API for executing natural language queries against databases.
+    """
+
+    def __init__(self, ctx: NL2SQLContext):
+        self._ctx = ctx
+
+    def run_query(
+        self,
+        natural_language: str,
+        datasource_id: Optional[str] = None,
+        execute: bool = True,
+        user_context: Optional[UserContext] = None,
+    ) -> QueryResult:
+        """
+        Execute a natural language query against the database.
+
+        Args:
+            natural_language: The natural language query to execute
+            datasource_id: Optional specific datasource to query (otherwise auto-resolved)
+            execute: Whether to actually execute the SQL against the database
+            user_context: Optional user context for permissions
+
+        Returns:
+            A :class:`QueryResult` built from the pipeline graph state.
+        """
+        state = run_with_graph(
+            self._ctx,
+            natural_language,
+            datasource_id=datasource_id,
+            execute=execute,
+            user_context=user_context
+        )
+        return result_from_state(state)

nl2sql/api/result_api.py

@@ -0,0 +1,24 @@
+"""
+Result API for NL2SQL
+
+Provides functionality for result management and storage.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from nl2sql.context import NL2SQLContext
+from nl2sql.execution.contracts import ArtifactRef
+from nl2sql_adapter_sdk.contracts import ResultFrame
+
+
+class ResultAPI:
+    """
+    API for result management and storage.
+    """
+    
+    def __init__(self, ctx: NL2SQLContext):
+        self._ctx = ctx
+    
+    

nl2sql/api/settings_api.py

@@ -0,0 +1,65 @@
+"""
+Settings API for NL2SQL
+
+Provides functionality for configuration and settings management.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Any
+
+from nl2sql.context import NL2SQLContext
+from nl2sql.common.settings import Settings, settings
+
+
+class SettingsAPI:
+    """
+    API for configuration and settings management.
+    """
+    
+    def __init__(self, ctx: NL2SQLContext):
+        self._ctx = ctx
+        # ``settings`` is a module singleton that ``reload_settings()`` updates in
+        # place, so binding it once here keeps reading the live values.
+        self._settings: Settings = settings
+    
+    def get_current_settings(self) -> Dict[str, Any]:
+        """
+        Get the current application settings.
+        
+        Returns:
+            Dictionary of current settings
+        """
+        if self._settings:
+            return self._settings.model_dump()
+        else:
+            # Fallback to the global settings
+            from nl2sql.common.settings import settings
+            return settings.model_dump()
+    
+    def get_setting(self, key: str) -> Any:
+        """
+        Get a specific setting value.
+        
+        Args:
+            key: Setting key to retrieve
+            
+        Returns:
+            Value of the setting
+        """
+        settings_dict = self.get_current_settings()
+        return settings_dict.get(key)
+    
+    def validate_configuration(self) -> bool:
+        """
+        Validate the current configuration.
+        
+        Returns:
+            True if configuration is valid, False otherwise
+        """
+        try:
+            # Try to access all settings to validate them
+            self.get_current_settings()
+            return True
+        except Exception:
+            return False

nl2sql/auth/__init__.py

@@ -0,0 +1,8 @@
+from .models import UserContext, RolePolicy
+from .rbac import RBAC
+
+__all__ = [
+    "UserContext",
+    "RBAC",
+    "RolePolicy"
+]

nl2sql/auth/models.py

@@ -0,0 +1,36 @@
+from pydantic import BaseModel, Field, field_validator
+from typing import List, Optional
+from pydantic import ConfigDict
+
+class UserContext(BaseModel):
+    """User identity and permission context."""
+    user_id: Optional[str] = Field(default=None, description="Unique identifier for the user.")
+    tenant_id: Optional[str] = Field(default=None, description="Organization/Tenant identifier.")
+    roles: List[str] = Field(default_factory=list, description="List of assigned roles.")
+    model_config = ConfigDict(extra="ignore")
+
+
+class RolePolicy(BaseModel):
+    """Defines access control rules for a specific role."""
+    
+    description: str = Field(..., description="Human-readable description of the role")
+    role: str = Field(..., description="Role ID used for logging and auditing")
+    allowed_datasources: List[str] = Field(default_factory=list, description="List of allowed datasource IDs or '*'")
+    allowed_tables: List[str] = Field(default_factory=list, description="List of allowed tables in 'datasource.table' format")
+
+    @field_validator("allowed_tables")
+    def validate_namespace(cls, v: List[str]) -> List[str]:
+        """Enforces strict namespacing for allowed tables."""
+        for table in v:
+            if table == "*":
+                continue
+            
+            if table.endswith(".*"):
+                if table.count(".") < 1:
+                     raise ValueError(f"Invalid wildcard '{table}'. Must be 'datasource.*'.")
+                continue
+
+            if "." not in table:
+                raise ValueError(f"Invalid table '{table}'. Policy requires explicit 'datasource.table' format to prevent ambiguity.")
+        
+        return v

nl2sql/auth/rbac.py

@@ -0,0 +1,25 @@
+from .models import RolePolicy
+from .models import UserContext
+from typing import List, Dict
+
+class RBAC:
+    def __init__(self, policies: Dict[str, RolePolicy]):
+        self.policies = policies
+
+    def is_allowed(self, user_ctx: UserContext, datasource_id: str, table: str) -> bool:
+        policy =  [self.policies.get(role) for role in user_ctx.roles]
+        if not policy:
+            return False
+        return any(p.is_allowed(datasource_id, table) for p in policy)
+    
+    def get_allowed_tables(self, user_ctx: UserContext) -> List[str]:
+        policy = [self.policies.get(role) for role in user_ctx.roles]
+        if not policy:
+            return []
+        return list(set().union(*[p.allowed_tables for p in policy]))
+    
+    def get_allowed_datasources(self, user_ctx: UserContext) -> List[str]:
+        policy = [self.policies.get(role) for role in user_ctx.roles]
+        if not policy:
+            return []
+        return list(set().union(*[p.allowed_datasources for p in policy]))

nl2sql/cli/checks.py

@@ -0,0 +1,53 @@
+import importlib.util
+from typing import Dict, Tuple
+import pathlib
+from rich.markup import escape
+from rich.table import Table
+from rich.text import Text
+from nl2sql.cli.console import console, print_success, print_error
+
+def check_package(name: str) -> bool:
+    """Checks if a python package is installed."""
+    import_name = name.replace("-", "_")
+    return importlib.util.find_spec(import_name) is not None
+
+def verify_connectivity(print_table: bool = True) -> bool:
+    """
+    Checks connectivity for all profiles in the default config.
+    Returns True if all checks passed, False otherwise.
+    """
+    from nl2sql.common.settings import settings
+    from nl2sql.datasources import load_profiles
+    from nl2sql.diagnostics import check_connectivity as core_check
+    
+    try:
+        profiles = load_profiles(pathlib.Path(settings.datasource_config_path))
+        
+        with console.status("[bold green]Verifying connectivity...[/bold green]"):
+            results = core_check(list(profiles.values()))
+        
+        all_ok = True
+        
+        if print_table:
+            conn_table = Table(show_header=True, header_style="bold cyan")
+            conn_table.add_column("Datasource ID")
+            conn_table.add_column("Status")
+            conn_table.add_column("Details")
+        
+        for ds_id, (success, msg) in results.items():
+            if not success:
+                all_ok = False
+            
+            if print_table:
+                status = "[green]OK[/green]" if success else "[red]Failed[/red]"
+                details = msg if not success else ""
+                conn_table.add_row(Text(str(ds_id)), status, Text(str(details)))
+        
+        if print_table:
+            console.print(conn_table)
+            
+        return all_ok
+
+    except Exception as e:
+         console.print(f"[red]Connectivity check failed: {escape(str(e))}[/red]")
+         return False

nl2sql/cli/commands/benchmark.py

@@ -0,0 +1,34 @@
+import sys
+from nl2sql import BenchmarkAPI, BenchmarkConfig
+from nl2sql.cli.reporting import ConsolePresenter
+from nl2sql.cli.common.decorators import handle_cli_errors
+
+
+@handle_cli_errors
+def run_benchmark(
+    config: BenchmarkConfig,
+) -> None:
+    """Runs the benchmark suite based on provided arguments.
+
+    Args:
+        config (BenchmarkConfig): Benchmark run configuration.
+    """
+    presenter = ConsolePresenter()
+    api = BenchmarkAPI()
+
+    try:
+        matrix_result = api.run_matrix(config, progress_callback=presenter.track)
+    except Exception as e:
+        presenter.print_error(f"Benchmark Failed: {e}")
+        sys.exit(1)
+
+    for name, result in matrix_result.results_by_config.items():
+        presenter.print_header(f"Evaluating Config: {name}")
+        presenter.print_dataset_benchmark_results(
+            result.results,
+            iterations=result.iterations,
+            routing_only=config.routing_only,
+        )
+        presenter.print_metrics_summary(result.metrics, result.results, routing_only=config.routing_only)
+        if config.export_path:
+            presenter.export_results(result.results, config.export_path)

nl2sql/cli/commands/doctor.py

@@ -0,0 +1,49 @@
+import sys
+import importlib.util
+from rich.markup import escape
+from rich.table import Table
+from rich.panel import Panel
+from nl2sql.cli.console import console, print_success, print_error
+from nl2sql.cli.config import ADAPTER_DRIVERS, KNOWN_ADAPTERS
+from nl2sql.cli.checks import check_package, verify_connectivity
+
+from nl2sql.cli.common.decorators import handle_cli_errors
+
+@handle_cli_errors
+def doctor_command():
+    console.print(Panel("[bold cyan]NL2SQL Doctor[/bold cyan]"))
+
+    # 1. Python Version
+    py_ver = sys.version.split()[0]
+    console.print(f"Python Version: {py_ver}")
+    if sys.version_info < (3, 9):
+        print_error("Python 3.9+ required.")
+    else:
+        print_success("Python version OK.")
+
+    # 2. Core Check
+    if importlib.util.find_spec("nl2sql"):
+        print_success("Core package (nl2sql) installed.")
+    else:
+        print_error("Core package (nl2sql) NOT found.")
+
+    # 3. Adapters
+    console.print("\n[bold]Adapters:[/bold]")
+    table = Table(show_header=True, header_style="bold magenta")
+    table.add_column("Database")
+    table.add_column("Package")
+    table.add_column("Status")
+
+    for name, pkg in KNOWN_ADAPTERS.items():
+        # The adapter module always ships with nl2sql; the driver is what an
+        # extra adds, so that is what decides whether the dialect is usable.
+        ok = check_package(ADAPTER_DRIVERS[name])
+        status = "[green]Installed[/green]" if ok else "[red]Missing[/red]"
+        table.add_row(name, escape(pkg), status)
+    
+    console.print(table)
+    
+    # 4. Connectivity Check
+    console.print("\n[bold]Connectivity:[/bold]")
+    verify_connectivity(print_table=True)
+