sql-code-graph 1.1.0__py3-none-any.whl → 1.2.2__py3-none-any.whl

sqlcg/lineage/aggregator.py

@@ -22,17 +22,33 @@ class CrossFileAggregator:
         # Maps lowercased table name (bare name) -> exp.Select body for CTAS statements.
         # Populated during register_pass1 and used to seed sources_map in pass 2.
         self.cross_file_sources: dict[str, Any] = {}
+        # #44 canonical-name index: bare name (lowercased) -> sole DDL-defined full_id.
+        # Built from DDL-defined tables only (defined_tables, not CTAS bodies).
+        # Used by _build_file_rows to rewrite an unqualified INSERT-target to the
+        # canonical full_id so INSERT-target nodes share identity with the DDL node.
+        self.canonical_by_bare: dict[str, str] = {}
+        # Bare names defined by >1 schema — do NOT rewrite (ambiguous).
+        self._ambiguous_bare: set[str] = set()
 
     def register_pass1(self, parsed: ParsedFile) -> None:
         """Register a pass-1 result and build view/table source map.
 
-        Also harvests CTAS bodies from statements for cross-file temp-table resolution.
+        Also harvests CTAS bodies from statements for cross-file temp-table resolution,
+        and builds the bare-name → canonical-full_id index (#44) from DDL tables.
 
         Args:
             parsed: ParsedFile from pass 1
         """
         for table in parsed.defined_tables:
             self.sources[table.full_id] = parsed
+            # #44: build canonical_by_bare index from DDL-defined tables
+            bare = (table.name or "").lower()
+            if bare:
+                if bare in self.canonical_by_bare and self.canonical_by_bare[bare] != table.full_id:
+                    # Same bare name defined in multiple schemas → ambiguous, never rewrite
+                    self._ambiguous_bare.add(bare)
+                else:
+                    self.canonical_by_bare[bare] = table.full_id
 
         # Harvest CTAS bodies from statements for cross-file resolution.
         # Key convention matches AnsiParser.parse_file line 109: lowercased bare name.
@@ -68,47 +84,3 @@ class CrossFileAggregator:
                 if bare and bare in self.cross_file_sources and bare not in same_file_bare_names:
                     return True
         return False
-
-    def resolve_pass2(self, parser, parsed: ParsedFile) -> ParsedFile:
-        """Re-parse with cross-file schema context.
-
-        Args:
-            parser: SqlParser instance
-            parsed: ParsedFile from pass 1
-
-        Returns:
-            ParsedFile from pass 2 with resolved cross-file references,
-            or the pass-1 result if skip predicate determines no re-parse is needed
-            or if the file cannot be re-read.
-
-        Raises:
-            No exceptions are raised; file read errors are logged as WARNING
-            and the pass-1 result is returned unchanged.
-
-        Note:
-            This method returns the exact same ParsedFile object (via `return parsed`)
-            on the skip path. This identity semantics are used by callers to track
-            which files were skipped (resolved is parsed). Do not introduce a .copy()
-            on the skip path — that would break the identity check.
-        """
-        # Register view sources for schema resolution
-        parser._schema.add_view_sources(self.sources)
-
-        if not self._needs_pass2(parsed):
-            # File has no cross-file dependencies — pass-1 result is already final.
-            return parsed
-
-        try:
-            sql = parsed.path.read_text(encoding="utf-8")
-        except (FileNotFoundError, OSError) as exc:
-            logger.warning(
-                "resolve_pass2: cannot re-read %s (%s) — returning pass-1 result",
-                parsed.path,
-                exc,
-            )
-            return parsed
-
-        # Filter cross-file CTAS bodies to what this file actually references —
-        # keeps exp.expand bounded by referenced_tables, not by corpus size.
-        ref_names = {(t.name or "").lower() for t in parsed.referenced_tables if t.name}
-        return parser.parse_file(parsed.path, sql, dependency_filter=ref_names)

sqlcg/parsers/ansi_parser.py

@@ -85,8 +85,8 @@ class AnsiParser(SqlParser):
                 the cross-file sources seeded into `sources_map` are filtered to only those
                 whose name is in the set. Pass-1 callers (and direct test callers) pass
                 `None` to disable filtering; pass-2 callers
-                (`CrossFileAggregator.resolve_pass2`) compute this from the pass-1
-                `ParsedFile.referenced_tables`.
+                (the `index_repo` pass-2 dispatch in `indexer.py`) compute this from
+                the pass-1 `ParsedFile.referenced_tables`.
             _precomputed_start_lines: optional list of 1-based start lines, one per
                 statement. When provided (e.g. by SnowflakeParser which computes the map
                 from the preprocessed SQL after ``_preprocess_snowflake_sql`` — which

sqlcg/parsers/base.py

@@ -967,10 +967,16 @@ class SqlParser(ABC):
                             if not isinstance(cte_body, (exp.Select, exp.Union)):
                                 continue
 
-                            # For Union bodies, use the left branch's projections.
+                            # For Union bodies, use the deepest left-branch Select's projections.
                             # Union.expressions is always empty; projections are on Union.this.
+                            # For N=2: cte_body.this is a Select — the while loop is a no-op.
+                            # For N≥3: cte_body.this is a nested Union (A UNION ALL B UNION ALL C
+                            # parses as Union(Union(A,B),C)), so we walk down to the deepest
+                            # left-branch Select (whose star qualify() already expanded in place).
                             if isinstance(cte_body, exp.Union):
                                 projection_source = cte_body.this
+                                while isinstance(projection_source, exp.Union):
+                                    projection_source = projection_source.this
                             else:
                                 projection_source = cte_body
 

sqlcg/server/read_client.py

@@ -0,0 +1,192 @@
+"""Client helper for routing CLI read commands through the live MCP server.
+
+When a server is live on the target DB, CLI read commands route their
+``run_read(cypher, params)`` calls over the Unix control socket instead of
+opening the DB directly.  This avoids "Database is locked" errors when the
+server holds KuzuDB's process-level write lock.
+
+With no server running (``query_via_server`` returns ``None``), the fallback
+opens the DB with ``get_backend(read_only=True)`` — zero-config small-repo
+invariant preserved.
+
+Framing protocol (v1.2.0):
+  Request:  ``<decimal-byte-length>\\n<json-body>``
+  Response: ``<decimal-byte-length>\\n<json-body>``
+  Only the ``query`` op uses this framing; legacy ops (status/stop/reindex)
+  keep their unframed ``{...}\\n`` protocol.
+
+Client receive strategy: after sending the framed request, read the length
+line with ``f.readline()`` (blocking, will not return a partial line) then
+read exactly that many bytes with ``f.read(n)``.  This is the recv-exactly
+pattern required by BLOCKER 2 — a single ``s.recv(65536)`` would silently
+truncate large result sets.  Do NOT copy reindex.py's single-recv pattern
+here.
+
+Server-busy behaviour (v1.1.0 F1 parity):
+  If the server is alive but the lock is held (timeout waiting for the
+  response), raise ``typer.Exit`` — a plain ``Exception`` subclass, NOT
+  ``SystemExit`` / ``BaseException``.  This ensures gain.py's
+  ``except Exception: pass`` handler catches it and degrades gracefully
+  (skips the parse-quality section) instead of crashing.  Other read
+  commands let the ``typer.Exit`` propagate to a clean non-zero CLI exit.
+  Do NOT fall back to a direct open on timeout — the server is alive and
+  holds the lock, so falling back would reproduce the "Database is locked"
+  error (mirrors the F1 fix in reindex.py:127–142).
+"""
+
+from __future__ import annotations
+
+import json
+import socket as _socket
+import sys
+from pathlib import Path
+
+import typer
+
+# Client-side socket timeout for the query control-socket path.
+# Sized to cover the longest in-flight reindex (~89 s DWH resync_changed)
+# with headroom.  This is a CLI transport constant, NOT a KuzuConfig value —
+# same convention as _NOTIFY_SOCKET_TIMEOUT_S in reindex.py.
+_QUERY_SOCKET_TIMEOUT_S = 300
+
+
+def query_via_server(
+    cypher: str,
+    params: dict,
+    db_path: Path | None = None,
+    timeout_s: float = _QUERY_SOCKET_TIMEOUT_S,
+) -> list[dict] | None:
+    """Send a read query over the control socket.
+
+    Uses length-prefixed framing (v1.2.0): ``<len>\\n<json-body>`` for both
+    request and response.  Reads the response with ``makefile`` + ``readline``
+    + ``read(n)`` — NOT a single ``recv`` — so arbitrarily large result sets
+    are returned in full without truncation (BLOCKER 2).
+
+    Args:
+        cypher: Cypher query string (must be read-only; server enforces).
+        params: Query parameter dict.
+        db_path: Explicit database path. Defaults to ``get_db_path()``.
+        timeout_s: Socket timeout in seconds.  On timeout the server is alive
+            and holds the lock — raises ``typer.Exit``, does NOT fall back to
+            a direct open (which would reproduce the lock error).
+
+    Returns:
+        Row list (list[dict]) on success.
+        None when NO server is live (caller should fall back to direct open).
+
+    Raises:
+        typer.Exit: Server is alive but busy (timeout waiting for response).
+            Exception-derived, NOT SystemExit — caught by gain.py's
+            ``except Exception: pass`` so parse-quality section degrades
+            gracefully (WARNING 3).
+        typer.Exit: Server returned ``{"error": ...}`` response.
+    """
+    from sqlcg.server.control import sock_path
+
+    if sys.platform == "win32":
+        # No Unix domain socket on Windows — fall through to direct open.
+        return None
+
+    sp = sock_path(db_path)
+    if not sp.exists():
+        return None
+
+    req = {"op": "query", "cypher": cypher, "params": params}
+    req_bytes = json.dumps(req).encode()
+    frame = f"{len(req_bytes)}\n".encode() + req_bytes
+
+    try:
+        with _socket.socket(_socket.AF_UNIX, _socket.SOCK_STREAM) as s:
+            s.settimeout(timeout_s)
+            s.connect(str(sp))
+            s.sendall(frame)
+
+            # Recv-exactly via makefile:
+            # - f.readline() reads the length line (``<int>\n``) — will not
+            #   return a partial line because makefile buffers internally.
+            # - f.read(n) reads exactly n bytes — accumulates until complete.
+            # A single s.recv(65536) would silently truncate large bodies
+            # (BLOCKER 2 guard: this is the recv-exactly implementation).
+            f = s.makefile("rb")
+            length_line = f.readline()
+            if not length_line:
+                return None  # server closed connection unexpectedly
+            try:
+                body_len = int(length_line.strip())
+            except ValueError:
+                # Server sent an unframed response — protocol mismatch.
+                return None
+            body = f.read(body_len)
+
+    except TimeoutError:
+        # Server is alive and holding the lock.  Do NOT fall back to a direct
+        # open — that would hit the held lock and produce "Database is locked"
+        # (mirrors v1.1.0 F1 fix in reindex.py:127–142).
+        from rich.console import Console
+
+        Console(stderr=True).print(
+            f"[red]Server is busy (reindex in progress); timed out after "
+            f"{timeout_s:.0f}s. The graph will update when it finishes — "
+            "check 'sqlcg mcp status'.[/red]"
+        )
+        raise typer.Exit(1) from None
+    except (FileNotFoundError, ConnectionRefusedError, OSError):
+        # Socket absent or refused — no live server; caller falls back to
+        # direct open.
+        return None
+
+    try:
+        resp = json.loads(body)
+    except (json.JSONDecodeError, ValueError):
+        return None  # malformed response; treat as no-server
+
+    if "error" in resp:
+        from rich.console import Console
+
+        Console(stderr=True).print(f"[red]Server query error: {resp['error']}[/red]")
+        raise typer.Exit(1)
+
+    return resp.get("rows", [])
+
+
+def run_read_routed(
+    cypher: str,
+    params: dict,
+    db_path: Path | None = None,
+) -> list[dict]:
+    """Route through a live server if present, else direct read-only open.
+
+    This is the single seam every CLI read command calls instead of building
+    its own backend.  Centralises the fallback semantics:
+
+    - ``query_via_server`` returns a list → server is live, use rows.
+    - ``query_via_server`` returns None → no server, open DB directly with
+      ``get_backend(read_only=True)`` (BLOCKER 1 — must pass read_only=True
+      or the fallback opens read-write and reproduces lock contention).
+    - ``query_via_server`` raises ``typer.Exit`` → server busy/error; let it
+      propagate (do NOT fall back — lock is held).
+
+    Args:
+        cypher: Cypher query string.
+        params: Query parameter dict.
+        db_path: Explicit database path. Defaults to ``get_db_path()``.
+
+    Returns:
+        Row list from the server or from a direct read-only DB open.
+
+    Raises:
+        typer.Exit: Server busy or server error (propagated from
+            ``query_via_server``).
+    """
+    rows = query_via_server(cypher, params, db_path=db_path)
+    if rows is not None:
+        return rows
+
+    # No server live — fall back to a direct read-only open.
+    # read_only=True is required: without it the fallback opens read-write
+    # and any concurrent writer will produce "Database is locked" (BLOCKER 1).
+    from sqlcg.core.config import get_backend
+
+    with get_backend(read_only=True) as backend:
+        return backend.run_read(cypher, params)

sqlcg/server/server.py

@@ -77,24 +77,39 @@ async def _control_socket_task(
     db_path: "Path",
     backend_ref: "Callable[[], GraphBackend | None]",
     stop_event: "anyio.Event",
-    reindex_lock: "anyio.Lock",
+    backend_lock: "anyio.Lock",
     start_time: float,
 ) -> None:
     """Accept control connections on ``<db>.sock`` and dispatch ops.
 
-    Supported ops (newline-delimited JSON request → response):
+    Supported ops:
 
     - ``{"op": "status"}`` → running state, pid, db_path, freshness, uptime.
+      Unframed (legacy single-recv protocol).
     - ``{"op": "stop"}`` → sends ``{"ok": true}`` then signals stop via
-      *stop_event*.  The ``_run_with_control`` coroutine watches this event
-      and closes stdin to trigger EOF in the MCP stdio loop.
+      *stop_event*.  Unframed.
     - ``{"op": "reindex", "root", "from", "to", "dialect"}`` → runs
       ``Indexer.resync_changed`` off the event-loop thread via
-      ``anyio.to_thread.run_sync``, serialised behind *reindex_lock* (R1, R2).
-
-    R2 (single connection): all backend mutations go through ``reindex_lock``
-    so concurrent notify calls never touch the single Kuzu connection
-    simultaneously.
+      ``anyio.to_thread.run_sync``, serialised behind *backend_lock* (R1, R2).
+      Unframed.
+    - ``{"op": "query", "cypher": ..., "params": ...}`` → executes a
+      read-only Cypher query on the single backend connection, serialised
+      behind *backend_lock*.  **Length-prefixed framing** (v1.2.0):
+      ``<decimal-byte-length>\\n<json-body>`` on both request and response.
+
+    Framing protocol (v1.2.0, ``query`` op only):
+      Request: ``b"<len>\\n" + json_body`` — server detects by sniffing the
+      first line; a bare decimal integer → framed.  Unframed requests always
+      start with ``{`` (never a digit), so the sniff is unambiguous.
+      Response: same ``<len>\\n<body>`` format for framed requests; unframed
+      response for unframed requests.  Old clients that use the unframed
+      ``s.recv(65536)`` + ``json.loads`` pattern will get a loud
+      ``json.JSONDecodeError`` if they accidentally receive a framed response —
+      NOT silent truncation.  Only the new ``read_client`` sends framed
+      requests, so this does not affect existing callers.
+
+    R2 (single connection): all backend operations go through ``backend_lock``
+    so concurrent calls never touch the single Kuzu connection simultaneously.
 
     R8 teardown ordering: the caller must cancel this task BEFORE calling
     ``shutdown_backend()``.  This is guaranteed by the ``anyio.CancelScope``
@@ -108,9 +123,25 @@ async def _control_socket_task(
     import anyio
     import anyio.abc as _anyio_abc
     import anyio.to_thread as _to_thread
+    from anyio.streams.buffered import BufferedByteReceiveStream
 
     from sqlcg.core.config import get_db_path as _get_db_path
 
+    # Read-only keyword allow-list for the ``query`` op.  Only these leading
+    # keywords are permitted — anything that starts with a write keyword is
+    # rejected before execution.  This is a guard against accidental mutation,
+    # not a security boundary (the socket is already 0o600 / owner-only).
+    _QUERY_ALLOWED_KEYWORDS = frozenset({"MATCH", "RETURN", "WITH", "CALL", "UNWIND", "OPTIONAL"})
+
+    def _is_read_only_cypher(cypher: str) -> bool:
+        """Return True iff the leading keyword is in the read-only allow-list."""
+        import re
+
+        m = re.match(r"\s*(?:--[^\n]*)?\s*(\w+)", cypher, re.IGNORECASE)
+        if not m:
+            return False
+        return m.group(1).upper() in _QUERY_ALLOWED_KEYWORDS
+
     sp = sock_path(db_path)
 
     listener = await anyio.create_unix_listener(str(sp))
@@ -119,8 +150,29 @@ async def _control_socket_task(
     async def _handle_connection(stream: _anyio_abc.SocketStream) -> None:
         async with stream:
             try:
-                raw = await stream.receive(4096)
-                req = json.loads(raw)
+                # Sniff for framed vs unframed request.
+                # Framed (query op, v1.2.0): ``<decimal-len>\n<json-body>``
+                # Unframed (legacy status/stop/reindex): JSON object starting with ``{``
+                # The sniff is unambiguous: unframed JSON always starts with ``{``,
+                # never a bare decimal digit.
+                buf = BufferedByteReceiveStream(stream)
+                first_line = await buf.receive_until(b"\n", max_bytes=64)
+
+                try:
+                    body_len = int(first_line.strip())
+                    framed = True
+                except ValueError:
+                    framed = False
+
+                if framed:
+                    # Framed request: read exactly body_len bytes then parse.
+                    raw_body = await buf.receive_exactly(body_len)
+                    req = json.loads(raw_body)
+                else:
+                    # Unframed request (legacy ops): first_line IS the JSON
+                    # (terminated by \n as sent by the client).
+                    req = json.loads(first_line)
+
                 op = req.get("op")
 
                 if op == "status":
@@ -195,15 +247,42 @@ async def _control_socket_task(
                                     dialect,
                                 )
 
-                            async with reindex_lock:
+                            async with backend_lock:
                                 # R1: run off event-loop thread; R2: lock serialises
                                 summary = await _to_thread.run_sync(_do_reindex)
                             resp = {"ok": True, "summary": summary}
 
+                elif op == "query":
+                    # Framed op (v1.2.0): read-only Cypher query over the socket.
+                    # Must only be called with a framed request (sniff above sets framed=True).
+                    cypher = req.get("cypher", "")
+                    params = req.get("params") or {}
+                    if not _is_read_only_cypher(cypher):
+                        resp = {"error": "query op is read-only"}
+                    else:
+                        db = backend_ref()
+                        if db is None:
+                            resp = {"error": "backend not available"}
+                        else:
+
+                            def _do_query() -> list:
+                                return db.run_read(cypher, params)
+
+                            async with backend_lock:
+                                # R1: run off event-loop thread; R2: lock serialises
+                                # reads and writes on the single Kuzu connection.
+                                rows = await _to_thread.run_sync(_do_query)
+                            resp = {"ok": True, "rows": rows}
+
                 else:
                     resp = {"error": f"unknown op: {op!r}"}
 
-                await stream.send(json.dumps(resp).encode() + b"\n")
+                # Send response: framed for framed requests, unframed for legacy ops.
+                resp_bytes = json.dumps(resp).encode()
+                if framed:
+                    await stream.send(f"{len(resp_bytes)}\n".encode() + resp_bytes)
+                else:
+                    await stream.send(resp_bytes + b"\n")
 
             except Exception as exc:
                 try:
@@ -289,16 +368,16 @@ async def _run_with_control(db_path: "Path", start_time: float) -> None:
     with ``abandon_on_cancel=False``).  We cannot interrupt it without
     killing the process; ``_stop_watcher`` does cleanup first.
 
-    ``reindex_lock`` is created once here and passed into
-    ``_control_socket_task`` so concurrent notify calls are serialised
-    behind a single lock (R2).
+    ``backend_lock`` is created once here and passed into
+    ``_control_socket_task`` so concurrent control ops (reindex, query) are
+    serialised behind a single lock on the Kuzu connection (R2).
     """
     import anyio
 
     import sqlcg.server.tools as _tools
 
     stop_event = anyio.Event()
-    reindex_lock = anyio.Lock()  # R2: serialise reindex ops (Kuzu not thread-safe)
+    backend_lock = anyio.Lock()  # R2: serialise all backend ops (Kuzu not thread-safe)
 
     async with anyio.create_task_group() as tg:
         if sys.platform != "win32":
@@ -308,7 +387,7 @@ async def _run_with_control(db_path: "Path", start_time: float) -> None:
                 db_path,
                 lambda: _tools._backend,
                 stop_event,
-                reindex_lock,
+                backend_lock,
                 start_time,
             )
             # Watch stop_event; shuts down and calls os._exit(0).