graflo 1.3.3__py3-none-any.whl

graflo/db/postgres/resource_mapping.py

@@ -0,0 +1,139 @@
+"""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 typing import Any
+
+from graflo.architecture.edge import EdgeConfig
+from graflo.architecture.resource import Resource
+from graflo.architecture.vertex import VertexConfig
+
+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: dict[str, Any],
+        vertex_config: VertexConfig,
+    ) -> 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
+
+        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"]
+
+        # 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"
+            )
+
+        # Create apply list with EdgeActor
+        # The actor wrapper will interpret {"source": source, "target": target} as EdgeActor
+        apply = [{"source": source_table, "target": target_table}]
+
+        resource = Resource(
+            resource_name=table_name,
+            apply=apply,
+        )
+
+        logger.debug(
+            f"Created edge resource '{table_name}' from {source_table} to {target_table}"
+        )
+
+        return resource
+
+    def map_tables_to_resources(
+        self,
+        introspection_result: dict[str, Any],
+        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 = []
+
+        # Map vertex tables to resources
+        vertex_tables = introspection_result.get("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.get("edge_tables", [])
+        for edge_table_info in edge_tables:
+            try:
+                resource = self.create_edge_resource(edge_table_info, vertex_config)
+                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

graflo/db/postgres/schema_inference.py

@@ -0,0 +1,245 @@
+"""Schema inference from PostgreSQL database introspection.
+
+This module provides functionality to infer graflo Schema objects from PostgreSQL
+3NF database schemas by analyzing table structures, relationships, and column types.
+"""
+
+import logging
+from typing import Any, Optional
+
+from graflo.architecture.edge import Edge, EdgeConfig, WeightConfig
+from graflo.architecture.onto import Index, IndexType
+from graflo.architecture.schema import Schema, SchemaMetadata
+from graflo.architecture.vertex import Field, Vertex, VertexConfig
+from graflo.onto import DBFlavor
+
+from .types import PostgresTypeMapper
+
+logger = logging.getLogger(__name__)
+
+
+class PostgresSchemaInferencer:
+    """Infers graflo Schema from PostgreSQL schema introspection results.
+
+    This class takes the output from PostgresConnection.introspect_schema() and
+    generates a complete graflo Schema with vertices, edges, and weights.
+    """
+
+    def __init__(self, db_flavor: DBFlavor = DBFlavor.ARANGO):
+        """Initialize the schema inferencer.
+
+        Args:
+            db_flavor: Target database flavor for the inferred schema
+        """
+        self.db_flavor = db_flavor
+        self.type_mapper = PostgresTypeMapper()
+
+    def infer_vertex_config(self, introspection_result: dict[str, Any]) -> VertexConfig:
+        """Infer VertexConfig from vertex tables.
+
+        Args:
+            introspection_result: Result from PostgresConnection.introspect_schema()
+
+        Returns:
+            VertexConfig: Inferred vertex configuration
+        """
+        vertex_tables = introspection_result.get("vertex_tables", [])
+        vertices = []
+
+        for table_info in vertex_tables:
+            table_name = table_info["name"]
+            columns = table_info["columns"]
+            pk_columns = table_info["primary_key"]
+
+            # Create fields from columns
+            fields = []
+            for col in columns:
+                field_name = col["name"]
+                field_type = self.type_mapper.map_type(col["type"])
+                fields.append(Field(name=field_name, type=field_type))
+
+            # Create indexes from primary key
+            indexes = []
+            if pk_columns:
+                indexes.append(
+                    Index(fields=pk_columns, type=IndexType.PERSISTENT, unique=True)
+                )
+
+            # Create vertex
+            vertex = Vertex(
+                name=table_name,
+                dbname=table_name,
+                fields=fields,
+                indexes=indexes,
+            )
+
+            vertices.append(vertex)
+            logger.debug(
+                f"Inferred vertex '{table_name}' with {len(fields)} fields and "
+                f"{len(indexes)} indexes"
+            )
+
+        return VertexConfig(vertices=vertices, db_flavor=self.db_flavor)
+
+    def infer_edge_weights(
+        self, edge_table_info: dict[str, Any]
+    ) -> Optional[WeightConfig]:
+        """Infer edge weights from edge table columns.
+
+        Args:
+            edge_table_info: Edge table information from introspection
+
+        Returns:
+            WeightConfig if there are weight columns, None otherwise
+        """
+        columns = edge_table_info["columns"]
+        pk_columns = set(edge_table_info["primary_key"])
+        fk_columns = {fk["column"] for fk in edge_table_info["foreign_keys"]}
+
+        # Find non-PK, non-FK columns (these become weights)
+        weight_columns = [
+            col
+            for col in columns
+            if col["name"] not in pk_columns and col["name"] not in fk_columns
+        ]
+
+        if not weight_columns:
+            return None
+
+        # Extract column names as direct weights
+        direct_weights = [col["name"] for col in weight_columns]
+
+        logger.debug(
+            f"Inferred {len(direct_weights)} weights for edge table "
+            f"'{edge_table_info['name']}': {direct_weights}"
+        )
+
+        return WeightConfig(direct=direct_weights)
+
+    def infer_edge_config(
+        self,
+        introspection_result: dict[str, Any],
+        vertex_config: VertexConfig,
+    ) -> EdgeConfig:
+        """Infer EdgeConfig from edge tables.
+
+        Args:
+            introspection_result: Result from PostgresConnection.introspect_schema()
+            vertex_config: Inferred vertex configuration
+
+        Returns:
+            EdgeConfig: Inferred edge configuration
+        """
+        edge_tables = introspection_result.get("edge_tables", [])
+        edges = []
+
+        vertex_names = vertex_config.vertex_set
+
+        for edge_table_info in edge_tables:
+            table_name = edge_table_info["name"]
+            source_table = edge_table_info["source_table"]
+            target_table = edge_table_info["target_table"]
+            fk_columns = edge_table_info["foreign_keys"]
+            pk_columns = edge_table_info["primary_key"]
+
+            # Verify source and target vertices exist
+            if source_table not in vertex_names:
+                logger.warning(
+                    f"Source vertex '{source_table}' for edge table '{table_name}' "
+                    f"not found in vertex config, skipping"
+                )
+                continue
+
+            if target_table not in vertex_names:
+                logger.warning(
+                    f"Target vertex '{target_table}' for edge table '{table_name}' "
+                    f"not found in vertex config, skipping"
+                )
+                continue
+
+            # Infer weights
+            weights = self.infer_edge_weights(edge_table_info)
+
+            # Create indexes from primary key and foreign keys
+            indexes = []
+            if pk_columns:
+                indexes.append(
+                    Index(fields=pk_columns, type=IndexType.PERSISTENT, unique=True)
+                )
+
+            # Add indexes for foreign keys (for efficient lookups)
+            # Note: Only add index if not already covered by primary key
+            pk_set = set(pk_columns)
+            for fk in fk_columns:
+                fk_column_name = fk["column"]
+                # Skip if FK column is part of primary key (already indexed)
+                if fk_column_name not in pk_set:
+                    indexes.append(
+                        Index(
+                            fields=[fk_column_name],
+                            type=IndexType.PERSISTENT,
+                            unique=False,
+                        )
+                    )
+
+            # Create edge
+            edge = Edge(
+                source=source_table,
+                target=target_table,
+                indexes=indexes,
+                weights=weights,
+                collection_name=table_name,
+            )
+
+            edges.append(edge)
+            logger.debug(
+                f"Inferred edge '{table_name}' from {source_table} to {target_table} "
+                f"with {len(indexes)} indexes"
+            )
+
+        return EdgeConfig(edges=edges)
+
+    def infer_schema(
+        self, introspection_result: dict[str, Any], schema_name: str | None = None
+    ) -> Schema:
+        """Infer complete Schema from PostgreSQL introspection.
+
+        Args:
+            introspection_result: Result from PostgresConnection.introspect_schema()
+            schema_name: Optional schema name (defaults to schema_name from introspection)
+
+        Returns:
+            Schema: Complete inferred schema with vertices, edges, and metadata
+        """
+        if schema_name is None:
+            schema_name = introspection_result.get("schema_name", "public")
+
+        logger.info(f"Inferring schema from PostgreSQL schema '{schema_name}'")
+
+        # Infer vertex configuration
+        vertex_config = self.infer_vertex_config(introspection_result)
+        logger.info(f"Inferred {len(vertex_config.vertices)} vertices")
+
+        # Infer edge configuration
+        edge_config = self.infer_edge_config(introspection_result, vertex_config)
+        edges_count = len(list(edge_config.edges_list()))
+        logger.info(f"Inferred {edges_count} edges")
+
+        # Create schema metadata
+        metadata = SchemaMetadata(name=schema_name)
+
+        # Create schema (resources will be added separately)
+        schema = Schema(
+            general=metadata,
+            vertex_config=vertex_config,
+            edge_config=edge_config,
+            resources=[],  # Resources will be created separately
+        )
+
+        logger.info(
+            f"Successfully inferred schema '{schema_name}' with "
+            f"{len(vertex_config.vertices)} vertices and "
+            f"{len(list(edge_config.edges_list()))} edges"
+        )
+
+        return schema

graflo/db/postgres/types.py

@@ -0,0 +1,148 @@
+"""Type mapping utilities for PostgreSQL to graflo type conversion.
+
+This module provides utilities for mapping PostgreSQL data types to graflo Field types,
+enabling automatic schema inference from PostgreSQL database schemas.
+"""
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class PostgresTypeMapper:
+    """Maps PostgreSQL data types to graflo Field types.
+
+    This class provides static methods for converting PostgreSQL type names
+    (from information_schema or pg_catalog) to graflo Field type strings.
+    """
+
+    # Mapping of PostgreSQL types to graflo Field types
+    TYPE_MAPPING = {
+        # Integer types
+        "integer": "INT",
+        "int": "INT",
+        "int4": "INT",
+        "smallint": "INT",
+        "int2": "INT",
+        "bigint": "INT",
+        "int8": "INT",
+        "serial": "INT",
+        "bigserial": "INT",
+        "smallserial": "INT",
+        # Floating point types
+        "real": "FLOAT",
+        "float4": "FLOAT",
+        "double precision": "FLOAT",
+        "float8": "FLOAT",
+        "numeric": "FLOAT",
+        "decimal": "FLOAT",
+        # Boolean
+        "boolean": "BOOL",
+        "bool": "BOOL",
+        # String types
+        "character varying": "STRING",
+        "varchar": "STRING",
+        "character": "STRING",
+        "char": "STRING",
+        "text": "STRING",
+        # Date/time types (mapped to DATETIME)
+        "timestamp": "DATETIME",
+        "timestamp without time zone": "DATETIME",
+        "timestamp with time zone": "DATETIME",
+        "timestamptz": "DATETIME",
+        "date": "DATETIME",
+        "time": "DATETIME",
+        "time without time zone": "DATETIME",
+        "time with time zone": "DATETIME",
+        "timetz": "DATETIME",
+        "interval": "STRING",  # Interval is duration, keep as STRING
+        # JSON types
+        "json": "STRING",
+        "jsonb": "STRING",
+        # Binary types
+        "bytea": "STRING",
+        # UUID
+        "uuid": "STRING",
+    }
+
+    @classmethod
+    def map_type(cls, postgres_type: str) -> str:
+        """Map PostgreSQL type to graflo Field type.
+
+        Args:
+            postgres_type: PostgreSQL type name (e.g., 'int4', 'varchar', 'timestamp')
+
+        Returns:
+            str: graflo Field type (INT, FLOAT, BOOL, or STRING)
+        """
+        # Normalize type name: lowercase and remove length specifications
+        normalized = postgres_type.lower().strip()
+
+        # Remove length specifications like (255) or (10,2)
+        if "(" in normalized:
+            normalized = normalized.split("(")[0].strip()
+
+        # Check direct mapping
+        if normalized in cls.TYPE_MAPPING:
+            return cls.TYPE_MAPPING[normalized]
+
+        # Check for partial matches (e.g., "character varying" might be stored as "varying")
+        for pg_type, graflo_type in cls.TYPE_MAPPING.items():
+            if pg_type in normalized or normalized in pg_type:
+                logger.debug(
+                    f"Mapped PostgreSQL type '{postgres_type}' to graflo type '{graflo_type}' "
+                    f"(partial match with '{pg_type}')"
+                )
+                return graflo_type
+
+        # Default to STRING for unknown types
+        logger.warning(
+            f"Unknown PostgreSQL type '{postgres_type}', defaulting to STRING"
+        )
+        return "STRING"
+
+    @classmethod
+    def is_datetime_type(cls, postgres_type: str) -> bool:
+        """Check if a PostgreSQL type is a datetime type.
+
+        Args:
+            postgres_type: PostgreSQL type name
+
+        Returns:
+            bool: True if the type is a datetime-related type
+        """
+        normalized = postgres_type.lower().strip()
+        datetime_types = [
+            "timestamp",
+            "date",
+            "time",
+            "interval",
+            "timestamptz",
+            "timetz",
+        ]
+        return any(dt_type in normalized for dt_type in datetime_types)
+
+    @classmethod
+    def is_numeric_type(cls, postgres_type: str) -> bool:
+        """Check if a PostgreSQL type is a numeric type.
+
+        Args:
+            postgres_type: PostgreSQL type name
+
+        Returns:
+            bool: True if the type is numeric
+        """
+        normalized = postgres_type.lower().strip()
+        numeric_types = [
+            "integer",
+            "int",
+            "bigint",
+            "smallint",
+            "serial",
+            "real",
+            "double precision",
+            "numeric",
+            "decimal",
+            "float",
+        ]
+        return any(nt_type in normalized for nt_type in numeric_types)

graflo/db/tigergraph/__init__.py

@@ -0,0 +1,9 @@
+"""TigerGraph database connection implementation.
+
+This package provides TigerGraph-specific database connection implementations
+and utilities for graph database operations.
+"""
+
+from .conn import TigerGraphConnection
+
+__all__ = ["TigerGraphConnection"]