resolvekit 0.0.1__py3-none-any.whl

resolvekit/data/db_manager.py

@@ -0,0 +1,196 @@
+"""Database connection manager."""
+
+from contextlib import AbstractContextManager, suppress
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.pool import StaticPool
+
+if TYPE_CHECKING:
+    from sqlalchemy.engine import Connection
+
+from resolvekit.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+class DatabaseManager:
+    """
+    Manages SQLite database connections with overlay support.
+
+    Handles:
+    - Connection creation with SQLAlchemy
+    - Performance PRAGMA application
+    - Overlay database attachment
+    - Transaction management
+    """
+
+    def __init__(
+        self,
+        base_path: Path,
+        overlays: list[Path] | None = None,
+        read_only: bool = True,
+    ):
+        """
+        Initialize database manager.
+
+        Args:
+            base_path: Path to base database file
+            overlays: Optional list of overlay databases (max 5, ordered by precedence)
+            read_only: Whether database is read-only
+
+        Raises:
+            ValueError: If database paths are not unique or too many overlays
+        """
+        # Validate overlay count
+        if overlays and len(overlays) > 5:
+            raise ValueError(
+                f"Maximum 5 overlays supported, got {len(overlays)}. "
+                "Too many overlays can impact performance."
+            )
+
+        # Validate that all paths are unique
+        paths = [base_path]
+        if overlays:
+            paths.extend(overlays)
+
+        # Resolve paths to catch symbolic links and relative paths
+        resolved_paths = [p.resolve() for p in paths]
+
+        if len(resolved_paths) != len(set(resolved_paths)):
+            raise ValueError(
+                f"Database paths must be unique. Got duplicates in: {paths}"
+            )
+
+        self.base_path = base_path
+        self.overlay_paths = overlays or []
+        self.read_only = read_only
+        self.overlays: list[tuple[str, int]] = []
+        self.engine: Engine | None = None
+
+    def connect(self) -> None:
+        """Initialize connection with pragmas and overlay attachment."""
+        # Create SQLAlchemy engine
+        uri = f"sqlite:///{self.base_path}"
+        self.engine = create_engine(
+            uri,
+            connect_args={"check_same_thread": False},
+            poolclass=StaticPool,
+            echo=False,
+        )
+
+        # Apply pragmas
+        self._apply_pragmas()
+
+        # Attach overlays
+        self._attach_overlays()
+
+    def _apply_pragmas(self) -> None:
+        """Apply performance pragmas from CLAUDE.md."""
+        if not self.engine:
+            raise RuntimeError("Engine not initialized")
+
+        pragmas = [
+            "PRAGMA foreign_keys=ON",
+            "PRAGMA journal_mode=OFF",
+            "PRAGMA synchronous=OFF",
+            "PRAGMA temp_store=MEMORY",
+            "PRAGMA mmap_size=268435456",
+            "PRAGMA cache_size=-100000",
+        ]
+
+        # Add read-only enforcement if requested
+        if self.read_only:
+            pragmas.append("PRAGMA query_only=ON")
+
+        with self.engine.connect() as conn:
+            for pragma in pragmas:
+                conn.execute(text(pragma))
+            conn.commit()
+
+    def _attach_overlays(self) -> None:
+        """
+        Attach overlay databases with precedence tracking.
+
+        Precedence is assigned based on list order:
+        - First overlay: precedence 100 (highest)
+        - Second overlay: precedence 99
+        - Third overlay: precedence 98
+        - etc.
+        """
+        if not self.engine:
+            raise RuntimeError("Engine not initialized")
+
+        # Clear overlays list to prevent duplicates on reconnect
+        self.overlays.clear()
+
+        with self.engine.connect() as conn:
+            # Attach overlays in order, assigning descending precedence
+            for idx, overlay_path in enumerate(self.overlay_paths):
+                schema_name = f"overlay_{idx}"
+                precedence = 100 - idx  # First=100, second=99, etc.
+
+                conn.execute(text(f"ATTACH DATABASE '{overlay_path}' AS {schema_name}"))
+                self.overlays.append((schema_name, precedence))
+
+            conn.commit()
+
+    def execute(self, query: str, params: dict[str, Any] | None = None) -> Any:
+        """
+        Execute a query.
+
+        Args:
+            query: SQL query string
+            params: Optional parameters dict
+
+        Returns:
+            List of Row objects for SELECT queries, empty list for non-SELECT
+        """
+        if not self.engine:
+            raise RuntimeError("Database not connected. Call connect() first.")
+
+        with self.engine.connect() as conn:
+            result = conn.execute(text(query), params or {})
+            conn.commit()
+            # Consume results before connection closes (if query returns rows)
+            if result.returns_rows:
+                return result.fetchall()
+            return []
+
+    def transaction(self) -> AbstractContextManager["Connection"]:
+        """
+        Get transaction context manager.
+
+        Returns:
+            Transaction context manager
+        """
+        if not self.engine:
+            raise RuntimeError("Database not connected. Call connect() first.")
+
+        return self.engine.begin()
+
+    def close(self) -> None:
+        """Close database connection and detach overlays."""
+        if self.engine:
+            # Detach overlays before disposing engine
+            if self.overlays:
+                try:
+                    with self.engine.connect() as conn:
+                        for overlay_name, _ in self.overlays:
+                            conn.execute(text(f"DETACH DATABASE {overlay_name}"))
+                        conn.commit()
+                except Exception as e:
+                    # Log but don't fail - connection may already be closed
+                    logger.debug(f"Error detaching overlays during close: {e}")
+                finally:
+                    self.overlays.clear()
+
+            self.engine.dispose()
+            self.engine = None
+
+    def __del__(self) -> None:
+        """Cleanup: Close database connection when object is garbage collected."""
+        with suppress(Exception):
+            self.close()

resolvekit/data/entity_repository.py

@@ -0,0 +1,466 @@
+"""Entity repository."""
+
+from datetime import date
+from typing import Any, ClassVar, overload
+
+from resolvekit.data.context_filters import ContextFilterBuilder
+from resolvekit.data.db_manager import DatabaseManager
+from resolvekit.data.query_builder import QueryBuilder
+from resolvekit.types import Entity, EntityRow, MatchContext
+
+
+class EntityRepository:
+    """Repository for entity operations."""
+
+    # Derive column list from EntityRow model (single source of truth)
+    ENTITY_COLUMNS: ClassVar[list[str]] = list(EntityRow.model_fields.keys())
+
+    def __init__(self, db_manager: DatabaseManager):
+        """
+        Initialize repository.
+
+        Args:
+            db_manager: Database manager instance
+        """
+        self.db = db_manager
+        self.query_builder = QueryBuilder(db_manager)
+
+    @overload
+    def find_by_dcid(
+        self,
+        dcid: str,
+        as_of: date | None = None,
+        include_codes: bool = False,
+    ) -> Entity | None: ...
+
+    @overload
+    def find_by_dcid(
+        self,
+        dcid: list[str],
+        as_of: date | None = None,
+        include_codes: bool = False,
+    ) -> dict[str, Entity]: ...
+
+    def find_by_dcid(
+        self,
+        dcid: str | list[str],
+        as_of: date | None = None,
+        include_codes: bool = False,
+    ) -> Entity | None | dict[str, Entity]:
+        """
+        Find entity by DCID.
+
+        Args:
+            dcid: Single DCID or list of DCIDs
+            as_of: Optional date for temporal filtering
+            include_codes: Whether to include code mappings
+
+        Returns:
+            Single entity, dict of entities, or None
+        """
+        if isinstance(dcid, str):
+            return self._find_single_by_dcid(dcid, as_of, include_codes)
+        else:
+            return self._find_batch_by_dcid(dcid, as_of, include_codes)
+
+    def _find_single_by_dcid(
+        self,
+        dcid: str,
+        as_of: date | None = None,
+        include_codes: bool = False,
+    ) -> Entity | None:
+        """Find single entity by DCID."""
+        sql, params = self.query_builder.build_union_query(
+            table="entities",
+            columns=self.ENTITY_COLUMNS,
+            where_clause="dcid = :dcid",
+            params={"dcid": dcid},
+            unique_key="dcid",
+            as_of=as_of,
+        )
+
+        rows = self.db.execute(sql, params)
+
+        if not rows:
+            return None
+
+        row = rows[0]
+        entity = self._row_to_entity(row)
+
+        # Load codes if requested
+        if include_codes:
+            self._load_codes_for_entities([entity])
+
+        return entity
+
+    def _find_batch_by_dcid(
+        self,
+        dcids: list[str],
+        as_of: date | None = None,
+        include_codes: bool = False,
+    ) -> dict[str, Entity]:
+        """Find multiple entities by DCID."""
+        if not dcids:
+            return {}
+
+        # Build IN clause
+        placeholders = ", ".join(f":dcid_{i}" for i in range(len(dcids)))
+        params = {f"dcid_{i}": dcid for i, dcid in enumerate(dcids)}
+
+        sql, params = self.query_builder.build_union_query(
+            table="entities",
+            columns=self.ENTITY_COLUMNS,
+            where_clause=f"dcid IN ({placeholders})",
+            params=params,
+            unique_key="dcid",
+            as_of=as_of,
+        )
+
+        result = self.db.execute(sql, params)
+
+        entities = {}
+        for row in result:
+            entity = self._row_to_entity(row)
+            entities[entity.dcid] = entity
+
+        # Load codes if requested
+        if include_codes and entities:
+            self._load_codes_for_entities(list(entities.values()))
+
+        return entities
+
+    @overload
+    def find_by_code(
+        self, code_system: str, code_value: str, context: MatchContext | None = None
+    ) -> Entity | None: ...
+
+    @overload
+    def find_by_code(
+        self,
+        code_system: str,
+        code_value: list[str],
+        context: MatchContext | None = None,
+    ) -> dict[str, Entity]: ...
+
+    def find_by_code(
+        self,
+        code_system: str,
+        code_value: str | list[str],
+        context: MatchContext | None = None,
+    ) -> Entity | None | dict[str, Entity]:
+        """
+        Find entity by code.
+
+        Args:
+            code_system: Code system (e.g., "iso2", "iso3")
+            code_value: Single code or list of codes
+            context: Optional filtering context
+
+        Returns:
+            Single entity, dict of entities, or None
+        """
+        if isinstance(code_value, str):
+            return self._find_single_by_code(code_system, code_value, context)
+        else:
+            return self._find_batch_by_code(code_system, code_value, context)
+
+    def _find_single_by_code(
+        self, code_system: str, code_value: str, context: MatchContext | None = None
+    ) -> Entity | None:
+        """Find single entity by code."""
+        sql, params = self.query_builder.build_code_lookup_union(
+            code_system=code_system,
+            code_values=code_value,
+            entity_columns=self.ENTITY_COLUMNS,
+            include_code_value=False,
+        )
+
+        # Apply context filters if provided
+        if context:
+            # Build context filters (no table prefix - subquery columns don't have aliases)
+            temporal_filter, type_filter, parent_filter = (
+                ContextFilterBuilder.build_filters(context, params, table_prefix="")
+            )
+            # Wrap the union query and apply filters
+            sql = f"""
+                SELECT * FROM ({sql}) AS filtered
+                WHERE 1=1
+                {temporal_filter}
+                {type_filter}
+                {parent_filter}
+            """
+
+        # Order by precedence DESC to respect overlay precedence
+        sql += "\nORDER BY precedence DESC LIMIT 1"
+
+        rows = self.db.execute(sql, params)
+
+        if not rows:
+            return None
+
+        row = rows[0]
+        return self._row_to_entity(row)
+
+    def _find_batch_by_code(
+        self,
+        code_system: str,
+        code_values: list[str],
+        context: MatchContext | None = None,
+    ) -> dict[str, Entity]:
+        """Find multiple entities by code."""
+        if not code_values:
+            return {}
+
+        sql, params = self.query_builder.build_code_lookup_union(
+            code_system=code_system,
+            code_values=code_values,
+            entity_columns=self.ENTITY_COLUMNS,
+            include_code_value=True,
+        )
+
+        # Apply context filters if provided
+        if context:
+            # Build context filters (no table prefix - subquery columns don't have aliases)
+            temporal_filter, type_filter, parent_filter = (
+                ContextFilterBuilder.build_filters(context, params, table_prefix="")
+            )
+            # Wrap the union query and apply filters
+            sql = f"""
+                SELECT * FROM ({sql}) AS filtered
+                WHERE 1=1
+                {temporal_filter}
+                {type_filter}
+                {parent_filter}
+            """
+
+        # Order by code_value and precedence to group overlays together
+        sql += "\nORDER BY code_value, precedence DESC"
+
+        result = self.db.execute(sql, params)
+
+        # Deduplicate by code_value, keeping highest precedence
+        entities = {}
+        for row in result:
+            code_val = row.code_value
+            # Only add if we haven't seen this code_value or this has higher precedence
+            if code_val not in entities:
+                entity = self._row_to_entity(row)
+                entities[code_val] = entity
+
+        return entities
+
+    def get_parent(self, dcid: str, as_of: date | None = None) -> Entity | None:
+        """
+        Get parent entity.
+
+        Args:
+            dcid: Entity DCID
+            as_of: Optional date for temporal filtering
+
+        Returns:
+            Parent entity or None
+        """
+        # Optimized single-query version using JOIN
+        entity_cols = ", ".join(f"parent.{col}" for col in self.ENTITY_COLUMNS)
+
+        # Build temporal filter if needed
+        temporal_filter = ""
+        params: dict[str, Any] = {"dcid": dcid}
+        if as_of is not None:
+            # Apply temporal filter to BOTH child and parent to ensure correct historical row selection
+            temporal_filter = """
+                AND (child.valid_from IS NULL OR child.valid_from <= :as_of)
+                AND (child.valid_until IS NULL OR child.valid_until >= :as_of)
+                AND (parent.valid_from IS NULL OR parent.valid_from <= :as_of)
+                AND (parent.valid_until IS NULL OR parent.valid_until >= :as_of)
+            """
+            params["as_of"] = as_of.isoformat()
+
+        # Query from main database
+        sql = f"""
+            SELECT {entity_cols}, 0 AS precedence, 'main' AS source_db
+            FROM main.entities child
+            INNER JOIN main.entities parent ON child.parent_dcid = parent.dcid
+            WHERE child.dcid = :dcid{temporal_filter}
+        """
+
+        # Add overlays if present (join child from any source with parent from any source)
+        for schema_name, precedence in self.db.overlays:
+            # Child in overlay, parent in overlay
+            sql += f"""
+                UNION ALL
+                SELECT {entity_cols}, {precedence} AS precedence, '{schema_name}' AS source_db
+                FROM {schema_name}.entities child
+                INNER JOIN {schema_name}.entities parent ON child.parent_dcid = parent.dcid
+                WHERE child.dcid = :dcid{temporal_filter}
+            """
+            # Child in overlay, parent in main
+            sql += f"""
+                UNION ALL
+                SELECT {entity_cols}, {precedence // 2} AS precedence, 'mixed' AS source_db
+                FROM {schema_name}.entities child
+                INNER JOIN main.entities parent ON child.parent_dcid = parent.dcid
+                WHERE child.dcid = :dcid{temporal_filter}
+            """
+            # Child in main, parent in overlay (respect overlay precedence)
+            sql += f"""
+                UNION ALL
+                SELECT {entity_cols}, {precedence} AS precedence, '{schema_name}' AS source_db
+                FROM main.entities child
+                INNER JOIN {schema_name}.entities parent ON child.parent_dcid = parent.dcid
+                WHERE child.dcid = :dcid{temporal_filter}
+            """
+
+        # Order by precedence to get highest precedence result
+        sql += "\nORDER BY precedence DESC LIMIT 1"
+
+        rows = self.db.execute(sql, params)
+
+        if not rows:
+            return None
+
+        row = rows[0]
+        return self._row_to_entity(row)
+
+    def get_children(self, dcid: str, as_of: date | None = None) -> list[Entity]:
+        """
+        Get child entities.
+
+        Args:
+            dcid: Parent entity DCID
+            as_of: Optional date for temporal filtering
+
+        Returns:
+            List of child entities
+        """
+        sql, params = self.query_builder.build_union_query(
+            table="entities",
+            columns=self.ENTITY_COLUMNS,
+            where_clause="parent_dcid = :parent_dcid",
+            params={"parent_dcid": dcid},
+            unique_key="dcid",
+            as_of=as_of,
+        )
+
+        result = self.db.execute(sql, params)
+
+        return [self._row_to_entity(row) for row in result]
+
+    def _row_to_entity(self, row: Any) -> Entity:
+        """Convert database row to Entity model via Pydantic validation."""
+        # Convert SQLAlchemy Row to dict using only entity columns
+        row_dict = {col: getattr(row, col) for col in self.ENTITY_COLUMNS}
+
+        # Validate through EntityRow model
+        entity_row = EntityRow.model_validate(row_dict)
+
+        # Convert EntityRow to Entity (adding computed fields)
+        return Entity(
+            **entity_row.model_dump(),
+            codes={},
+            provenance={},
+        )
+
+    def find_by_canonical_name(
+        self, normalized_canonical: str, context: MatchContext | None = None
+    ) -> Entity | None:
+        """
+        Find entity by normalized canonical name.
+
+        Uses the indexed normalized_canonical column for efficient lookups.
+
+        Args:
+            normalized_canonical: Normalized canonical name
+            context: Optional filtering context
+
+        Returns:
+            Entity if found, None otherwise
+        """
+        params: dict[str, Any] = {"normalized": normalized_canonical}
+
+        # Build context filters using shared utility
+        temporal_filter, type_filter, parent_filter = (
+            ContextFilterBuilder.build_filters(context, params, table_prefix="")
+        )
+
+        # Query from main database using indexed normalized_canonical column
+        entity_cols = ", ".join(self.ENTITY_COLUMNS)
+        sql = f"""
+        SELECT {entity_cols}, 0 AS precedence
+        FROM main.entities
+        WHERE normalized_canonical = :normalized
+        {temporal_filter}
+        {type_filter}
+        {parent_filter}
+    """
+
+        # Add overlays if present
+        for schema_name, precedence in self.db.overlays:
+            sql += f"""
+            UNION ALL
+            SELECT {entity_cols}, {precedence} AS precedence
+            FROM {schema_name}.entities
+            WHERE normalized_canonical = :normalized
+            {temporal_filter}
+            {type_filter}
+            {parent_filter}
+        """
+
+        # Order by precedence to get highest precedence result
+        sql += "\nORDER BY precedence DESC LIMIT 1"
+
+        rows = self.db.execute(sql, params)
+
+        if not rows:
+            return None
+
+        row = rows[0]
+        return self._row_to_entity(row)
+
+    def _load_codes_for_entities(self, entities: list[Entity]) -> None:
+        """
+        Load codes for entities and populate their codes dict in place.
+
+        Args:
+            entities: List of entities to load codes for
+        """
+        if not entities:
+            return
+
+        # Get all DCIDs
+        dcids = [e.dcid for e in entities]
+
+        # Build placeholders for IN clause
+        placeholders = ", ".join(f":dcid_{i}" for i in range(len(dcids)))
+        params = {f"dcid_{i}": dcid for i, dcid in enumerate(dcids)}
+
+        # Query codes from main database
+        sql = f"""
+            SELECT entity_dcid, code_system, code_value
+            FROM main.codes
+            WHERE entity_dcid IN ({placeholders})
+        """
+
+        # Add overlays if present
+        for schema_name, _ in self.db.overlays:
+            sql += f"""
+                UNION ALL
+                SELECT entity_dcid, code_system, code_value
+                FROM {schema_name}.codes
+                WHERE entity_dcid IN ({placeholders})
+            """
+
+        result = self.db.execute(sql, params)
+
+        # Group codes by entity_dcid
+        codes_by_dcid: dict[str, dict[str, str]] = {}
+        for row in result:
+            dcid = row.entity_dcid
+            if dcid not in codes_by_dcid:
+                codes_by_dcid[dcid] = {}
+            codes_by_dcid[dcid][row.code_system] = row.code_value
+
+        # Populate codes in entities
+        for entity in entities:
+            if entity.dcid in codes_by_dcid:
+                entity.codes = codes_by_dcid[entity.dcid]