natquery 1.0.0__py3-none-any.whl

natquery/cli/commands.py

@@ -0,0 +1,256 @@
+import json
+from pathlib import Path
+from getpass import getpass
+from urllib.parse import urlparse
+from natquery.config.settings import Settings
+from natquery.observability.performance_api import (
+    get_last_run,
+    compare_last_two_runs,
+    get_slowest_queries,
+)
+
+
+BASE_DIR = Path.home() / ".natquery"
+CONFIG_FILE = BASE_DIR / "config.json"
+
+
+def extract_db_name_from_dsn(dsn: str) -> str:
+    """
+    Extract database name from a PostgreSQL DSN.
+    Supports both URL format (postgresql://) and key-value format (dbname=).
+    """
+    parsed = urlparse(dsn)
+
+    if parsed.scheme:  # URL format
+        db_name = parsed.path.lstrip("/").split("?")[
+            0
+        ]  # Remove leading / and query params
+        return db_name if db_name else "unknown_db"
+
+    # Fallback: key-value DSN (e.g., "dbname=mydb user=postgres host=localhost")
+    parts = dsn.split()
+    for part in parts:
+        if part.startswith("dbname="):
+            return part.split("=", 1)[1]
+
+    return "unknown_db"
+
+
+def connect_command(dsn_arg=None):
+    """
+    Setup for NatQuery.
+    If dsn_arg is provided, extracts DB name and uses it as workspace.
+    Otherwise, runs interactive setup for standard or DSN connections.
+    """
+
+    print("============ NatQuery Setup ============")
+
+    BASE_DIR.mkdir(exist_ok=True)
+
+    # If DSN was provided as argument
+    if dsn_arg:
+        db_name = extract_db_name_from_dsn(dsn_arg)
+        workspace_name = db_name
+
+        print(f"Your Datbase Name: {db_name}")
+        print(f"Workspace Name: {workspace_name}")
+
+        config = {
+            "connection_type": "dsn",
+            "db_dsn": dsn_arg,
+            "db_name": db_name,
+            "workspace_name": workspace_name,
+        }
+
+    else:
+        db_name = input("Database Name: ").strip()
+        workspace_name = db_name
+
+        db_dir = BASE_DIR / workspace_name
+        db_dir.mkdir(parents=True, exist_ok=True)
+
+        ssl_input = input("Use SSL? (yes/no): ").strip().lower()
+        use_ssl = ssl_input in {"yes", "y"}
+
+        config = {
+            "connection_type": "standard",
+            "db_host": input("DB Host (e.g. 127.0.0.1): ").strip(),
+            "db_port": input("DB Port (e.g. 5432): ").strip(),
+            "db_name": db_name,
+            "db_user": input("DB User: ").strip(),
+            "db_password": getpass("DB Password: "),
+            "workspace_name": workspace_name,
+        }
+
+        if use_ssl:
+            config["sslmode"] = "require"
+
+    db_dir = BASE_DIR / workspace_name
+    db_dir.mkdir(parents=True, exist_ok=True)
+
+    # LLM config (common to both)
+    config.update(
+        {
+            "llm_provider": input("LLM Provider (groq): ").strip(),
+            "llm_api_key": getpass("LLM API Key: "),
+            "llm_model": input("LLM Model (e.g. llama-3.1-70b-versatile): ").strip(),
+        }
+    )
+
+    config_file = db_dir / "config.json"
+
+    with open(config_file, "w") as f:
+        json.dump(config, f, indent=2)
+
+    # Set active database
+    (BASE_DIR / "current_db").write_text(workspace_name)
+
+    print(f"\nConfiguration saved for workspace: {workspace_name}")
+    print("You can now run: natquery\n")
+
+
+def reset_command():
+    """
+    Deletes the active database configuration.
+    """
+
+    try:
+        config_path = Settings._get_config_path()
+    except RuntimeError:
+        print("No active database configured.")
+        return
+
+    if config_path.exists():
+        config_path.unlink()
+        print("Configuration removed for active database.")
+    else:
+        print("No configuration found.")
+
+
+def show_config_command():
+    """
+    Display current active DB configuration (without secrets).
+    """
+
+    try:
+        config = Settings.load_config()
+    except RuntimeError:
+        print("NatQuery not configured.")
+        print("Run: natquery connect")
+        return
+
+    safe_config = config.copy()
+
+    if "db_password" in safe_config:
+        safe_config["db_password"] = "********"
+
+    if "llm_api_key" in safe_config:
+        safe_config["llm_api_key"] = "********"
+
+    print(json.dumps(safe_config, indent=2))
+
+
+def select_database_command():
+    """
+    List all configured databases and allow user to select one as the active database.
+    Updates the current_db file without requiring reconnection.
+    """
+
+    BASE_DIR.mkdir(exist_ok=True)
+
+    # Find all database directories (those containing config.json)
+    databases = []
+
+    if BASE_DIR.exists():
+        for item in BASE_DIR.iterdir():
+            if item.is_dir() and (item / "config.json").exists():
+                databases.append(item.name)
+
+    if not databases:
+        print("No configured databases found.")
+        print("Run: natquery connect")
+        return
+
+    # Get current database
+    current_db_file = BASE_DIR / "current_db"
+    current_db = (
+        current_db_file.read_text().strip() if current_db_file.exists() else None
+    )
+
+    # Display databases
+    print("\n========== Configured Databases ==========")
+    for i, db in enumerate(databases, 1):
+        marker = " (current)" if db == current_db else ""
+        print(f"{i}. {db}{marker}")
+
+    print("\n")
+
+    # Get user selection
+    while True:
+        try:
+            selection = input("Select a database (enter number): ").strip()
+            index = int(selection) - 1
+
+            if 0 <= index < len(databases):
+                selected_db = databases[index]
+                break
+            else:
+                print(
+                    f"Invalid selection. Please enter a number between 1 and {len(databases)}."
+                )
+        except ValueError:
+            print("Invalid input. Please enter a number.")
+
+    # Update current database
+    current_db_file.write_text(selected_db)
+    print(f"\nActive database switched to: {selected_db}")
+
+
+def performance_summary_command():
+    db_name = Settings.get_db_config()["dbname"]
+
+    last = get_last_run(db_name)
+
+    if not last:
+        print("No performance data available.")
+        return
+
+    print("\n====== Performance Summary ======")
+    print(f"Execution Time: {last['execution_time_ms']} ms")
+    print(f"Total Cost: {last['total_cost']}")
+    print(f"Planning Time: {last['planning_time_ms']} ms")
+
+    print("\n--- Suggestions ---")
+    for s in last["suggestions"]:
+        print(f"- {s}")
+
+
+def performance_compare_command():
+    db_name = Settings.get_db_config()["dbname"]
+
+    comp = compare_last_two_runs(db_name)
+
+    if not comp:
+        print("Oops! Not enough data to compare.")
+        return
+
+    print("\n====== Performance Comparison ======")
+
+    print(f"Execution Time Change: {comp['execution_time_diff']} ms")
+    print(f"Cost Change: {comp['cost_diff']}")
+
+    if comp["execution_time_diff"] < 0:
+        print("Booyah! Query improved!")
+    else:
+        print("Oops! Query got slower.")
+
+
+def performance_slowest_command():
+    db_name = Settings.get_db_config()["dbname"]
+
+    slow = get_slowest_queries(db_name)
+
+    print("\n====== Slowest Queries ======")
+
+    for i, q in enumerate(slow, 1):
+        print(f"{i}. {q['execution_time_ms']} ms → {q['sql']}")

natquery/cli/formatter.py

@@ -0,0 +1,126 @@
+from rich.table import Table
+from rich.console import Console
+from rich.panel import Panel
+from rich import box
+import math
+
+console = Console()
+
+
+def _is_number(value):
+    return isinstance(value, (int, float))
+
+
+def _infer_alignment(data, column):
+    """
+    Infer alignment based on column values.
+    """
+    for row in data:
+        val = row.get(column)
+        if val is not None:
+            if _is_number(val):
+                return "right"
+            else:
+                return "left"
+    return "left"
+
+
+def _truncate(text, max_length=50):
+    """
+    Soft truncate long values (for extremely wide columns).
+    """
+    text = str(text)
+    if len(text) > max_length:
+        return text[: max_length - 3] + "..."
+    return text
+
+
+def _format_value(value):
+    """
+    Clean formatting of cell values.
+    """
+    if value is None:
+        return ""
+    return str(value)
+
+
+def print_table(
+    data: list[dict],
+    title: str = "#### Result ####",
+    page_size: int = 20,
+    wrap: bool = True,
+):
+    """
+    Advanced table printer:
+    - smart alignment
+    - pagination
+    - wrapping
+    - adaptive formatting
+    """
+
+    if not data:
+        console.print("[yellow]No results found.[/yellow]")
+        return
+
+    columns = list(data[0].keys())
+
+    # Special case: single column → vertical list
+    if len(columns) == 1:
+        col = columns[0]
+        console.print(Panel(f"[bold cyan]{col}[/bold cyan]", title=title))
+
+        for row in data:
+            console.print(f"- {row.get(col)}")
+
+        return
+
+    total_rows = len(data)
+    total_pages = math.ceil(total_rows / page_size)
+
+    for page in range(total_pages):
+
+        start = page * page_size
+        end = start + page_size
+        chunk = data[start:end]
+
+        table = Table(
+            title=f"{title} (Page {page+1}/{total_pages})",
+            show_header=True,
+            header_style="bold magenta",
+            box=box.SIMPLE_HEAVY,
+        )
+
+        # Add columns with smart alignment
+        for col in columns:
+            justify = _infer_alignment(chunk, col)
+
+            table.add_column(
+                col,
+                justify=justify,
+                overflow="fold" if wrap else "ellipsis",
+                no_wrap=not wrap,
+            )
+
+        # Add rows
+        for row in chunk:
+            row_values = []
+
+            for col in columns:
+                val = _format_value(row.get(col))
+                val = _truncate(val, 80)
+                row_values.append(val)
+
+            table.add_row(*row_values)
+
+        console.print(table)
+
+        # Pagination pause (non-blocking for automation)
+        if page < total_pages - 1:
+            try:
+                console.print(
+                    "[dim]Press Enter for next page... (or Ctrl+C to stop)[/dim]"
+                )
+                input()
+            except (EOFError, KeyboardInterrupt):
+                console.print("[yellow]Pagination stopped.[/yellow]")
+                break

natquery/cli/main.py

@@ -0,0 +1,51 @@
+import sys
+from rich import print
+from natquery.cli.commands import (
+    connect_command,
+    reset_command,
+    show_config_command,
+    select_database_command,
+)
+from natquery.cli.shell import start_shell
+from natquery.config.settings import Settings
+from natquery.orchestration.workspace import initialize_workspace
+
+
+def main():
+
+    if len(sys.argv) > 1:
+
+        cmd = sys.argv[1]
+
+        if cmd == "connect":
+            # Pass DSN if provided: natquery connect <dsn>
+            dsn = sys.argv[2] if len(sys.argv) > 2 else None
+            connect_command(dsn)
+            return
+
+        elif cmd == "reset":
+            reset_command()
+            return
+
+        elif cmd == "config":
+            show_config_command()
+            return
+
+        elif cmd == "select":
+            select_database_command()
+            return
+
+        else:
+            print(f"Unknown command: {cmd}")
+            return
+
+    # No subcommand → start REPL
+
+    if not Settings.exists():
+        print("NatQuery not configured.")
+        print("Run: natquery connect")
+        return
+
+    print("\n[bold cyan]Starting NatQuery...[/bold cyan]\n")
+    initialize_workspace()
+    start_shell()

natquery/cli/shell.py

@@ -0,0 +1,135 @@
+from rich import print
+from rich.panel import Panel
+from rich.text import Text
+from rich.console import Console
+from natquery.orchestration.pipeline import run_query
+from natquery.llm.client import generate_sql
+from natquery.cli.formatter import print_table
+
+console = Console()
+
+
+def show_banner():
+    banner = Text(
+        """
+███╗   ██╗ █████╗ ████████╗ ██████╗ ██╗   ██╗███████╗██████╗ ██╗   ██╗
+████╗  ██║██╔══██╗╚══██╔══╝██╔═══██╗██║   ██║██╔════╝██╔══██╗╚██╗ ██╔╝
+██╔██╗ ██║███████║   ██║   ██║   ██║██║   ██║█████╗  ██████╔╝ ╚████╔╝ 
+██║╚██╗██║██╔══██║   ██║   ██║▄▄ ██║██║   ██║██╔══╝  ██╔══██╗  ╚██╔╝  
+██║ ╚████║██║  ██║   ██║   ╚██████╔╝╚██████╔╝███████╗██║  ██║   ██║   
+╚═╝  ╚═══╝╚═╝  ╚═╝   ╚═╝    ╚══▀▀═╝  ╚═════╝ ╚══════╝╚═╝  ╚═╝   ╚═╝   
+""",
+        style="bold cyan",
+    )
+
+    print(Panel(banner, title="⚡ NatQuery CLI", border_style="green"))
+
+
+def format_plan_output(summary, suggestions):
+    """Format execution plan and index recommendations for display."""
+    if not summary and not suggestions:
+        print("[dim]No plan data available[/dim]")
+        return
+
+    # Display execution metrics
+    if summary:
+        metrics = []
+        if summary.get("execution_time_ms"):
+            metrics.append(f"Execution Time: {summary['execution_time_ms']:.2f}ms")
+        if summary.get("planning_time_ms"):
+            metrics.append(f"Planning Time: {summary['planning_time_ms']:.2f}ms")
+        if summary.get("total_cost"):
+            metrics.append(f"Total Cost: {summary['total_cost']:.2f}")
+        if summary.get("plan_rows"):
+            metrics.append(f"Planned Rows: {summary['plan_rows']}")
+
+        if metrics:
+            print(
+                Panel(
+                    "\n".join(metrics),
+                    title="Execution Plan",
+                    border_style="yellow",
+                )
+            )
+
+    # Display index recommendations
+    if suggestions and len(suggestions) > 0:
+        suggestion_lines = []
+        for i, sugg in enumerate(suggestions, 1):
+            sugg_type = sugg.get("type", "unknown")
+            reason = sugg.get("reason", "")
+            suggestion_lines.append(f"[bold]{i}. {sugg_type.upper()}[/bold]")
+            if reason:
+                suggestion_lines.append(f"   Reason: {reason}")
+            if sugg.get("sql"):
+                suggestion_lines.append(f"   SQL: {sugg['sql']}")
+            suggestion_lines.append("")
+
+        print(
+            Panel(
+                "\n".join(suggestion_lines),
+                title="Index Recommendations",
+                border_style="cyan",
+            )
+        )
+    elif suggestions is not None and len(suggestions) == 0:
+        print("[dim]No optimization recommendations[/dim]")
+
+
+def start_shell():
+    show_banner()
+    show_sql = False
+    show_plan = False
+    print("[bold green]NatQuery ready.[/bold green]")
+    print(
+        "[dim]Commands: /sql (toggle SQL), /plan (toggle execution plan), exit[/dim]\n"
+    )
+
+    while True:
+        query = console.input("[bold blue]> [/bold blue]")
+
+        if query.lower() in ["exit", "quit"]:
+            print("[bold red]Goodbye![/bold red]")
+            break
+
+        # Toggle SQL display
+        if query.strip() == "/sql":
+            show_sql = not show_sql
+            status = "ON" if show_sql else "OFF"
+            print(f"[yellow]Show SQL:[/yellow] {status}")
+            continue
+
+        # Toggle plan display
+        if query.strip() == "/plan":
+            show_plan = not show_plan
+            status = "ON" if show_plan else "OFF"
+            print(f"[yellow]Show Plan:[/yellow] {status}")
+            continue
+
+        if not query.strip():
+            continue
+
+        try:
+            query_output = run_query(query)
+
+            # Handle both old (list) and new (dict) return formats
+            if isinstance(query_output, dict):
+                result = query_output["result"]
+                summary = query_output.get("summary")
+                suggestions = query_output.get("suggestions")
+            else:
+                result = query_output
+                summary = None
+                suggestions = None
+
+            if show_sql:
+                generated_sql = generate_sql(query)
+                print(Panel(generated_sql, title="Generated SQL", border_style="cyan"))
+
+            print_table(result)
+
+            if show_plan:
+                format_plan_output(summary, suggestions)
+
+        except Exception as e:
+            print(f"[bold red]Error:[/bold red] {e}")

natquery/config/connection.py

@@ -0,0 +1,123 @@
+import psycopg2
+from psycopg2.extensions import connection as PGConnection
+from psycopg2.extras import RealDictCursor
+from natquery.config.settings import Settings
+from natquery.observability.logger import NatQueryLogger
+
+
+def get_connection() -> PGConnection:
+    """
+    Creates and returns a PostgreSQL connection.
+    Supports:
+        - Standard host/port configuration
+        - DSN (connection string)
+        - SSL (optional)
+    """
+
+    db_config = Settings.get_db_config()
+
+    # Determine DB name safely (for logging)
+    db_name = db_config.get("dbname", "unknown_db")
+
+    # Log connection attempt
+    NatQueryLogger.log_event(
+        level="INFO",
+        event="db_connection_attempt",
+        db_name=db_name,
+        conv_id=None,
+        details=db_config,
+    )
+
+    try:
+
+        # ---- DSN MODE ----
+        if db_config["type"] == "dsn":
+
+            conn = psycopg2.connect(
+                db_config["dsn"],
+                connect_timeout=5,
+            )
+
+        # ---- STANDARD MODE ----
+        else:
+
+            connect_args = {
+                "host": db_config["host"],
+                "port": db_config["port"],
+                "dbname": db_config["dbname"],
+                "user": db_config["user"],
+                "password": db_config["password"],
+                "connect_timeout": 5,
+            }
+
+            # Optional SSL
+            sslmode = db_config.get("sslmode")
+            if sslmode:
+                connect_args["sslmode"] = sslmode
+
+            conn = psycopg2.connect(**connect_args)
+
+        conn.autocommit = False
+
+        # Log success
+        NatQueryLogger.log_event(
+            level="INFO",
+            event="db_connection_success",
+            db_name=db_name,
+            conv_id=None,
+        )
+
+        return conn
+
+    except psycopg2.OperationalError as e:
+
+        # Log failure
+        NatQueryLogger.log_event(
+            level="ERROR",
+            event="db_connection_failed",
+            db_name=db_name,
+            conv_id=None,
+            details={"error": str(e)},
+        )
+
+        raise ConnectionError(
+            f"Unable to connect to PostgreSQL server.\nDetails: {e}"
+        ) from e
+
+
+def get_cursor(conn: PGConnection):
+    """
+    Returns a dictionary-style cursor.
+    """
+    return conn.cursor(cursor_factory=RealDictCursor)
+
+
+def close_connection(conn: PGConnection) -> None:
+    """
+    Safely closes a PostgreSQL connection.
+    """
+    if conn:
+        try:
+            # Safely extract db_name; use fallback if connection is partially closed
+            try:
+                db_name = conn.get_dsn_parameters().get("dbname")
+            except (AttributeError, psycopg2.Error):
+                db_name = "unknown_db"
+
+            conn.close()
+
+            NatQueryLogger.log_event(
+                level="INFO",
+                event="db_connection_closed",
+                db_name=db_name,
+                conv_id=None,
+            )
+
+        except Exception as e:
+            NatQueryLogger.log_event(
+                level="ERROR",
+                event="db_connection_close_failed",
+                db_name=None,
+                conv_id=None,
+                details={"error": str(e)},
+            )