feat: add support for postgres schema selection (#475)

* feat: add support for postgres schema selection

Add support for selecting a PostgreSQL schema instead of always using
'public'. The schema is extracted from the connection URL's options
parameter (search_path), following PostgreSQL's native libpq format.

Changes:
- Add _parse_schema_from_url() to extract schema from connection URL
- Thread schema parameter through all extraction methods with 'public' default
- Add pg_namespace JOINs for correct cross-schema disambiguation
- Add schema input field in DatabaseModal (PostgreSQL only)
- Add comprehensive unit tests for URL schema parsing
- Update documentation with custom schema configuration guide

Based on PR #373 by sirudog with the following fixes:
- Fix pg_namespace JOIN order in extract_columns_info to prevent
  duplicate rows when same-named tables exist across schemas
- Fix regex to require '=' separator (prevents mis-capture edge cases)
- Improve $user handling to loop through all schemas instead of only
  checking first two positions
- Fix pylint line-too-long in test file

Co-authored-by: sirudog <1550561+sirudog@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: make parse_schema_from_url public to fix CI pylint

Rename _parse_schema_from_url to parse_schema_from_url since the
method is already documented for external use and tested directly.
This eliminates W0212 (protected-access) warnings that cause CI
pylint to fail with exit code 4.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: address review comments on PR #475

- Add constraint_schema qualifier to key_column_usage JOINs in
  extract_columns_info to prevent cross-schema constraint name
  collisions
- Sanitize schema input in DatabaseModal to strip non-identifier
  characters before building the URL options
- Add edge case tests: empty tokens, blank quoted tokens, repeated
  $user entries

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* chore: remove accidentally committed build artifacts

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: address copilot reviewer comments on PR #475

- Fix regex to capture search_path values with spaces after commas
  (e.g. $user, public) by matching up to next -c option or EOL
- Set session search_path explicitly after connecting so sample
  queries resolve to the correct schema
- Use versionless PostgreSQL docs link (/docs/current/)
- Clarify case-sensitivity note for schema names in troubleshooting

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* chore: gitignore build artifacts

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: replace ReDoS-vulnerable regex in parse_schema_from_url

Replace (.+?)(?=\s+-c|\s*$) with [^\s,]+(?:\s*,\s*[^\s,]+)* to
eliminate polynomial backtracking flagged by CodeQL. The new pattern
uses unambiguous character classes with no overlapping quantifiers.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: validate schema input instead of silent sanitization, fix doc URL encoding

- DatabaseModal: Show validation error for invalid schema characters instead
  of silently stripping them. Throw error on submit if invalid chars present.
- docs: URL-encode the example URL to prevent copy/paste connection failures.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: revert doc URL to readable form to fix spellcheck

The URL-encoded form (-csearch_path%3Dmy_schema) inside the Liquid
capture block triggers spellcheck failures ('csearch', 'Dmy'). Reverted
to readable form since Python's urlparse handles both formats fine.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: add missing tech terms to spellcheck wordlist

Add terms from AGENTS.md/CLAUDE.md (added in staging merge) to the
spellcheck wordlist: config, docstring, dotenv, ESLint, HSTS, init,
Middleware, monorepo, PRs, pylint, pytest, Radix, Zod, and error class
names. Also fix DockerHub capitalization.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: ensure DB connection cleanup on error and add cursor type hints

- Wrap psycopg2 connection/cursor in try/finally so they are always
  closed, even when extract_tables_info or extract_relationships raises
- Set conn/cursor to None after explicit close to avoid double-close
  in the finally block
- Add Any type hints to cursor parameters on extract_tables_info,
  extract_columns_info, extract_foreign_keys, extract_relationships,
  and _execute_sample_query

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix: increase timeout for multi-step E2E chat tests

Mark three tests that perform multiple LLM round-trips with
test.slow() to triple their timeout (60s → 180s), preventing
spurious CI failures when LLM responses are slow:

- multiple sequential queries maintain conversation history
- switching databases clears chat history
- duplicate record shows user-friendly error message

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: sirudog <1550561+sirudog@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: 3 people

.github/wordlist.txt

@@ -8,6 +8,7 @@ schemas
 psycopg
 html
 PostgreSQLLoader
+PostgresLoader
 api
 postgres
 postgresql
@@ -77,6 +78,7 @@ LLM
 Ollama
 OpenAI
 OpenAI's
+DockerHub
 Dockerhub
 FDE
 github
@@ -98,4 +100,21 @@ Sanitization
 JOINs
 subqueries
 subquery
-TTL
+TTL
+
+config
+docstring
+dotenv
+ESLint
+GraphNotFoundError
+HSTS
+init
+InternalError
+InvalidArgumentError
+Middleware
+monorepo
+PRs
+pylint
+pytest
+Radix
+Zod

.gitignore

@@ -31,4 +31,7 @@ demo_tokens.py
 /blob-report/
 /playwright/.cache/
 /playwright/.auth/
-e2e/.auth/
+e2e/.auth/
+# Build artifacts
+clients/python/queryweaver_client.egg-info/
+clients/ts/dist/

README.md

@@ -231,6 +231,7 @@ with requests.post(url, headers=headers, json={"chat": ["Count orders last week"
                         continue
                   obj = json.loads(part)
                   print('STREAM:', obj)
+```
 
 Notes & tips
 - Graph IDs are namespaced per-user. When calling the API directly use the plain graph id (the server will namespace by the authenticated user). For uploaded files the `database` field determines the saved graph id.

api/loaders/postgres_loader.py

@@ -5,6 +5,7 @@
 import decimal
 import logging
 from typing import AsyncGenerator, Dict, Any, List, Tuple
+from urllib.parse import urlparse, parse_qs, unquote
 
 import psycopg2
 from psycopg2 import sql
@@ -52,7 +53,7 @@ class PostgresLoader(BaseLoader):
 
     @staticmethod
     def _execute_sample_query(
-        cursor, table_name: str, col_name: str, sample_size: int = 3
+        cursor: Any, table_name: str, col_name: str, sample_size: int = 3
     ) -> List[Any]:
         """
         Execute query to get random sample values for a column.
@@ -96,39 +97,96 @@ def _serialize_value(value):
             return None
         return value
 
+    @staticmethod
+    def parse_schema_from_url(connection_url: str) -> str:
+        """
+        Parse the search_path from the connection URL's options parameter.
+
+        The options parameter follows PostgreSQL's libpq format:
+        postgresql://user:pass@host:port/db?options=-csearch_path%3Dschema_name
+
+        Args:
+            connection_url: PostgreSQL connection URL
+
+        Returns:
+            The first schema from search_path, or 'public' if not specified
+        """
+        try:
+            parsed = urlparse(connection_url)
+            query_params = parse_qs(parsed.query)
+
+            options = query_params.get('options', [])
+            if not options:
+                return 'public'
+
+            options_str = unquote(options[0])
+
+            # Parse -c search_path=value from options
+            # Format can be: -csearch_path=schema or -c search_path=schema
+            # Match comma-separated schema tokens (supports spaces after commas).
+            match = re.search(r'-c\s*search_path\s*=\s*([^\s,]+(?:\s*,\s*[^\s,]+)*)', options_str, re.IGNORECASE)
+            if match:
+                search_path = match.group(1)
+                schemas = search_path.split(',')
+                for s in schemas:
+                    s = s.strip().strip('"\'')
+                    if s and s != '$user':
+                        return s
+                return 'public'
+
+            return 'public'
+
+        except Exception:  # pylint: disable=broad-exception-caught
+            return 'public'
+
     @staticmethod
     async def load(prefix: str, connection_url: str) -> AsyncGenerator[tuple[bool, str], None]:
         """
         Load the graph data from a PostgreSQL database into the graph database.
 
         Args:
             connection_url: PostgreSQL connection URL in format:
-                          postgresql://username:password@host:port/database
+                        postgresql://username:password@host:port/database
+                        Optionally with schema via options parameter:
+                        postgresql://...?options=-csearch_path%3Dschema_name
 
         Returns:
             Tuple[bool, str]: Success status and message
         """
+        conn = None
+        cursor = None
         try:
+            # Parse schema from connection URL (defaults to 'public')
+            schema = PostgresLoader.parse_schema_from_url(connection_url)
+
             # Connect to PostgreSQL database
             conn = psycopg2.connect(connection_url)
             cursor = conn.cursor()
 
+            # Set the session search_path to the parsed schema so unqualified
+            # table references (e.g. in sample queries) resolve correctly.
+            cursor.execute(
+                sql.SQL("SET search_path TO {}").format(sql.Identifier(schema))
+            )
+
             # Extract database name from connection URL
             db_name = connection_url.split('/')[-1]
             if '?' in db_name:
                 db_name = db_name.split('?')[0]
 
             # Get all table information
             yield True, "Extracting table information..."
-            entities = PostgresLoader.extract_tables_info(cursor)
+            entities = PostgresLoader.extract_tables_info(cursor, schema)
 
             yield True, "Extracting relationship information..."
             # Get all relationship information
-            relationships = PostgresLoader.extract_relationships(cursor)
+            relationships = PostgresLoader.extract_relationships(cursor, schema)
 
-            # Close database connection
+            # Close database connection before graph loading
             cursor.close()
+            cursor = None
             conn.close()
+            conn = None
 
             yield True, "Loading data into graph..."
             # Load data into graph
@@ -144,46 +202,53 @@ async def load(prefix: str, connection_url: str) -> AsyncGenerator[tuple[bool, s
         except Exception as e:  # pylint: disable=broad-exception-caught
             logging.error("Error loading PostgreSQL schema: %s", e)
             yield False, "Failed to load PostgreSQL database schema"
+        finally:
+            if cursor is not None:
+                cursor.close()
+            if conn is not None:
+                conn.close()
 
     @staticmethod
-    def extract_tables_info(cursor) -> Dict[str, Any]:
+    def extract_tables_info(cursor: Any, schema: str = 'public') -> Dict[str, Any]:
         """
         Extract table and column information from PostgreSQL database.
 
         Args:
             cursor: Database cursor
+            schema: Database schema to extract tables from (default: 'public')
 
         Returns:
             Dict containing table information
         """
         entities = {}
 
-        # Get all tables in public schema
+        # Get all tables in the specified schema
         cursor.execute("""
             SELECT table_name, table_comment
             FROM information_schema.tables t
             LEFT JOIN (
                 SELECT schemaname, tablename, description as table_comment
                 FROM pg_tables pt
                 JOIN pg_class pc ON pc.relname = pt.tablename
+                JOIN pg_namespace pn ON pn.oid = pc.relnamespace AND pn.nspname = pt.schemaname
                 JOIN pg_description pd ON pd.objoid = pc.oid AND pd.objsubid = 0
-                WHERE pt.schemaname = 'public'
+                WHERE pt.schemaname = %s
             ) tc ON tc.tablename = t.table_name
-            WHERE t.table_schema = 'public'
+            WHERE t.table_schema = %s
             AND t.table_type = 'BASE TABLE'
             ORDER BY t.table_name;
-        """)
+        """, (schema, schema))
 
         tables = cursor.fetchall()
 
         for table_name, table_comment in tqdm.tqdm(tables, desc="Extracting table information"):
             table_name = table_name.strip()
 
             # Get column information for this table
-            columns_info = PostgresLoader.extract_columns_info(cursor, table_name)
+            columns_info = PostgresLoader.extract_columns_info(cursor, table_name, schema)
 
             # Get foreign keys for this table
-            foreign_keys = PostgresLoader.extract_foreign_keys(cursor, table_name)
+            foreign_keys = PostgresLoader.extract_foreign_keys(cursor, table_name, schema)
 
             # Generate table description
             table_description = table_comment if table_comment else f"Table: {table_name}"
@@ -201,13 +266,14 @@ def extract_tables_info(cursor) -> Dict[str, Any]:
         return entities
 
     @staticmethod
-    def extract_columns_info(cursor, table_name: str) -> Dict[str, Any]:
+    def extract_columns_info(cursor: Any, table_name: str, schema: str = 'public') -> Dict[str, Any]:
         """
         Extract column information for a specific table.
 
         Args:
             cursor: Database cursor
             table_name: Name of the table
+            schema: Database schema (default: 'public')
 
         Returns:
             Dict containing column information
@@ -230,24 +296,29 @@ def extract_columns_info(cursor, table_name: str) -> Dict[str, Any]:
                 FROM information_schema.table_constraints tc
                 JOIN information_schema.key_column_usage ku
                     ON tc.constraint_name = ku.constraint_name
+                    AND tc.constraint_schema = ku.constraint_schema
                 WHERE tc.table_name = %s
+                AND tc.table_schema = %s
                 AND tc.constraint_type = 'PRIMARY KEY'
             ) pk ON pk.column_name = c.column_name
             LEFT JOIN (
                 SELECT ku.column_name
                 FROM information_schema.table_constraints tc
                 JOIN information_schema.key_column_usage ku
                     ON tc.constraint_name = ku.constraint_name
+                    AND tc.constraint_schema = ku.constraint_schema
                 WHERE tc.table_name = %s
+                AND tc.table_schema = %s
                 AND tc.constraint_type = 'FOREIGN KEY'
             ) fk ON fk.column_name = c.column_name
-            LEFT JOIN pg_class pc ON pc.relname = c.table_name
+            LEFT JOIN pg_namespace pn ON pn.nspname = c.table_schema
+            LEFT JOIN pg_class pc ON pc.relname = c.table_name AND pc.relnamespace = pn.oid
             LEFT JOIN pg_attribute pa ON pa.attrelid = pc.oid AND pa.attname = c.column_name
             LEFT JOIN pg_description pgd ON pgd.objoid = pc.oid AND pgd.objsubid = pa.attnum
             WHERE c.table_name = %s
-            AND c.table_schema = 'public'
+            AND c.table_schema = %s
             ORDER BY c.ordinal_position;
-        """, (table_name, table_name, table_name))
+        """, (table_name, schema, table_name, schema, table_name, schema))
 
         columns = cursor.fetchall()
         columns_info = {}
@@ -289,13 +360,14 @@ def extract_columns_info(cursor, table_name: str) -> Dict[str, Any]:
         return columns_info
 
     @staticmethod
-    def extract_foreign_keys(cursor, table_name: str) -> List[Dict[str, str]]:
+    def extract_foreign_keys(cursor: Any, table_name: str, schema: str = 'public') -> List[Dict[str, str]]:
         """
         Extract foreign key information for a specific table.
 
         Args:
             cursor: Database cursor
             table_name: Name of the table
+            schema: Database schema (default: 'public')
 
         Returns:
             List of foreign key dictionaries
@@ -315,8 +387,8 @@ def extract_foreign_keys(cursor, table_name: str) -> List[Dict[str, str]]:
                 AND ccu.table_schema = tc.table_schema
             WHERE tc.constraint_type = 'FOREIGN KEY'
             AND tc.table_name = %s
-            AND tc.table_schema = 'public';
-        """, (table_name,))
+            AND tc.table_schema = %s;
+        """, (table_name, schema))
 
         foreign_keys = []
         for constraint_name, column_name, foreign_table, foreign_column in cursor.fetchall():
@@ -330,12 +402,13 @@ def extract_foreign_keys(cursor, table_name: str) -> List[Dict[str, str]]:
         return foreign_keys
 
     @staticmethod
-    def extract_relationships(cursor) -> Dict[str, List[Dict[str, str]]]:
+    def extract_relationships(cursor: Any, schema: str = 'public') -> Dict[str, List[Dict[str, str]]]:
         """
         Extract all relationship information from the database.
 
         Args:
             cursor: Database cursor
+            schema: Database schema (default: 'public')
 
         Returns:
             Dict containing relationship information
@@ -355,9 +428,9 @@ def extract_relationships(cursor) -> Dict[str, List[Dict[str, str]]]:
                 ON ccu.constraint_name = tc.constraint_name
                 AND ccu.table_schema = tc.table_schema
             WHERE tc.constraint_type = 'FOREIGN KEY'
-            AND tc.table_schema = 'public'
+            AND tc.table_schema = %s
             ORDER BY tc.table_name, tc.constraint_name;
-        """)
+        """, (schema,))
 
         relationships = {}
         for (table_name, constraint_name, column_name,