compass-code 0.1.0__py3-none-any.whl

compass_core/graph/db.py

@@ -0,0 +1,205 @@
+"""LadybugDB wrapper — the primary interface for graph database operations.
+
+This module provides a thin wrapper around the real_ladybug database driver,
+ensuring consistent connection handling and query execution. It is the only
+file in the codebase permitted to import real_ladybug directly.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from pathlib import Path
+from types import TracebackType
+from typing import Any, TypeAlias, cast
+
+import real_ladybug as lb
+
+from compass_core.errors import IndexBusyError, QueryTimeoutError
+from compass_core.settings import settings
+
+logger = logging.getLogger(__name__)
+
+
+QueryParams: TypeAlias = dict[str, Any]
+"""Prepared-statement parameters for ``GraphDB.execute``."""
+
+Row: TypeAlias = Sequence[Any]
+"""One positional result row; only ``query/rows.py`` should interpret its columns."""
+
+
+def as_row(result: lb.QueryResult) -> Sequence[Any]:
+    """Return the next result row as a positionally-indexable sequence.
+
+    Every query executed through this module selects columns positionally
+    (never ``RETURN * AS map``), so ``QueryResult``'s ``list[Any] |
+    dict[str, Any]`` row type never actually resolves to the dict variant
+    here — that's the one shape ``row[0]``-style indexing can't handle.
+    """
+    row = result.get_next()
+    assert not isinstance(row, dict), "expected a positional row, got a dict-row result"
+    return row
+
+
+def as_rows(result: lb.QueryResult) -> list[Row]:
+    """Return every remaining result row as a positionally-indexable sequence. See `as_row`."""
+    rows = result.get_all()
+    for row in rows:
+        assert not isinstance(row, dict), "expected a positional row, got a dict-row result"
+    return cast("list[Row]", rows)
+
+
+class GraphDB:
+    """Thin wrapper around a LadybugDB database and its connection pair.
+
+    This class manages the lifecycle of the database handle and its associated
+    connection, providing a simplified interface for executing Cypher queries.
+
+    Attributes:
+        _db: The underlying LadybugDB database instance.
+        _conn: The active connection to the database instance.
+    """
+
+    def __init__(
+        self,
+        db_path: str | Path | None = None,
+        *,
+        read_only: bool = False,
+        timeout_ms: int | None = None,
+    ) -> None:
+        """Initializes the GraphDB wrapper and opens a database connection.
+
+        If no database path is provided, it defaults to the value specified
+        in the LADYBUG_DB_PATH environment variable via the settings module.
+
+        Args:
+            db_path: Path to the LadybugDB database directory. Defaults to
+                settings.LADYBUG_DB_PATH.
+            read_only: Open the database read-only; write statements fail.
+            timeout_ms: Per-query timeout in milliseconds. Defaults to
+                settings.QUERY_TIMEOUT_MS.
+
+        Raises:
+            IndexBusyError: Another process holds the database file lock.
+        """
+        path = db_path or settings.LADYBUG_DB_PATH
+        try:
+            self._db = lb.Database(str(path), read_only=read_only)
+        except RuntimeError as exc:
+            # Real driver text (verified against a second process holding the
+            # write lock; handles within one process do not conflict):
+            #   "IO exception: Could not set lock on file : <path>\nSee the docs: ..."
+            if "Could not set lock on file" in str(exc):
+                raise IndexBusyError(retry_after_s=2) from exc
+            raise
+        self._timeout_ms = timeout_ms or settings.QUERY_TIMEOUT_MS
+        try:
+            self._conn = lb.Connection(self._db)
+            self._conn.set_query_timeout(self._timeout_ms)
+        except BaseException:
+            self._db.close()
+            raise
+
+    # ------------------------------------------------------------------
+    # Core query interface
+    # ------------------------------------------------------------------
+
+    def execute(self, query: str, params: QueryParams | None = None) -> lb.QueryResult:
+        """Executes a single Cypher query.
+
+        Args:
+            query: The Cypher query string to execute.
+            params: Optional dictionary of parameters to pass as
+                prepared-statement arguments.
+
+        Returns:
+            The raw Result object containing the query output.
+
+        Raises:
+            QueryTimeoutError: The query exceeded the configured timeout.
+        """
+        try:
+            result = self._conn.execute(query, params) if params else self._conn.execute(query)
+        except RuntimeError as exc:
+            # Real driver text for an exceeded set_query_timeout is exactly "Interrupted."
+            if "interrupted" in str(exc).lower():
+                raise QueryTimeoutError(timeout_ms=self._timeout_ms) from exc
+            raise
+        assert isinstance(result, lb.QueryResult), (
+            "a single-statement query must return a single QueryResult, not a list"
+        )
+        return result
+
+    def execute_many(self, statements: list[str]) -> None:
+        """Executes multiple Cypher statements in sequence.
+
+        Results from these statements are ignored. High-level DDL or bulk
+        initialisation should use this method.
+
+        Args:
+            statements: A list of Cypher query strings to execute.
+        """
+        for stmt in statements:
+            self._conn.execute(stmt)
+
+    # ------------------------------------------------------------------
+    # Convenience helpers
+    # ------------------------------------------------------------------
+
+    def node_table_names(self) -> list[str]:
+        """Retrieves the names of all node tables (labels) in the database.
+
+        Returns:
+            A list of strings representing the node table names.
+        """
+        return self._conn._get_node_table_names()
+
+    def rel_table_names(self) -> list[str]:
+        """Retrieves the names of all relationship tables in the database.
+
+        Returns:
+            A list of strings representing the relationship table names.
+        """
+        return [table["name"] for table in self._conn._get_rel_table_names()]
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def close(self) -> None:
+        """Safely closes the connection and the database handle.
+
+        Any errors during closure are logged and suppressed to ensure
+        cleanup completes without crashing the caller.
+        """
+        try:
+            self._conn.close()
+        except Exception:
+            logger.exception("Error closing database connection")
+        try:
+            self._db.close()
+        except Exception:
+            logger.exception("Error closing database handle")
+
+    def __enter__(self) -> GraphDB:
+        """Enables usage of GraphDB as a context manager.
+
+        Returns:
+            The GraphDB instance itself.
+        """
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        """Automatically closes the database on exit from a context block.
+
+        Args:
+            exc_type: The type of the exception raised, if any.
+            exc_val: The exception instance raised, if any.
+            tb: The traceback object, if any.
+        """
+        self.close()

compass_core/graph/schema.py

@@ -0,0 +1,165 @@
+"""Cypher DDL statements to initialise the Compass knowledge graph schema.
+
+This module defines the graph database schema for storing Python source code
+information, including nodes for files, classes, and functions, and their
+inter-relationships.
+"""
+
+from __future__ import annotations
+
+from compass_core.graph.db import GraphDB
+
+# ---------------------------------------------------------------------------
+# DDL — node tables (must be created before rel tables)
+# ---------------------------------------------------------------------------
+
+_NODE_TABLES: list[str] = [
+    # File: A source file entry.
+    #   - path (PK): Repo-relative POSIX path of the source file.
+    #   - module_name: Dotted Python module name (e.g., "os.path").
+    #   - content_hash: SHA-256 hex digest for change detection.
+    """
+    CREATE NODE TABLE IF NOT EXISTS File(
+        path         STRING,
+        module_name  STRING,
+        content_hash STRING,
+        PRIMARY KEY(path)
+    )
+    """,
+    # Class: A class definition extracted from source.
+    #   - id (PK): Public id, format: '<repo-relative path>::ClassName'.
+    #   - name: The unqualified name of the class.
+    #   - qualified_name: Dotted module plus class name, e.g. 'app.models.User'.
+    #   - file_path: Reference to the file where the class is defined.
+    #   - line_start: 1-indexed line where the class definition starts.
+    #   - line_end: 1-indexed line where the class definition ends.
+    #   - source: The entire source code block of the class.
+    #   - docstring: The extracted class-level docstring.
+    #   - bases: A list of raw names for base classes (STRING[]).
+    """
+    CREATE NODE TABLE IF NOT EXISTS Class(
+        id             STRING,
+        name           STRING,
+        qualified_name STRING,
+        file_path      STRING,
+        line_start     INT64,
+        line_end       INT64,
+        source         STRING,
+        docstring      STRING,
+        bases          STRING[],
+        PRIMARY KEY(id)
+    )
+    """,
+    # Function: A function or method definition.
+    #   - id (PK): Public id, format: '<repo-relative path>::[ClassName.]name'.
+    #   - name: The name of the function or method.
+    #   - qualified_name: Dotted module plus '[ClassName.]name', e.g. 'app.models.User.save'.
+    #   - signature: The full 'def' line including arguments and return type.
+    #   - file_path: Reference to the file where the function is defined.
+    #   - line_start: 1-indexed line (including decorators) where it starts.
+    #   - line_end: 1-indexed line where the function body ends.
+    #   - source: The entire source code block of the function.
+    #   - docstring: The extracted function-level docstring.
+    #   - is_method: Boolean indicating if it's defined inside a class.
+    #   - parent_class: The name of the owning class if is_method=True.
+    #   - decorators: A list of decorator names applied (STRING[]).
+    #   - raw_calls: Unresolved call strings within the body (STRING[]).
+    """
+    CREATE NODE TABLE IF NOT EXISTS Function(
+        id             STRING,
+        name           STRING,
+        qualified_name STRING,
+        signature      STRING,
+        file_path      STRING,
+        line_start     INT64,
+        line_end       INT64,
+        source         STRING,
+        docstring      STRING,
+        is_method      BOOLEAN,
+        parent_class   STRING,
+        decorators     STRING[],
+        raw_calls      STRING[],
+        PRIMARY KEY(id)
+    )
+    """,
+    # ExternalFunction: A call to an external library function.
+    #   - full_ref (PK): Fully qualified name (e.g., "jwt.decode").
+    #   - name: The name of the function.
+    #   - library: The name of the library or package.
+    """
+    CREATE NODE TABLE IF NOT EXISTS ExternalFunction(
+        full_ref STRING,
+        name     STRING,
+        library  STRING,
+        PRIMARY KEY(full_ref)
+    )
+    """,
+]
+
+# ---------------------------------------------------------------------------
+# DDL — relationship tables
+# ---------------------------------------------------------------------------
+
+_REL_TABLES: list[str] = [
+    # CONTAINS: Unified containment model (File→Class, File→Function, Class→Function).
+    """
+    CREATE REL TABLE IF NOT EXISTS CONTAINS(
+        FROM File TO Class,
+        FROM File TO Function,
+        FROM Class TO Function
+    )
+    """,
+    # IMPORTS: Dependency between files when one module imports another.
+    """
+    CREATE REL TABLE IF NOT EXISTS IMPORTS(
+        FROM File TO File
+    )
+    """,
+    # INHERITS: Class inheritance relationship (Subclass → Superclass).
+    """
+    CREATE REL TABLE IF NOT EXISTS INHERITS(
+        FROM Class TO Class
+    )
+    """,
+    # CALLS: Internal function execution flow with metadata.
+    #   - confidence (FLOAT): Probability score from the call graph analyzer.
+    #   - call_type (STRING): A ``CallType`` value (e.g., "DIRECT_IMPORT", "MRO").
+    #   - line (INT64): 1-indexed line number where the call occurs.
+    """
+    CREATE REL TABLE IF NOT EXISTS CALLS(
+        FROM Function TO Function,
+        confidence FLOAT,
+        call_type  STRING,
+        line       INT64
+    )
+    """,
+    # CALLS_EXTERN: Function execution flow to external library functions.
+    #   - line (INT64): 1-indexed line number where the call occurs.
+    """
+    CREATE REL TABLE IF NOT EXISTS CALLS_EXTERN(
+        FROM Function TO ExternalFunction,
+        line INT64
+    )
+    """,
+]
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def create_schema(db: GraphDB) -> None:
+    """Create all node and relationship tables if they do not already exist.
+
+    All statements use the 'IF NOT EXISTS' clause to ensure safety when
+    running against an initialized database. Node tables are always
+    created before relationship tables to avoid schema violations.
+
+    Args:
+        db: An instance of GraphDB to execute the DDL statements against.
+    """
+    for stmt in _NODE_TABLES:
+        db.execute(stmt)
+    for stmt in _REL_TABLES:
+        db.execute(stmt)

compass_core/graph/stored_ids.py

@@ -0,0 +1,133 @@
+"""The ids the graph stores for classes and functions: the public ids queries return.
+
+Parsing and resolution key definitions by an internal, language-namespaced id
+(``py::app.models::User.save``). The graph stores the public id instead
+(``app/models.py::User.save``), plus the dotted qualified name, so every lookup by a
+public id is a primary-key match and no reader has to translate between the two.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, replace
+from pathlib import PurePosixPath
+
+from compass_core.ir.graph import ResolvedGraph
+from compass_core.ir.models import Definition
+
+SEP = "::"
+"""Separator between a public id's file path and its ``[Class.]name`` tail."""
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class StoredGraph:
+    """A resolved graph re-keyed by public ids, ready to write.
+
+    Attributes:
+        graph: The graph with every definition id and edge endpoint replaced by its
+            public id. File paths and external references are unchanged.
+        qualified_names: Dotted qualified name of every definition, by public id.
+
+    Example:
+        ``class User`` with method ``save`` in ``app/models.py`` gives definitions with ids
+        ``"app/models.py::User"`` and ``"app/models.py::User.save"``, and
+        ``qualified_names["app/models.py::User.save"] == "app.models.User.save"``.
+    """
+
+    graph: ResolvedGraph
+    qualified_names: Mapping[str, str]
+
+
+def _module_from_path(file_path: str) -> str:
+    parts = list(PurePosixPath(file_path).with_suffix("").parts)
+    if parts and parts[-1] == "__init__":
+        parts.pop()
+    return ".".join(parts)
+
+
+def public_id(definition: Definition) -> str:
+    """The public id of a definition: ``<repo-relative path>::<[Class.]name>``.
+
+    Args:
+        definition: A parsed class, function or method.
+
+    Returns:
+        The public id.
+
+    Example:
+        ``User.save`` in ``app/models.py`` (internal id ``py::app.models::User.save``)
+        gives ``"app/models.py::User.save"``.
+    """
+    return f"{definition.file_path}{SEP}{definition.id.rsplit(SEP, 1)[-1]}"
+
+
+def qualified_name(definition: Definition) -> str:
+    """The dotted qualified name of a definition: module plus ``[Class.]name``.
+
+    Args:
+        definition: A parsed class, function or method.
+
+    Returns:
+        The qualified name; the module comes from the internal id, or from the file
+        path when the id carries none.
+
+    Example:
+        ``User.save`` in ``app/models.py`` gives ``"app.models.User.save"``.
+    """
+    head, _, tail = definition.id.rpartition(SEP)
+    module = head.split(SEP, 1)[1] if SEP in head else _module_from_path(definition.file_path)
+    return f"{module}.{tail}" if module else tail
+
+
+def with_stored_ids(graph: ResolvedGraph) -> StoredGraph:
+    """Re-key a resolved graph by public ids.
+
+    Edge endpoints that are not definitions of the graph keep their value, so they
+    still match no node and are dropped by the writer, as before.
+
+    Args:
+        graph: The graph produced by resolution.
+
+    Returns:
+        A copy with public ids, and the qualified name of every definition.
+
+    Example:
+        A ``CallEdge(caller_id="py::app.api::handle", callee_id="py::app.auth::parse")``
+        becomes ``CallEdge(caller_id="app/api.py::handle", callee_id="app/auth.py::parse")``.
+    """
+    definitions = [d for f in graph.files for d in f.definitions]
+    ids = {d.id: public_id(d) for d in definitions}
+
+    def new(old: str) -> str:
+        return ids.get(old, old)
+
+    files = [
+        replace(
+            f,
+            definitions=[
+                replace(
+                    d, id=new(d.id), parent_id=None if d.parent_id is None else new(d.parent_id)
+                )
+                for d in f.definitions
+            ],
+        )
+        for f in graph.files
+    ]
+    return StoredGraph(
+        graph=ResolvedGraph(
+            files=files,
+            call_edges=[
+                replace(e, caller_id=new(e.caller_id), callee_id=new(e.callee_id))
+                for e in graph.call_edges
+            ],
+            external_call_edges=[
+                replace(e, caller_id=new(e.caller_id)) for e in graph.external_call_edges
+            ],
+            import_edges=list(graph.import_edges),
+            inherit_edges=[
+                replace(e, child_id=new(e.child_id), parent_id=new(e.parent_id))
+                for e in graph.inherit_edges
+            ],
+        ),
+        qualified_names={ids[d.id]: qualified_name(d) for d in definitions},
+    )