nao-core 0.0.29__py3-none-any.whl → 0.0.31__py3-none-any.whl

nao_core/__init__.py

@@ -1,2 +1,2 @@
 # nao Core CLI package
-__version__ = "0.0.29"
+__version__ = "0.0.31"

nao_core/bin/fastapi/main.py

@@ -7,6 +7,9 @@ import os
 import sys
 import numpy as np
 from pathlib import Path
+from dotenv import load_dotenv
+
+load_dotenv()
 
 cli_path = Path(__file__).parent.parent.parent / "cli"
 sys.path.insert(0, str(cli_path))
@@ -85,8 +88,21 @@ async def execute_sql(request: ExecuteSQLRequest):
 
         connection = db_config.connect()
         
-        cursor = connection.raw_sql(request.sql)  # type: ignore[attr-defined]
-        df = cursor.to_dataframe()
+        # Use raw_sql to execute arbitrary SQL (including CTEs)
+        cursor = connection.raw_sql(request.sql)
+        
+        # Handle different cursor types from different backends
+        if hasattr(cursor, 'fetchdf'):
+            # DuckDB returns a cursor with fetchdf()
+            df = cursor.fetchdf()
+        elif hasattr(cursor, 'to_dataframe'):
+            # Some backends return cursors with to_dataframe()
+            df = cursor.to_dataframe()
+        else:
+            # Fallback: try to use pandas read_sql or fetchall
+            import pandas as pd
+            columns = [desc[0] for desc in cursor.description]
+            df = pd.DataFrame(cursor.fetchall(), columns=columns)
         
         def convert_value(v):
             if isinstance(v, (np.integer,)):
@@ -113,4 +129,7 @@ async def execute_sql(request: ExecuteSQLRequest):
         raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
+    nao_project_folder = os.getenv('NAO_PROJECT_FOLDER')
+    if nao_project_folder:
+        os.chdir(nao_project_folder)
     uvicorn.run("main:app", host="0.0.0.0", port=port, reload=True)

nao_core/bin/fastapi/test_main.py

@@ -0,0 +1,156 @@
+import tempfile
+from pathlib import Path
+
+import pytest
+import yaml
+from fastapi.testclient import TestClient
+
+from main import app
+
+
+def assert_sql_result(data: dict, *, row_count: int, columns: list[str], expected_data: list[dict]):
+    """Assert that SQL response data matches expected values."""
+    assert data["row_count"] == row_count
+    assert data["columns"] == columns
+    assert len(data["data"]) == row_count
+    assert data["data"] == expected_data
+
+
+@pytest.fixture
+def duckdb_project_folder():
+    """Create a temporary project folder with a DuckDB config."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "project_name": "test-project",
+            "databases": [
+                {
+                    "name": "test-duckdb",
+                    "type": "duckdb",
+                    "path": ":memory:",
+                }
+            ],
+        }
+        config_path = Path(tmpdir) / "nao_config.yaml"
+        with config_path.open("w") as f:
+            yaml.dump(config, f)
+        yield tmpdir
+
+
+def test_execute_sql_simple_duckdb(duckdb_project_folder):
+    """Test execute_sql endpoint with a DuckDB in-memory database."""
+    client = TestClient(app)
+
+    response = client.post(
+        "/execute_sql",
+        json={
+            "sql": "SELECT 1 AS id, 'hello' AS message",
+            "nao_project_folder": duckdb_project_folder,
+        },
+    )
+
+    assert response.status_code == 200
+    assert_sql_result(
+        response.json(),
+        row_count=1,
+        columns=["id", "message"],
+        expected_data=[{"id": 1, "message": "hello"}],
+    )
+
+
+def test_execute_sql_with_cte_duckdb(duckdb_project_folder):
+    """Test execute_sql endpoint with a DuckDB in-memory database."""
+    client = TestClient(app)
+
+    response = client.post(
+        "/execute_sql",
+        json={
+            "sql": "WITH test AS (SELECT 1 AS id, 'hello' AS message) SELECT * FROM test",
+            "nao_project_folder": duckdb_project_folder,
+        },
+    )
+
+    assert response.status_code == 200
+    assert_sql_result(
+        response.json(),
+        row_count=1,
+        columns=["id", "message"],
+        expected_data=[{"id": 1, "message": "hello"}],
+    )
+
+
+# BigQuery tests (requires SSO authentication)
+
+@pytest.fixture
+def bigquery_project_folder():
+    """Create a temporary project folder with a BigQuery config using SSO."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "project_name": "test-project",
+            "databases": [
+                {
+                    "name": "nao-bigquery",
+                    "type": "bigquery",
+                    "project_id": "nao-corp",
+                    "sso": True,
+                }
+            ],
+        }
+        config_path = Path(tmpdir) / "nao_config.yaml"
+        with config_path.open("w") as f:
+            yaml.dump(config, f)
+        yield tmpdir
+
+
+def test_execute_sql_simple_bigquery(bigquery_project_folder):
+    """Test execute_sql endpoint with BigQuery using SSO."""
+    client = TestClient(app)
+
+    response = client.post(
+        "/execute_sql",
+        json={
+            "sql": "SELECT 1 AS id, 'hello' AS message",
+            "nao_project_folder": bigquery_project_folder,
+        },
+    )
+
+    assert response.status_code == 200
+    assert_sql_result(
+        response.json(),
+        row_count=1,
+        columns=["id", "message"],
+        expected_data=[{"id": 1, "message": "hello"}],
+    )
+
+
+def test_execute_sql_with_cte_bigquery(bigquery_project_folder):
+    """Test execute_sql endpoint with a CTE query on BigQuery."""
+    client = TestClient(app)
+
+    cte_sql = """
+    WITH users AS (
+        SELECT 1 AS id, 'Alice' AS name
+        UNION ALL SELECT 2, 'Bob'
+        UNION ALL SELECT 3, 'Charlie'
+    )
+    SELECT * FROM users
+    """
+
+    response = client.post(
+        "/execute_sql",
+        json={
+            "sql": cte_sql,
+            "nao_project_folder": bigquery_project_folder,
+        },
+    )
+
+    assert response.status_code == 200
+    assert_sql_result(
+        response.json(),
+        row_count=3,
+        columns=["id", "name"],
+        expected_data=[
+            {"id": 1, "name": "Alice"},
+            {"id": 2, "name": "Bob"},
+            {"id": 3, "name": "Charlie"},
+        ],
+    )

nao_core/bin/migrations-postgres/0004_input_and_output_tokens.sql

@@ -0,0 +1,8 @@
+ALTER TABLE "message_part" ADD COLUMN "input_total_tokens" integer;--> statement-breakpoint
+ALTER TABLE "message_part" ADD COLUMN "input_no_cache_tokens" integer;--> statement-breakpoint
+ALTER TABLE "message_part" ADD COLUMN "input_cache_read_tokens" integer;--> statement-breakpoint
+ALTER TABLE "message_part" ADD COLUMN "input_cache_write_tokens" integer;--> statement-breakpoint
+ALTER TABLE "message_part" ADD COLUMN "output_total_tokens" integer;--> statement-breakpoint
+ALTER TABLE "message_part" ADD COLUMN "output_text_tokens" integer;--> statement-breakpoint
+ALTER TABLE "message_part" ADD COLUMN "output_reasoning_tokens" integer;--> statement-breakpoint
+ALTER TABLE "message_part" ADD COLUMN "total_tokens" integer;

nao_core/bin/migrations-postgres/0005_add_project_tables.sql

@@ -0,0 +1,39 @@
+CREATE TABLE "project" (
+	"id" text PRIMARY KEY NOT NULL,
+	"name" text NOT NULL,
+	"type" text NOT NULL,
+	"path" text,
+	"slack_bot_token" text,
+	"slack_signing_secret" text,
+	"created_at" timestamp DEFAULT now() NOT NULL,
+	"updated_at" timestamp DEFAULT now() NOT NULL,
+	CONSTRAINT "local_project_path_required" CHECK (CASE WHEN "project"."type" = 'local' THEN "project"."path" IS NOT NULL ELSE TRUE END)
+);
+--> statement-breakpoint
+CREATE TABLE "project_llm_config" (
+	"id" text PRIMARY KEY NOT NULL,
+	"project_id" text NOT NULL,
+	"provider" text NOT NULL,
+	"api_key" text NOT NULL,
+	"created_at" timestamp DEFAULT now() NOT NULL,
+	"updated_at" timestamp DEFAULT now() NOT NULL,
+	CONSTRAINT "project_llm_config_unique" UNIQUE("id","project_id","provider")
+);
+--> statement-breakpoint
+CREATE TABLE "project_member" (
+	"project_id" text NOT NULL,
+	"user_id" text NOT NULL,
+	"role" text NOT NULL,
+	"created_at" timestamp DEFAULT now() NOT NULL,
+	CONSTRAINT "project_member_project_id_user_id_pk" PRIMARY KEY("project_id","user_id")
+);
+--> statement-breakpoint
+ALTER TABLE "chat" ADD COLUMN "project_id" text NOT NULL;--> statement-breakpoint
+ALTER TABLE "project_llm_config" ADD CONSTRAINT "project_llm_config_project_id_project_id_fk" FOREIGN KEY ("project_id") REFERENCES "public"."project"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
+ALTER TABLE "project_member" ADD CONSTRAINT "project_member_project_id_project_id_fk" FOREIGN KEY ("project_id") REFERENCES "public"."project"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
+ALTER TABLE "project_member" ADD CONSTRAINT "project_member_user_id_user_id_fk" FOREIGN KEY ("user_id") REFERENCES "public"."user"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
+CREATE INDEX "project_llm_config_projectId_idx" ON "project_llm_config" USING btree ("project_id");--> statement-breakpoint
+CREATE INDEX "project_member_userId_idx" ON "project_member" USING btree ("user_id");--> statement-breakpoint
+DELETE FROM "chat";--> statement-breakpoint
+ALTER TABLE "chat" ADD CONSTRAINT "chat_project_id_project_id_fk" FOREIGN KEY ("project_id") REFERENCES "public"."project"("id") ON DELETE cascade ON UPDATE no action;--> statement-breakpoint
+CREATE INDEX "chat_projectId_idx" ON "chat" USING btree ("project_id");