nl2sql-engine 0.1.0__py3-none-any.whl

nl2sql/__init__.py

@@ -0,0 +1,38 @@
+# nl2sql package
+
+from .public_api import NL2SQL, QueryResult
+
+# Also expose individual API modules for more granular access
+from .api.query_api import QueryAPI
+from .api.datasource_api import DatasourceAPI
+from .api.llm_api import LLM_API
+from .api.indexing_api import IndexingAPI
+from .api.auth_api import AuthAPI
+from .api.settings_api import SettingsAPI
+from .api.result_api import ResultAPI
+from .api.policy_api import PolicyAPI
+from .api.benchmark_api import BenchmarkAPI
+
+# Also expose core models and enums
+from .common.errors import ErrorSeverity, ErrorCode, PipelineError
+from .auth.models import UserContext
+from .evaluation.types import BenchmarkConfig
+
+__all__ = [
+    "NL2SQL",
+    "QueryResult",
+    "QueryAPI",
+    "DatasourceAPI",
+    "LLM_API",
+    "IndexingAPI",
+    "AuthAPI",
+    "SettingsAPI",
+    "ResultAPI",
+    "PolicyAPI",
+    "BenchmarkAPI",
+    "ErrorSeverity",
+    "ErrorCode",
+    "PipelineError",
+    "UserContext",
+    "BenchmarkConfig",
+]

nl2sql/adapters/duckdb/adapter.py

@@ -0,0 +1,71 @@
+from typing import Any, Dict
+
+from nl2sql.adapters.sqlalchemy_base import (
+    CostEstimate,
+    DryRunResult,
+    QueryPlan,
+    BaseSQLAlchemyAdapter
+)
+
+from pydantic import BaseModel, Field
+
+
+class DuckdbConnectionConfig(BaseModel):
+    """Strict configuration schema for DuckDB adapter."""
+    type: str
+    database: str = Field(..., description="Path to DuckDB database file, or ':memory:'")
+    options: Dict[str, Any] = Field(default_factory=dict)
+
+    model_config = {"extra": "ignore"}
+
+
+class DuckdbAdapter(BaseSQLAlchemyAdapter):
+
+    def construct_uri(self, args: Dict[str, Any]) -> str:
+        """Constructs the DuckDB connection URI.
+
+        Args:
+            args: The raw connection arguments dictionary.
+
+        Returns:
+            str: The fully constructed SQLAlchemy connection URI.
+
+        Raises:
+            ValidationError: If the configuration is invalid.
+        """
+        config = DuckdbConnectionConfig(**args)
+        return f"duckdb:///{config.database}"
+
+    def dry_run(self, query: str) -> DryRunResult:
+        try:
+            self.execute_sql(f"EXPLAIN {query}")
+            return DryRunResult(is_valid=True, error_message=None)
+        except Exception as e:
+            return DryRunResult(is_valid=False, error_message=str(e))
+
+    def explain(self, query: str) -> QueryPlan:
+        try:
+            res = self.execute_sql(f"EXPLAIN {query}")
+            # EXPLAIN yields (explain_key, explain_value) rows; the plan is the value.
+            plan_text = "\n".join(str(row[-1]) for row in res.rows)
+            return QueryPlan(plan_text=plan_text)
+        except Exception:
+            return QueryPlan(plan_text="Could not retrieve plan")
+
+    def cost_estimate(self, query: str) -> CostEstimate:
+        # DuckDB's EXPLAIN prints an operator tree with no cost or cardinality
+        # figures, so there is nothing real to report here.
+        try:
+            self.execute_sql(f"EXPLAIN {query}")
+            return CostEstimate(estimated_cost=1.0, estimated_rows=10) # Stub
+        except Exception:
+            return CostEstimate(estimated_cost=-1.0, estimated_rows=0)
+
+    def get_dialect(self) -> str:
+        return "duckdb"
+
+    @property
+    def exclude_schemas(self) -> set[str]:
+        # DuckDB qualifies schema names with their catalog; only the attached
+        # database's own schemas carry user tables.
+        return {"system.main", "system.information_schema", "temp.main"}

nl2sql/adapters/mssql/adapter.py

@@ -0,0 +1,122 @@
+from typing import Any, List, Dict
+from sqlalchemy import create_engine, text, inspect
+from sqlalchemy.dialects import mssql
+from nl2sql.adapters.sqlalchemy_base import (
+    CostEstimate,
+    DryRunResult,
+    QueryPlan,
+)
+from nl2sql.adapters.sqlalchemy_base import BaseSQLAlchemyAdapter
+
+from pydantic import BaseModel, Field, SecretStr
+from typing import Optional
+
+class MssqlConnectionConfig(BaseModel):
+    """Strict configuration schema for MSSQL adapter."""
+    type: str = Field("mssql", description="Connection type")
+    host: str = Field(..., description="Server hostname")
+    user: Optional[str] = None
+    password: Optional[SecretStr] = None
+    port: int = 1433
+    database: str = Field(..., description="Database name")
+    driver: str = "ODBC Driver 17 for SQL Server"
+    trusted_connection: bool = False
+    options: Dict[str, Any] = Field(default_factory=dict)
+    
+    model_config = {"extra": "ignore"}
+
+class MssqlAdapter(BaseSQLAlchemyAdapter):
+
+    def construct_uri(self, args: Dict[str, Any]) -> str:
+        """Constructs the MSSQL connection URI.
+
+        Args:
+            args: The raw connection arguments dictionary.
+
+        Returns:
+            str: The fully constructed SQLAlchemy connection URI.
+        
+        Raises:
+            ValidationError: If the configuration is invalid.
+        """
+        config = MssqlConnectionConfig(**args)
+        
+        user = config.user or ""
+        password = config.password.get_secret_value() if config.password else ""
+        host = config.host
+        port = config.port
+        database = config.database
+        driver = config.driver
+        
+        options = config.options.copy()
+        options["driver"] = driver
+        
+        if config.trusted_connection:
+            options["Trusted_Connection"] = "yes"
+            
+        creds = f"{user}:{password}@" if user or password else ""
+        netloc = f"{host}:{port}"
+        
+        from urllib.parse import urlencode
+        query_str = "?" + urlencode(options)
+        
+        return f"mssql+pyodbc://{creds}{netloc}/{database}{query_str}"
+
+    def dry_run(self, sql: str) -> DryRunResult:
+        try:
+            with self.engine.connect() as conn:
+                conn.execute(text("SET NOEXEC ON"))
+                try:
+                    conn.execute(text(sql))
+                    valid = True
+                    msg = None
+                except Exception as e:
+                    valid = False
+                    msg = str(e)
+                finally:
+                    conn.execute(text("SET NOEXEC OFF"))
+            return DryRunResult(is_valid=valid, error_message=msg)
+        except Exception as e:
+            return DryRunResult(is_valid=False, error_message=str(e))
+
+    def explain(self, sql: str) -> QueryPlan:
+        try:
+            with self.engine.connect() as conn:
+                conn.execute(text("SET SHOWPLAN_XML ON"))
+        except Exception as e:
+            return QueryPlan(plan_text=f"Error: {e}")
+
+    def get_dialect(self) -> str:
+        """MSSQL uses T-SQL dialect."""
+        return mssql.dialect.name
+
+    def cost_estimate(self, sql: str) -> CostEstimate:
+        import re
+        try:
+            with self.engine.connect() as conn:
+                conn.execute(text("SET SHOWPLAN_XML ON"))
+                try:
+                    res = conn.execute(text(sql)).fetchone()
+                finally:
+                    conn.execute(text("SET SHOWPLAN_XML OFF"))
+                
+                if res and res[0]:
+                    xml_str = res[0]
+                    # Extract cost and rows using regex to avoid namespace complexity
+                    # StatementSubTreeCost="0.00328" StatementEstRows="1"
+                    cost_match = re.search(r'StatementSubTreeCost="([^"]+)"', xml_str)
+                    rows_match = re.search(r'StatementEstRows="([^"]+)"', xml_str)
+                    
+                    return CostEstimate(
+                        estimated_cost=float(cost_match.group(1)) if cost_match else 0.0,
+                        estimated_rows=float(rows_match.group(1)) if rows_match else 0
+                    )
+        except Exception:
+             pass
+        return CostEstimate(estimated_cost=0.0, estimated_rows=0)
+
+    
+    @property
+    def exclude_schemas(self) -> set[str]:
+        return {"sys", "INFORMATION_SCHEMA"}
+

nl2sql/adapters/mysql/adapter.py

@@ -0,0 +1,123 @@
+from typing import Any, List, Dict
+from sqlalchemy import create_engine, text, inspect
+from sqlalchemy.dialects import mysql
+from nl2sql.adapters.sqlalchemy_base import (
+    CostEstimate,
+    DryRunResult,
+    QueryPlan,
+    BaseSQLAlchemyAdapter
+)
+
+from pydantic import BaseModel, Field, SecretStr
+from typing import Optional
+
+class MysqlConnectionConfig(BaseModel):
+    """Strict configuration schema for MySQL adapter."""
+    type: str
+    host: str = Field(..., description="MySQL server hostname")
+    user: str = Field(..., description="Username")
+    password: SecretStr = Field(..., description="Password")
+    port: int = 3306
+    database: str = Field(..., description="Database name")
+    options: Dict[str, Any] = Field(default_factory=dict)
+    
+    model_config = {"extra": "ignore"}
+
+class MysqlAdapter(BaseSQLAlchemyAdapter):
+
+    def construct_uri(self, args: Dict[str, Any]) -> str:
+        """Constructs the MySQL connection URI.
+
+        Args:
+            args: The raw connection arguments dictionary.
+
+        Returns:
+            str: The fully constructed SQLAlchemy connection URI.
+        
+        Raises:
+            ValidationError: If the configuration is invalid.
+        """
+        config = MysqlConnectionConfig(**args)
+        
+        user = config.user
+        password = config.password.get_secret_value()
+        host = config.host
+        port = config.port
+        database = config.database
+        options = config.options.copy()
+        
+        creds = f"{user}:{password}@" if user or password else ""
+        netloc = f"{host}:{port}"
+        
+        query_str = ""
+        if options:
+            from urllib.parse import urlencode
+            query_str = "?" + urlencode(options)
+            
+        return f"mysql+pymysql://{creds}{netloc}/{database}{query_str}"
+
+    def connect(self) -> None:
+        """MySQL-specific connection with Native Server-Side Timeout."""
+        if not self.connection_string:
+             raise ValueError(f"Connection string is required for {self}")
+             
+        connect_args = {}
+        if self.statement_timeout_ms:
+            # Native MySQL Timeout (server-side)
+            # SET MAX_EXECUTION_TIME={ms}
+            connect_args["init_command"] = f"SET MAX_EXECUTION_TIME={self.statement_timeout_ms}"
+
+        try:
+            self.engine = create_engine(
+                self.connection_string, 
+                pool_pre_ping=True,
+                execution_options=self.execution_options,
+                connect_args=connect_args
+            )
+        except Exception as e:
+            import logging
+            logging.getLogger(__name__).error(f"Failed to connect to MySQL: {e}")
+            raise
+
+    def dry_run(self, sql: str) -> DryRunResult:
+        try:
+            with self.engine.connect() as conn:
+                trans = conn.begin()
+                conn.execute(text(sql))
+                trans.rollback()
+            return DryRunResult(is_valid=True)
+        except Exception as e:
+            return DryRunResult(is_valid=False, error_message=str(e))
+
+    def explain(self, sql: str) -> QueryPlan:
+        try:
+            with self.engine.connect() as conn:
+                res = conn.execute(text(f"EXPLAIN FORMAT=JSON {sql}")).scalar()
+                return QueryPlan(plan_text=str(res))
+        except Exception as e:
+            return QueryPlan(plan_text=f"Error: {e}")
+
+
+    def cost_estimate(self, sql: str) -> CostEstimate:
+        import json
+        try:
+            with self.engine.connect() as conn:
+                res = conn.execute(text(f"EXPLAIN FORMAT=JSON {sql}")).scalar()
+                if res:
+                    data = json.loads(res)
+                    cost_info = data.get('query_block', {}).get('cost_info', {})
+                    return CostEstimate(
+                        estimated_cost=float(cost_info.get('query_cost', 0.0)),
+                        estimated_rows=0 # MySQL doesn't give a single total rows estimate easily
+                    )
+        except Exception:
+             pass
+        return CostEstimate(estimated_cost=0.0, estimated_rows=0)
+
+    def get_dialect(self) -> str:
+        return mysql.dialect.name
+
+
+    @property
+    def exclude_schemas(self) -> set[str]:
+        return {"mysql", "INFORMATION_SCHEMA", "performance_schema", "sys"}

nl2sql/adapters/postgres/adapter.py

@@ -0,0 +1,115 @@
+from typing import Dict, Any
+from sqlalchemy import create_engine, inspect, text
+from sqlalchemy.dialects import postgresql
+from nl2sql.adapters.sqlalchemy_base import (
+    DryRunResult,
+    QueryPlan,
+    CostEstimate,
+    BaseSQLAlchemyAdapter
+)
+
+from pydantic import BaseModel, Field, SecretStr
+from typing import Optional
+
+class PostgresConnectionConfig(BaseModel):
+    """Strict configuration schema for Postgres adapter."""
+    type: str = Field("postgresql", description="Connection type")
+    host: str = Field(..., description="Postgres server hostname")
+    user: str = Field(..., description="Username")
+    password: SecretStr = Field(..., description="Password")
+    port: int = 5432
+    database: str = Field(..., description="Database name")
+    options: Dict[str, Any] = Field(default_factory=dict)
+    
+    model_config = {"extra": "ignore"}
+
+class PostgresAdapter(BaseSQLAlchemyAdapter):
+
+    def construct_uri(self, args: Dict[str, Any]) -> str:
+        """Constructs the Postgres connection URI.
+
+        Args:
+            args: The raw connection arguments dictionary.
+
+        Returns:
+            str: The fully constructed SQLAlchemy connection URI.
+        
+        Raises:
+            ValidationError: If the configuration is invalid.
+        """
+        config = PostgresConnectionConfig(**args)
+        
+        user = config.user
+        password = config.password.get_secret_value()
+        host = config.host
+        port = config.port
+        database = config.database
+        options = config.options.copy()
+        
+        creds = f"{user}:{password}@" if user or password else ""
+        netloc = f"{host}:{port}"
+        
+        query_str = ""
+        if options:
+            from urllib.parse import urlencode
+            query_str = "?" + urlencode(options)
+            
+        return f"postgresql://{creds}{netloc}/{database}{query_str}"
+
+    def connect(self) -> None:
+        """Postgres-specific connection with Native Server-Side Timeout."""
+        if not self.connection_string:
+             raise ValueError(f"Connection string is required for {self}")
+             
+        connect_args = {}
+        if self.statement_timeout_ms:
+            connect_args["options"] = f"-c statement_timeout={self.statement_timeout_ms}"
+
+        try:
+            self.engine = create_engine(
+                self.connection_string, 
+                pool_pre_ping=True,
+                execution_options=self.execution_options, # Pass standard options too
+                connect_args=connect_args
+            )
+        except Exception as e:
+            import logging
+            logging.getLogger(__name__).error(f"Failed to connect to Postgres: {e}")
+            raise
+
+    def dry_run(self, sql: str) -> DryRunResult:
+        try:
+            self.execute_sql(f"EXPLAIN {sql}")
+            return DryRunResult(is_valid=True)
+        except Exception as e:
+            return DryRunResult(is_valid=False, error_message=str(e))
+
+    def explain(self, sql: str) -> QueryPlan:
+        try:
+            res = self.execute_sql(f"EXPLAIN (FORMAT JSON) {sql}")
+            return QueryPlan(plan_text=str(res.rows))
+        except Exception:            # Fallback
+            return QueryPlan(plan_text="Could not retrieve plan")
+
+    def cost_estimate(self, sql: str) -> CostEstimate:
+        try:
+            res = self.execute_sql(f"EXPLAIN (FORMAT JSON) {sql}")
+            if res.rows and res.rows[0]:
+                plan_data = res.rows[0][0] # The JSON object/list
+                if isinstance(plan_data, list) and len(plan_data) > 0:
+                    root = plan_data[0].get('Plan', {})
+                    return CostEstimate(
+                        estimated_cost=float(root.get('Total Cost', 0.0)),
+                        estimated_rows=int(root.get('Plan Rows', 0))
+                    )
+            return CostEstimate(estimated_cost=0.0, estimated_rows=0)
+        except Exception:
+            return CostEstimate(estimated_cost=0.0, estimated_rows=0)
+
+    def get_dialect(self) -> str:
+        return postgresql.dialect.name
+
+
+    @property
+    def exclude_schemas(self) -> set[str]:
+        return {"pg_catalog", "information_schema"}

nl2sql/adapters/sqlalchemy_base/__init__.py

@@ -0,0 +1,17 @@
+from .adapter import BaseSQLAlchemyAdapter
+from .models import (
+    QueryResult,
+    CostEstimate,
+    DryRunResult,
+    QueryPlan,
+    AdapterError,
+)
+
+__all__ = [
+    "BaseSQLAlchemyAdapter",
+    "QueryResult",
+    "CostEstimate",
+    "DryRunResult",
+    "QueryPlan",
+    "AdapterError",
+]