graflo 1.3.7__py3-none-any.whl

graflo/db/postgres/inference_utils.py

@@ -0,0 +1,428 @@
+"""Inference utilities for PostgreSQL schema analysis.
+
+This module provides utility functions for inferring relationships and patterns
+from PostgreSQL table and column names using heuristics and fuzzy matching.
+"""
+
+from typing import Any
+
+from .fuzzy_matcher import FuzzyMatchCache, FuzzyMatcher
+
+
+def fuzzy_match_fragment(
+    fragment: str, vertex_names: list[str], threshold: float = 0.6
+) -> str | None:
+    """Fuzzy match a fragment to vertex names.
+
+    Backward-compatible wrapper function that uses the improved FuzzyMatcher.
+
+    Args:
+        fragment: Fragment to match
+        vertex_names: List of vertex table names to match against
+        threshold: Similarity threshold (0.0 to 1.0)
+
+    Returns:
+        Best matching vertex name or None if no match above threshold
+    """
+    matcher = FuzzyMatcher(vertex_names, threshold)
+    match, _ = matcher.match(fragment)
+    return match
+
+
+def detect_separator(text: str) -> str:
+    """Detect the most common separator character in a text.
+
+    Args:
+        text: Text to analyze
+
+    Returns:
+        Most common separator character, defaults to '_'
+    """
+    # Common separators
+    separators = ["_", "-", "."]
+    counts = {sep: text.count(sep) for sep in separators}
+
+    if max(counts.values()) > 0:
+        return max(counts, key=counts.get)
+    return "_"  # Default separator
+
+
+def split_by_separator(text: str, separator: str) -> list[str]:
+    """Split text by separator, handling multiple consecutive separators.
+
+    Args:
+        text: Text to split
+        separator: Separator character
+
+    Returns:
+        List of non-empty fragments
+    """
+    # Split and filter out empty strings
+    parts = [p for p in text.split(separator) if p]
+    return parts
+
+
+def infer_edge_vertices_from_table_name(
+    table_name: str,
+    pk_columns: list[str],
+    fk_columns: list[dict[str, Any]],
+    vertex_table_names: list[str] | None = None,
+    match_cache: FuzzyMatchCache | None = None,
+) -> tuple[str | None, str | None, str | None]:
+    """Infer source and target vertex names from table name and structure.
+
+    Uses fuzzy matching to identify vertex names in table name fragments and key names.
+    Handles patterns like:
+    - rel_cluster_containment_host -> cluster, host, containment
+    - rel_cluster_containment_cluster_2 -> cluster, cluster, containment (self-reference)
+    - user_follows_user -> user, user, follows (self-reference)
+    - product_category_mapping -> product, category, mapping
+
+    Args:
+        table_name: Name of the table
+        pk_columns: List of primary key column names
+        fk_columns: List of foreign key dictionaries with 'column' and 'references_table' keys
+        vertex_table_names: Optional list of known vertex table names for fuzzy matching
+        match_cache: Optional pre-computed fuzzy match cache for better performance
+
+    Returns:
+        Tuple of (source_table, target_table, relation_name) or (None, None, None) if cannot infer
+    """
+    if vertex_table_names is None:
+        vertex_table_names = []
+
+    # Use cache if provided, otherwise create a temporary one
+    if match_cache is None:
+        match_cache = FuzzyMatchCache(vertex_table_names)
+
+    # Step 1: Detect separator
+    separator = detect_separator(table_name)
+
+    # Step 2: Split table name by separator
+    table_fragments = split_by_separator(table_name, separator)
+
+    # Initialize relation_name - will be set if we identify a relation fragment
+    relation_name = None
+
+    # Step 3: Extract fragments from keys (preserve order for PK columns)
+    key_fragments_list = []  # Preserve order
+    key_fragments_set = set()  # For deduplication
+
+    # Extract fragments from PK columns in order
+    for pk_col in pk_columns:
+        pk_fragments = split_by_separator(pk_col, separator)
+        for frag in pk_fragments:
+            if frag not in key_fragments_set:
+                key_fragments_list.append(frag)
+                key_fragments_set.add(frag)
+
+    # Extract fragments from FK columns
+    for fk in fk_columns:
+        fk_col = fk.get("column", "")
+        fk_fragments = split_by_separator(fk_col, separator)
+        for frag in fk_fragments:
+            if frag not in key_fragments_set:
+                key_fragments_list.append(frag)
+                key_fragments_set.add(frag)
+
+    # Step 4: Match table name fragments to vertices
+    # Strategy: match source from left, target from right
+    # Stop when we have 2 matches OR target_index > source_index + 1
+    source_match_idx: int | None = None
+    target_match_idx: int | None = None
+    source_vertex: str | None = None
+    target_vertex: str | None = None
+    matched_vertices_set = set()  # For deduplication
+    matched_fragment_indices = {}  # Track which fragment indices matched which vertices
+
+    # Match source starting from the left
+    for i, fragment in enumerate(table_fragments):
+        matched = match_cache.get_match(fragment)
+        if matched and matched not in matched_vertices_set:
+            source_match_idx = i
+            source_vertex = matched
+            matched_vertices_set.add(matched)
+            matched_fragment_indices[i] = matched
+            break  # Found source, stop searching left
+
+    # Match target starting from the right
+    for i in range(
+        len(table_fragments) - 1,
+        source_match_idx if source_match_idx is not None else -1,
+        -1,
+    ):
+        fragment = table_fragments[i]
+        matched = match_cache.get_match(fragment)
+        if matched and matched not in matched_vertices_set:
+            target_match_idx = i
+            target_vertex = matched
+            matched_vertices_set.add(matched)
+            matched_fragment_indices[i] = matched
+            break  # Found target, stop searching right
+
+    # Match key fragments to fill in missing vertices (keys are primary source of truth)
+    matched_vertices = []
+    key_matched_vertices = []  # Track vertices matched from keys (higher priority)
+
+    if source_vertex:
+        matched_vertices.append(source_vertex)
+    if target_vertex and target_vertex != source_vertex:
+        matched_vertices.append(target_vertex)
+
+    for fragment in key_fragments_list:
+        matched = match_cache.get_match(fragment)
+        if matched:
+            if matched not in matched_vertices_set:
+                matched_vertices.append(matched)
+                matched_vertices_set.add(matched)
+            # Track key-matched vertices separately for priority
+            if matched not in key_matched_vertices:
+                key_matched_vertices.append(matched)
+
+    # Step 5: Use foreign keys to confirm or infer vertices
+    fk_vertex_names = []
+    if fk_columns:
+        for fk in fk_columns:
+            ref_table = fk.get("references_table")
+            if ref_table:
+                fk_vertex_names.append(ref_table)
+
+    # Step 6: Form hypothesis
+    source_table = None
+    target_table = None
+
+    # Priority 1: Use FK references if available (most reliable)
+    if len(fk_vertex_names) >= 2:
+        source_table = fk_vertex_names[0]
+        target_table = fk_vertex_names[1]
+    elif len(fk_vertex_names) == 1:
+        # Self-reference case
+        source_table = fk_vertex_names[0]
+        target_table = fk_vertex_names[0]
+
+    # Priority 2: Use matched vertices from fuzzy matching
+    # Prefer vertices from keys (primary source of truth) over table matches
+    if not source_table or not target_table:
+        # If we have vertices from keys, prefer those
+        if len(key_matched_vertices) >= 2:
+            source_table = key_matched_vertices[0]
+            target_table = key_matched_vertices[1]
+        elif len(key_matched_vertices) == 1:
+            # Use key vertex for source, try to find target from all matched vertices
+            source_table = key_matched_vertices[0]
+            if len(matched_vertices) >= 2:
+                # Find target that's not the source
+                for v in matched_vertices:
+                    if v != source_table:
+                        target_table = v
+                        break
+                if not target_table:
+                    target_table = source_table  # Self-reference
+            else:
+                target_table = source_table  # Self-reference
+        elif len(matched_vertices) >= 2:
+            source_table = matched_vertices[0]
+            target_table = matched_vertices[1]
+        elif len(matched_vertices) == 1:
+            # Self-reference case
+            source_table = matched_vertices[0]
+            target_table = matched_vertices[0]
+
+    # Priority 3: Fill in missing vertex from remaining options
+    if source_table and not target_table:
+        # Try to find target from remaining fragments or keys
+        if fk_vertex_names and len(fk_vertex_names) > 1:
+            # Use second FK if available
+            target_table = fk_vertex_names[1]
+        elif matched_vertices and len(matched_vertices) > 1:
+            target_table = matched_vertices[1]
+        elif fk_vertex_names:
+            # Self-reference case
+            target_table = fk_vertex_names[0]
+        elif matched_vertices:
+            target_table = matched_vertices[0]
+
+    if target_table and not source_table:
+        # Try to find source from remaining fragments or keys
+        if fk_vertex_names:
+            source_table = fk_vertex_names[0]
+        elif matched_vertices:
+            source_table = matched_vertices[0]
+
+    # Step 7: Identify relation from table fragments (after we know source/target)
+    # Relation is derived from table name fragments that are neither source nor target
+    # Patterns: bla_SOURCE_<relation>_TARGET, bla_SOURCE_TARGET_<relation>,
+    #           bla_<relation>_SOURCE_TARGET, and other combinations
+    # We allow relation to appear anywhere except as source/target fragments
+    if relation_name is None and source_table and target_table:
+        source_lower = source_table.lower()
+        target_lower = target_table.lower()
+
+        # Use the matched fragment indices if available (more precise)
+        # Otherwise, find by matching fragment text
+        source_idx = source_match_idx
+        target_idx = target_match_idx
+
+        # If we don't have indices from matching, find them by text matching
+        if source_idx is None or target_idx is None:
+            for idx, fragment in enumerate(table_fragments):
+                fragment_lower = fragment.lower()
+                if source_idx is None and (
+                    fragment_lower == source_lower
+                    or source_lower in fragment_lower
+                    or fragment_lower in source_lower
+                ):
+                    source_idx = idx
+                if target_idx is None and (
+                    fragment_lower == target_lower
+                    or target_lower in fragment_lower
+                    or fragment_lower in target_lower
+                ):
+                    target_idx = idx
+
+        relation_candidates = []
+
+        # Collect all fragments that are not source or target
+        # Allow relation to appear anywhere: before, between, or after source/target
+        for idx, fragment in enumerate(table_fragments):
+            fragment_lower = fragment.lower()
+
+            # Skip if it's a source or target fragment
+            if (
+                fragment_lower == source_lower
+                or source_lower in fragment_lower
+                or fragment_lower in source_lower
+                or fragment_lower == target_lower
+                or target_lower in fragment_lower
+                or fragment_lower in target_lower
+            ):
+                continue
+
+            # Include all non-source/target fragments as relation candidates
+            relation_candidates.append((len(fragment), idx, fragment))
+
+        # Select candidate using scoring system:
+        # - Score = fragment_length + (position_index * 5) if fragment_length >= 3
+        # - Score = fragment_length if fragment_length < 3
+        # - Prefer candidates further to the right and longer
+        if relation_candidates:
+
+            def score_candidate(candidate: tuple[int, int, str]) -> int:
+                fragment_length, position_idx, _ = candidate
+                if fragment_length >= 3:
+                    # Position bonus: each position to the right counts as 5 extra characters
+                    return fragment_length + (position_idx * 5)
+                else:
+                    # Fragments below 3 symbols don't get position bonus
+                    return fragment_length
+
+            _, _, relation_name = max(relation_candidates, key=score_candidate)
+        elif len(table_fragments) >= 2:
+            # Fallback: if we have 2+ fragments and one doesn't match source/target, it might be the relation
+            for fragment in table_fragments:
+                fragment_lower = fragment.lower()
+                # Use if it doesn't match source or target
+                if (
+                    fragment_lower != source_lower
+                    and source_lower not in fragment_lower
+                    and fragment_lower not in source_lower
+                    and fragment_lower != target_lower
+                    and target_lower not in fragment_lower
+                    and fragment_lower not in target_lower
+                ):
+                    relation_name = fragment
+                    break
+
+    return (source_table, target_table, relation_name)
+
+
+def infer_vertex_from_column_name(
+    column_name: str,
+    vertex_table_names: list[str] | None = None,
+    match_cache: FuzzyMatchCache | None = None,
+) -> str | None:
+    """Infer vertex table name from a column name using robust pattern matching.
+
+    Uses the same logic as infer_edge_vertices_from_table_name but focused on
+    extracting vertex names from column names. Handles patterns like:
+    - user_id -> user
+    - product_id -> product
+    - customer_fk -> customer
+    - source_vertex -> source_vertex (if matches)
+
+    Args:
+        column_name: Name of the column
+        vertex_table_names: Optional list of known vertex table names for fuzzy matching
+        match_cache: Optional pre-computed fuzzy match cache for better performance
+
+    Returns:
+        Inferred vertex table name or None if cannot infer
+    """
+    if vertex_table_names is None:
+        vertex_table_names = []
+
+    # Use cache if provided, otherwise create a temporary one
+    if match_cache is None:
+        match_cache = FuzzyMatchCache(vertex_table_names)
+
+    if not column_name:
+        return None
+
+    # Step 1: Detect separator
+    separator = detect_separator(column_name)
+
+    # Step 2: Split column name by separator
+    fragments = split_by_separator(column_name, separator)
+
+    if not fragments:
+        return None
+
+    # Step 3: Try to match fragments to vertex names
+    # Common suffixes to remove: id, fk, key, pk, ref
+    common_suffixes = {"id", "fk", "key", "pk", "ref", "reference"}
+
+    # Try matching full column name first
+    matched = match_cache.get_match(column_name)
+    if matched:
+        return matched
+
+    # Try matching fragments (excluding common suffixes)
+    for fragment in fragments:
+        fragment_lower = fragment.lower()
+        # Skip common suffixes
+        if fragment_lower in common_suffixes:
+            continue
+
+        matched = match_cache.get_match(fragment)
+        if matched:
+            return matched
+
+    # Step 4: If no match found, try removing common suffixes and matching again
+    # Remove last fragment if it's a common suffix
+    if len(fragments) > 1:
+        last_fragment = fragments[-1].lower()
+        if last_fragment in common_suffixes:
+            # Try matching the remaining fragments
+            remaining = separator.join(fragments[:-1])
+            matched = match_cache.get_match(remaining)
+            if matched:
+                return matched
+
+    # Step 5: As last resort, try exact match against vertex names (case-insensitive)
+    column_lower = column_name.lower()
+    for vertex_name in vertex_table_names:
+        vertex_lower = vertex_name.lower()
+        # Check if column name contains vertex name or vice versa
+        if vertex_lower in column_lower:
+            # Remove common suffixes from column name and check if it matches
+            for suffix in common_suffixes:
+                if column_lower.endswith(f"_{suffix}") or column_lower.endswith(suffix):
+                    base = (
+                        column_lower[: -len(f"_{suffix}")]
+                        if column_lower.endswith(f"_{suffix}")
+                        else column_lower[: -len(suffix)]
+                    )
+                    if base == vertex_lower:
+                        return vertex_name
+
+    return None

graflo/db/postgres/resource_mapping.py

@@ -0,0 +1,273 @@
+"""Resource mapping from PostgreSQL tables to graflo Resources.
+
+This module provides functionality to map PostgreSQL tables (both vertex and edge tables)
+to graflo Resource objects that can be used for data ingestion.
+"""
+
+import logging
+
+from graflo.architecture.edge import EdgeConfig
+from graflo.architecture.resource import Resource
+from graflo.architecture.vertex import VertexConfig
+
+from .conn import EdgeTableInfo, SchemaIntrospectionResult
+from .fuzzy_matcher import FuzzyMatchCache
+from .inference_utils import (
+    detect_separator,
+    split_by_separator,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class PostgresResourceMapper:
+    """Maps PostgreSQL tables to graflo Resources.
+
+    This class creates Resource objects that map PostgreSQL tables to graph vertices
+    and edges, enabling ingestion of relational data into graph databases.
+    """
+
+    def create_vertex_resource(self, table_name: str, vertex_name: str) -> Resource:
+        """Create a Resource for a vertex table.
+
+        Args:
+            table_name: Name of the PostgreSQL table
+            vertex_name: Name of the vertex type (typically same as table_name)
+
+        Returns:
+            Resource: Resource configured to ingest vertex data
+        """
+        # Create apply list with VertexActor
+        # The actor wrapper will interpret {"vertex": vertex_name} as VertexActor
+        apply = [{"vertex": vertex_name}]
+
+        resource = Resource(
+            resource_name=table_name,
+            apply=apply,
+        )
+
+        logger.debug(
+            f"Created vertex resource '{table_name}' for vertex '{vertex_name}'"
+        )
+
+        return resource
+
+    def create_edge_resource(
+        self,
+        edge_table_info: EdgeTableInfo,
+        vertex_config: VertexConfig,
+        match_cache: FuzzyMatchCache | None = None,
+    ) -> Resource:
+        """Create a Resource for an edge table.
+
+        Args:
+            edge_table_info: Edge table information from introspection
+            vertex_config: Vertex configuration for source/target validation
+            match_cache: Optional fuzzy match cache for better performance
+
+        Returns:
+            Resource: Resource configured to ingest edge data
+        """
+        table_name = edge_table_info.name
+        source_table = edge_table_info.source_table
+        target_table = edge_table_info.target_table
+        source_column = edge_table_info.source_column
+        target_column = edge_table_info.target_column
+        relation = edge_table_info.relation
+
+        # Verify source and target vertices exist
+        if source_table not in vertex_config.vertex_set:
+            raise ValueError(
+                f"Source vertex '{source_table}' for edge table '{table_name}' "
+                f"not found in vertex config"
+            )
+
+        if target_table not in vertex_config.vertex_set:
+            raise ValueError(
+                f"Target vertex '{target_table}' for edge table '{table_name}' "
+                f"not found in vertex config"
+            )
+
+        # Get primary key fields for source and target vertices
+        source_vertex_obj = vertex_config._vertices_map[source_table]
+        target_vertex_obj = vertex_config._vertices_map[target_table]
+
+        # Get the primary key field(s) from the first index (primary key)
+        source_pk_fields = (
+            source_vertex_obj.indexes[0].fields if source_vertex_obj.indexes else []
+        )
+        target_pk_fields = (
+            target_vertex_obj.indexes[0].fields if target_vertex_obj.indexes else []
+        )
+
+        # Use heuristics to infer PK field names from column names
+        # This handles cases like "bla_user" -> "user" vertex -> use "id" or matched field
+        vertex_names = list(vertex_config.vertex_set)
+        source_pk_field = self._infer_pk_field_from_column(
+            source_column, source_table, source_pk_fields, vertex_names, match_cache
+        )
+        target_pk_field = self._infer_pk_field_from_column(
+            target_column, target_table, target_pk_fields, vertex_names, match_cache
+        )
+
+        # Create apply list using source_vertex and target_vertex pattern
+        # This pattern explicitly specifies which vertex type each mapping targets,
+        # avoiding attribute collisions between different vertex types
+        apply = []
+
+        # First mapping: map source foreign key column to source vertex's primary key field
+        if source_column:
+            source_map_config = {
+                "target_vertex": source_table,
+                "map": {source_column: source_pk_field},
+            }
+            apply.append(source_map_config)
+
+        # Second mapping: map target foreign key column to target vertex's primary key field
+        if target_column:
+            target_map_config = {
+                "target_vertex": target_table,
+                "map": {target_column: target_pk_field},
+            }
+            apply.append(target_map_config)
+
+        resource = Resource(
+            resource_name=table_name,
+            apply=apply,
+        )
+
+        relation_info = f" with relation '{relation}'" if relation else ""
+        logger.debug(
+            f"Created edge resource '{table_name}' from {source_table} to {target_table}"
+            f"{relation_info} "
+            f"(source_col: {source_column} -> {source_pk_field}, "
+            f"target_col: {target_column} -> {target_pk_field})"
+        )
+
+        return resource
+
+    @staticmethod
+    def _infer_pk_field_from_column(
+        column_name: str,
+        vertex_name: str,
+        pk_fields: list[str],
+        vertex_names: list[str],
+        match_cache: FuzzyMatchCache | None = None,
+    ) -> str:
+        """Infer primary key field name from column name using heuristics.
+
+        Uses fuzzy matching to identify vertex name fragments in column names,
+        then matches to the appropriate PK field. Handles cases like:
+        - "user_id" -> "user" vertex -> use first PK field (e.g., "id")
+        - "bla_user" -> "user" vertex -> use first PK field
+        - "user_id_2" -> "user" vertex -> use first PK field
+        - "source_user_id" -> "user" vertex -> use first PK field
+        - "bla_user" and "bla_user_2" -> both map to "user" vertex PK field
+
+        The heuristic works by:
+        1. Splitting the column name into fragments
+        2. Fuzzy matching fragments to vertex names
+        3. If a fragment matches the target vertex_name, use the vertex's PK field
+        4. Otherwise, fall back to first PK field or "id"
+
+        Args:
+            column_name: Name of the column (e.g., "user_id", "bla_user", "bla_user_2")
+            vertex_name: Name of the target vertex (already known from edge table info)
+            pk_fields: List of primary key field names for the vertex
+            vertex_names: List of all vertex names for fuzzy matching
+            match_cache: Optional fuzzy match cache for better performance
+
+        Returns:
+            Primary key field name (defaults to first PK field or "id" if no match)
+        """
+        # Split column name into fragments
+        separator = detect_separator(column_name)
+        fragments = split_by_separator(column_name, separator)
+
+        # Try to find a fragment that matches the target vertex name
+        # This confirms that the column is indeed related to this vertex
+        for fragment in fragments:
+            # Fuzzy match fragment to vertex names
+            if match_cache:
+                matched_vertex = match_cache.get_match(fragment)
+            else:
+                # Fallback: create temporary matcher if cache not provided
+                from .fuzzy_matcher import FuzzyMatcher
+
+                matcher = FuzzyMatcher(vertex_names)
+                matched_vertex, _ = matcher.match(fragment)
+
+            # If we found a match to our target vertex, use its PK field
+            if matched_vertex == vertex_name:
+                if pk_fields:
+                    # Use the first PK field (most common case is single-column PK)
+                    return pk_fields[0]
+                else:
+                    # No PK fields available, use "id" as default
+                    return "id"
+
+        # No fragment matched the target vertex, but we still have vertex_name
+        # This might happen if the column name doesn't contain the vertex name fragment
+        # In this case, trust that vertex_name is correct and use its PK field
+        if pk_fields:
+            return pk_fields[0]
+
+        # Last resort: use "id" as default
+        # This is better than failing, but ideally pk_fields should always be available
+        logger.debug(
+            f"No PK fields found for vertex '{vertex_name}', using 'id' as default "
+            f"for column '{column_name}'"
+        )
+        return "id"
+
+    def map_tables_to_resources(
+        self,
+        introspection_result: SchemaIntrospectionResult,
+        vertex_config: VertexConfig,
+        edge_config: EdgeConfig,
+    ) -> list[Resource]:
+        """Map all PostgreSQL tables to Resources.
+
+        Creates Resources for both vertex and edge tables, enabling ingestion
+        of the entire database schema.
+
+        Args:
+            introspection_result: Result from PostgresConnection.introspect_schema()
+            vertex_config: Inferred vertex configuration
+            edge_config: Inferred edge configuration
+
+        Returns:
+            list[Resource]: List of Resources for all tables
+        """
+        resources = []
+
+        # Create fuzzy match cache once for all edge tables (significant performance improvement)
+        vertex_names = list(vertex_config.vertex_set)
+        match_cache = FuzzyMatchCache(vertex_names)
+
+        # Map vertex tables to resources
+        vertex_tables = introspection_result.vertex_tables
+        for table_info in vertex_tables:
+            table_name = table_info.name
+            vertex_name = table_name  # Use table name as vertex name
+            resource = self.create_vertex_resource(table_name, vertex_name)
+            resources.append(resource)
+
+        # Map edge tables to resources
+        edge_tables = introspection_result.edge_tables
+        for edge_table_info in edge_tables:
+            try:
+                resource = self.create_edge_resource(
+                    edge_table_info, vertex_config, match_cache
+                )
+                resources.append(resource)
+            except ValueError as e:
+                logger.warning(f"Skipping edge resource creation: {e}")
+                continue
+
+        logger.info(
+            f"Mapped {len(vertex_tables)} vertex tables and {len(edge_tables)} edge tables "
+            f"to {len(resources)} resources"
+        )
+
+        return resources