natquery 1.0.0__py3-none-any.whl

natquery/config/settings.py

@@ -0,0 +1,100 @@
+import json
+from pathlib import Path
+from typing import Dict
+
+
+class Settings:
+    """
+    Configuration manager:
+    Reads configuration from:
+        ~/.natquery/<active_db>/config.json
+    """
+
+    BASE_DIR = Path.home() / ".natquery"
+
+    @classmethod
+    def _get_active_db_name(cls) -> str:
+        current_file = cls.BASE_DIR / "current_db"
+
+        if not current_file.exists():
+            raise RuntimeError("No active database configured.\nRun: natquery connect")
+
+        return current_file.read_text().strip()
+
+    @classmethod
+    def _get_config_path(cls) -> Path:
+        db_name = cls._get_active_db_name()
+        return cls.BASE_DIR / db_name / "config.json"
+
+    @classmethod
+    def exists(cls) -> bool:
+        return cls._get_config_path().exists()
+
+    @classmethod
+    def load_config(cls) -> Dict:
+        config_path = cls._get_config_path()
+
+        if not config_path.exists():
+            raise RuntimeError(
+                "NatQuery not configured properly.\nRun: natquery connect"
+            )
+
+        with open(config_path, "r") as f:
+            return json.load(f)
+
+    @classmethod
+    def get_db_config(cls) -> Dict:
+        config = cls.load_config()
+
+        if config.get("connection_type") == "dsn":
+            if "db_dsn" not in config or not config["db_dsn"]:
+                raise ValueError("Missing database DSN.")
+
+            return {
+                "type": "dsn",
+                "dsn": config["db_dsn"],
+                "dbname": config.get("db_name", "unknown_db"),
+            }
+
+        # Default: standard connection
+        required = [
+            "db_host",
+            "db_port",
+            "db_name",
+            "db_user",
+            "db_password",
+        ]
+
+        for key in required:
+            if key not in config or not config[key]:
+                raise ValueError(f"Missing database config field: {key}")
+
+        return {
+            "type": "standard",
+            "host": config["db_host"],
+            "port": config["db_port"],
+            "dbname": config["db_name"],
+            "user": config["db_user"],
+            "password": config["db_password"],
+            "sslmode": config.get("sslmode"),
+        }
+
+    @classmethod
+    def get_llm_config(cls) -> Dict[str, str]:
+        config = cls.load_config()
+
+        required = [
+            "llm_provider",
+            "llm_api_key",
+            "llm_model",
+        ]
+
+        for key in required:
+            if key not in config or not config[key]:
+                raise ValueError(f"Missing LLM config field: {key}")
+
+        return {
+            "provider": config["llm_provider"],
+            "api_key": config["llm_api_key"],
+            "model": config["llm_model"],
+        }

natquery/execution/engine.py

@@ -0,0 +1,29 @@
+from natquery.config.connection import get_connection, close_connection, get_cursor
+
+
+def execute_sql(sql: str):
+    conn = get_connection()
+
+    try:
+        cursor = get_cursor(conn)
+        cursor.execute(sql)
+
+        if cursor.description:
+            rows = cursor.fetchall()
+            # RealDictCursor rows are already dict-like, convert to plain dict
+            result = [
+                (
+                    dict(row)
+                    if hasattr(row, "keys")
+                    else dict(zip([desc[0] for desc in cursor.description], row))
+                )
+                for row in rows
+            ]
+        else:
+            result = []
+
+        conn.commit()
+        return result
+
+    finally:
+        close_connection(conn)

natquery/execution/explain.py

@@ -0,0 +1,51 @@
+from natquery.config.connection import get_connection, close_connection
+
+
+def run_explain(query: str):
+    """
+    Runs EXPLAIN (FORMAT JSON) on a query.
+    Returns the plan as JSON.
+    """
+    conn = get_connection()
+
+    try:
+        cursor = conn.cursor()
+
+        explain_query = f"EXPLAIN (FORMAT JSON) {query}"
+        cursor.execute(explain_query)
+
+        result = cursor.fetchone()[0]  # JSON result
+
+        # PostgreSQL returns list with single JSON object
+        if isinstance(result, list):
+            return result[0]
+
+        return result
+
+    finally:
+        close_connection(conn)
+
+
+def run_explain_analyze(query: str):
+    """
+    Runs EXPLAIN ANALYZE (FORMAT JSON) for actual execution metrics.
+    Returns the plan with actual execution stats as JSON.
+    """
+    conn = get_connection()
+
+    try:
+        cursor = conn.cursor()
+
+        explain_query = f"EXPLAIN (ANALYZE, FORMAT JSON) {query}"
+        cursor.execute(explain_query)
+
+        result = cursor.fetchone()[0]  # JSON result
+
+        # PostgreSQL returns list with single JSON object
+        if isinstance(result, list):
+            return result[0]
+
+        return result
+
+    finally:
+        close_connection(conn)

natquery/llm/client.py

@@ -0,0 +1,37 @@
+from groq import Groq
+from natquery.config.settings import Settings
+from natquery.schema.formatter import format_schema
+from natquery.prompt.builder import build_prompt
+from pathlib import Path
+import json
+
+
+def generate_sql(user_query: str) -> str:
+
+    llm_config = Settings.get_llm_config()
+    client = Groq(api_key=llm_config["api_key"])
+
+    db_config = Settings.get_db_config()
+    db_name = db_config["dbname"]
+
+    schema_path = Path.home() / ".natquery" / db_name / "schema.json"
+
+    if schema_path.exists():
+        with open(schema_path, "r") as f:
+            schema = json.load(f)
+    else:
+        schema = {}
+
+    schema_text = format_schema(schema)
+    messages = build_prompt(user_query, schema_text)
+
+    response = client.chat.completions.create(
+        model=llm_config["model"],
+        messages=messages,
+        temperature=0.0,
+    )
+
+    sql = response.choices[0].message.content.strip()
+    sql = sql.replace("```sql", "").replace("```", "").strip()
+
+    return sql

natquery/llm/self_corrector.py

@@ -0,0 +1,84 @@
+from groq import Groq
+from natquery.config.settings import Settings
+from natquery.schema.formatter import format_schema
+from natquery.orchestration.error_classifier import classify_error
+from pathlib import Path
+import json
+
+
+def correct_sql(user_query: str, failed_sql: str, error_message: str) -> str:
+
+    llm_config = Settings.get_llm_config()
+    client = Groq(api_key=llm_config["api_key"])
+
+    db_config = Settings.get_db_config()
+    db_name = db_config["dbname"]
+
+    schema_path = Path.home() / ".natquery" / db_name / "schema.json"
+
+    if schema_path.exists():
+        with open(schema_path, "r") as f:
+            schema = json.load(f)
+    else:
+        schema = {}
+
+    schema_text = format_schema(schema)
+
+    error_type = classify_error(error_message)
+
+    # Dynamic instruction based on error type
+    correction_hint = {
+        "SYNTAX_ERROR": "Fix SQL syntax.",
+        "UNDEFINED_TABLE": "Use only valid table names from schema.",
+        "UNDEFINED_COLUMN": "Fix incorrect column names using schema.",
+        "AMBIGUOUS_COLUMN": "Disambiguate columns using table aliases.",
+        "TYPE_ERROR": "Fix type mismatches and operators.",
+        "GROUP_BY_ERROR": "Ensure proper GROUP BY usage.",
+        "JOIN_ERROR": "Fix joins using foreign key relationships.",
+        "UNKNOWN": "Fix the SQL query to make it executable.",
+    }[error_type]
+
+    system_prompt = f"""
+You are a PostgreSQL expert specializing in fixing SQL queries.
+
+ERROR TYPE: {error_type}
+
+INSTRUCTION:
+{correction_hint}
+
+RULES:
+- Output ONLY corrected SQL.
+- Do NOT explain anything.
+- Use only schema provided.
+- Ensure query executes correctly.
+
+DATABASE SCHEMA:
+{schema_text}
+"""
+
+    user_prompt = f"""
+User Query:
+{user_query}
+
+Previous SQL:
+{failed_sql}
+
+Error:
+{error_message}
+
+Return corrected SQL:
+"""
+
+    response = client.chat.completions.create(
+        model=llm_config["model"],
+        messages=[
+            {"role": "system", "content": system_prompt.strip()},
+            {"role": "user", "content": user_prompt.strip()},
+        ],
+        temperature=0.0,
+    )
+
+    sql = response.choices[0].message.content.strip()
+    sql = sql.replace("```sql", "").replace("```", "").strip()
+
+    return sql

natquery/observability/cost_analyzer.py

@@ -0,0 +1,79 @@
+import re
+
+
+def extract_plan_root(plan_json):
+    if isinstance(plan_json, list):
+        plan_json = plan_json[0]
+    return plan_json
+
+
+def extract_summary(plan_json):
+    plan_json = extract_plan_root(plan_json)
+
+    root = plan_json.get("Plan", {})
+
+    return {
+        "execution_time_ms": plan_json.get("Execution Time"),
+        "planning_time_ms": plan_json.get("Planning Time"),
+        "total_cost": root.get("Total Cost"),
+        "plan_rows": root.get("Plan Rows"),
+    }
+
+
+def extract_columns_from_condition(condition: str):
+    """
+    Extract column names from filter/index conditions.
+    Example:
+        "((email = 'x') AND (status = 'active'))"
+    """
+
+    if not condition:
+        return []
+
+    # Simple heuristic (robust enough for now)
+    matches = re.findall(r"([a-zA-Z_][a-zA-Z0-9_]*)\s*(=|>|<|>=|<=)", condition)
+
+    return list(set([m[0] for m in matches]))
+
+
+def flatten_plan(node, nodes=None):
+    if nodes is None:
+        nodes = []
+
+    filter_cond = node.get("Filter")
+    index_cond = node.get("Index Cond")
+
+    extracted_columns = []
+    extracted_columns += extract_columns_from_condition(filter_cond)
+    extracted_columns += extract_columns_from_condition(index_cond)
+
+    node_info = {
+        "node_type": node.get("Node Type"),
+        "relation": node.get("Relation Name"),
+        "startup_cost": node.get("Startup Cost"),
+        "total_cost": node.get("Total Cost"),
+        "plan_rows": node.get("Plan Rows"),
+        "actual_rows": node.get("Actual Rows"),
+        "actual_time": node.get("Actual Total Time"),
+        "filter": filter_cond,
+        "index_cond": index_cond,
+        "columns": extracted_columns,
+        "join_type": node.get("Join Type"),
+    }
+
+    nodes.append(node_info)
+
+    for child in node.get("Plans", []) or []:
+        flatten_plan(child, nodes)
+
+    return nodes
+
+
+def analyze_cost(plan_json):
+    plan_json = extract_plan_root(plan_json)
+    root = plan_json.get("Plan", {})
+
+    return {
+        "summary": extract_summary(plan_json),
+        "nodes": flatten_plan(root),
+    }

natquery/observability/index_recommender.py

@@ -0,0 +1,58 @@
+def suggest_indexes(plan_analysis):
+    nodes = plan_analysis.get("nodes", [])
+    suggestions = []
+
+    for node in nodes:
+        node_type = node.get("node_type")
+        table = node.get("relation")
+        columns = node.get("columns", [])
+
+        # SEQ SCAN (main case)
+        if node_type == "Seq Scan" and table:
+
+            if columns:
+                if len(columns) == 1:
+                    col = columns[0]
+
+                    suggestions.append(
+                        {
+                            "type": "index",
+                            "table": table,
+                            "columns": columns,
+                            "sql": f"CREATE INDEX idx_{table}_{col} ON {table}({col});",
+                            "reason": "Sequential scan with filter condition",
+                        }
+                    )
+
+                else:
+                    cols = ", ".join(columns)
+
+                    suggestions.append(
+                        {
+                            "type": "composite_index",
+                            "table": table,
+                            "columns": columns,
+                            "sql": f"CREATE INDEX idx_{table}_{'_'.join(columns)} ON {table}({cols});",
+                            "reason": "Multiple filter conditions detected",
+                        }
+                    )
+
+            else:
+                suggestions.append(
+                    {
+                        "type": "seq_scan_warning",
+                        "table": table,
+                        "reason": "Full table scan detected (no filter)",
+                    }
+                )
+
+        # Join optimization
+        if node.get("join_type"):
+            suggestions.append(
+                {
+                    "type": "join",
+                    "reason": "Ensure indexes exist on join keys",
+                }
+            )
+
+    return suggestions

natquery/observability/logger.py

@@ -0,0 +1,109 @@
+import json
+import uuid
+import hashlib
+from datetime import datetime, timezone
+from pathlib import Path
+
+
+class NatQueryLogger:
+
+    BASE_DIR = Path.home() / ".natquery"
+    SYSTEM_LOG = BASE_DIR / "system.log"
+
+    @staticmethod
+    def _timestamp():
+        return datetime.now(timezone.utc).isoformat() + "Z"
+
+    @classmethod
+    def generate_conv_id(cls) -> str:
+        return str(uuid.uuid4())
+
+    @classmethod
+    def log_event(
+        cls,
+        level: str,
+        event: str,
+        db_name: str = None,
+        conv_id: str = None,
+        details: dict = None,
+    ):
+
+        cls.BASE_DIR.mkdir(exist_ok=True)
+
+        log_entry = {
+            "timestamp": cls._timestamp(),
+            "level": level,
+            "event": event,
+            "conv_id": conv_id,
+            "db_name": db_name,
+            "details": details or {},
+        }
+
+        with open(cls.SYSTEM_LOG, "a") as f:
+            f.write(json.dumps(log_entry) + "\n")
+
+    @classmethod
+    def log_conversation(
+        cls,
+        db_name: str,
+        conv_id: str,
+        user_query: str,
+        generated_sql: str,
+        rows_returned: int,
+        execution_time_ms: float,
+    ):
+
+        db_log_dir = cls.BASE_DIR / db_name / "logs"
+        db_log_dir.mkdir(parents=True, exist_ok=True)
+
+        conv_file = db_log_dir / "conversations.jsonl"
+
+        entry = {
+            "timestamp": cls._timestamp(),
+            "conv_id": conv_id,
+            "user_query": user_query,
+            "generated_sql": generated_sql,
+            "rows_returned": rows_returned,
+            "execution_time_ms": execution_time_ms,
+        }
+
+        with open(conv_file, "a") as f:
+            f.write(json.dumps(entry) + "\n")
+
+    @staticmethod
+    def _hash_query(sql: str) -> str:
+        return hashlib.sha256(sql.encode()).hexdigest()[:16]
+
+    @classmethod
+    def log_performance(
+        cls,
+        db_name: str,
+        conv_id: str,
+        sql: str,
+        summary: dict,
+        suggestions: list,
+    ):
+        """
+        Stores structured performance history.
+        """
+
+        perf_dir = cls.BASE_DIR / db_name / "logs"
+        perf_dir.mkdir(parents=True, exist_ok=True)
+
+        perf_file = perf_dir / "performance.jsonl"
+
+        query_hash = cls._hash_query(sql)
+
+        entry = {
+            "timestamp": cls._timestamp(),
+            "conv_id": conv_id,
+            "query_hash": query_hash,
+            "sql": sql,
+            "execution_time_ms": summary.get("execution_time_ms"),
+            "total_cost": summary.get("total_cost"),
+            "planning_time_ms": summary.get("planning_time_ms"),
+            "suggestions": suggestions,
+        }
+
+        with open(perf_file, "a") as f:
+            f.write(json.dumps(entry) + "\n")

natquery/observability/performance_api.py

@@ -0,0 +1,58 @@
+import json
+from pathlib import Path
+
+
+BASE_DIR = Path.home() / ".natquery"
+
+
+def load_performance(db_name: str):
+    perf_file = BASE_DIR / db_name / "logs" / "performance.jsonl"
+
+    if not perf_file.exists():
+        return []
+
+    entries = []
+    with open(perf_file, "r") as f:
+        for line in f:
+            entries.append(json.loads(line))
+
+    return entries
+
+
+def get_last_run(db_name: str):
+    data = load_performance(db_name)
+    return data[-1] if data else None
+
+
+def get_query_history(db_name: str, query_hash: str):
+    data = load_performance(db_name)
+    return [d for d in data if d["query_hash"] == query_hash]
+
+
+def compare_last_two_runs(db_name: str):
+    data = load_performance(db_name)
+
+    if len(data) < 2:
+        return None
+
+    last = data[-1]
+    prev = data[-2]
+
+    return {
+        "execution_time_diff": last["execution_time_ms"] - prev["execution_time_ms"],
+        "cost_diff": last["total_cost"] - prev["total_cost"],
+        "last": last,
+        "previous": prev,
+    }
+
+
+def get_slowest_queries(db_name: str, top_k=5):
+    data = load_performance(db_name)
+
+    sorted_data = sorted(
+        data,
+        key=lambda x: x.get("execution_time_ms", 0),
+        reverse=True,
+    )
+
+    return sorted_data[:top_k]

natquery/orchestration/error_classifier.py

@@ -0,0 +1,29 @@
+def classify_error(error_message: str) -> str:
+    """
+    Classify PostgreSQL errors into categories.
+    """
+
+    msg = error_message.lower()
+
+    if "syntax error" in msg:
+        return "SYNTAX_ERROR"
+
+    if "does not exist" in msg:
+        if "column" in msg:
+            return "UNDEFINED_COLUMN"
+        if "relation" in msg or "table" in msg:
+            return "UNDEFINED_TABLE"
+
+    if "ambiguous" in msg:
+        return "AMBIGUOUS_COLUMN"
+
+    if "operator does not exist" in msg or "type" in msg:
+        return "TYPE_ERROR"
+
+    if "must appear in the group by clause" in msg:
+        return "GROUP_BY_ERROR"
+
+    if "missing from clause" in msg:
+        return "JOIN_ERROR"
+
+    return "UNKNOWN"