graflo.db.postgres

PostgreSQL database implementation.

This package provides PostgreSQL-specific implementations for schema introspection and connection management. It focuses on reading and analyzing 3NF schemas to identify vertex-like and edge-like tables, and inferring graflo Schema objects.

Key Components

• PostgresConnection: PostgreSQL connection and schema introspection implementation
• PostgresSchemaInferencer: Infers graflo Schema from PostgreSQL schemas
• PostgresResourceMapper: Maps PostgreSQL tables to graflo Resources

Example

from graflo.hq import GraphEngine from graflo.connections.onto import PostgresConfig config = PostgresConfig.from_docker_env() engine = GraphEngine() manifest = engine.infer_manifest(config, schema_name="public")

Modules:

Name	Description
conn	PostgreSQL connection implementation for schema introspection.
resource_mapping	Resource mapping from PostgreSQL tables to graflo Resources.
schema_inference	Schema inference from PostgreSQL database introspection.
target_write	PostgreSQL graph target write operations (DDL/DML for vertices and edge tables).
types	PostgreSQL type mapping.
util

Attributes

__all__ = ['PostgresConnection', 'PostgresResourceMapper', 'PostgresSchemaInferencer'] module-attribute

Classes

PostgresConnection

Bases: PostgresTargetWriteMixin, Connection

PostgreSQL connection for schema introspection and graph target writes.

This class provides PostgreSQL-specific functionality for connecting to databases and introspecting 3NF schemas to identify vertex-like and edge-like tables. It also supports writing graph data to vertex and edge junction tables.

Attributes:

Name	Type	Description
config		PostgreSQL connection configuration
conn		psycopg2 connection instance

Source code in graflo/db/postgres/conn.py

class PostgresConnection(PostgresTargetWriteMixin, Connection):
    """PostgreSQL connection for schema introspection and graph target writes.

    This class provides PostgreSQL-specific functionality for connecting to databases
    and introspecting 3NF schemas to identify vertex-like and edge-like tables.
    It also supports writing graph data to vertex and edge junction tables.

    Attributes:
        config: PostgreSQL connection configuration
        conn: psycopg2 connection instance
    """

    def __init__(self, config: PostgresConfig):
        """Initialize PostgreSQL connection.

        Args:
            config: PostgreSQL connection configuration containing URI and credentials
        """
        super().__init__()
        self.config = config

        # Validate required config values
        if config.uri is None:
            raise ValueError("PostgreSQL connection requires a URI to be configured")
        if config.database is None:
            raise ValueError(
                "PostgreSQL connection requires a database name to be configured"
            )

        # Use config properties directly - all fallbacks are handled in PostgresConfig
        host = config.hostname or "localhost"
        port = int(config.port) if config.port else 5432
        database = config.database
        user = config.username or "postgres"
        password = config.password

        # Build connection parameters dict
        conn_params = {
            "host": host,
            "port": port,
            "database": database,
            "user": user,
        }

        if password:
            conn_params["password"] = password

        try:
            self.conn = psycopg2.connect(**conn_params)
            logger.info(f"Successfully connected to PostgreSQL database '{database}'")
        except Exception as e:
            logger.error(f"Failed to connect to PostgreSQL: {e}", exc_info=True)
            raise

    def read(
        self, query: str, params: tuple | dict[str, Any] | None = None
    ) -> list[dict[str, Any]]:
        """Execute a SELECT query and return results as a list of dictionaries.

        Args:
            query: SQL SELECT query to execute
            params: Optional parameters — a tuple for ``%s`` placeholders, or a
                mapping for named ``%(name)s`` ones.

        Returns:
            List of dictionaries, where each dictionary represents a row with column names as keys.
            Decimal values are converted to float for compatibility with graph databases.
        """
        from decimal import Decimal

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            if params:
                cursor.execute(query, params)
            else:
                cursor.execute(query)

            # Convert rows to dictionaries and convert Decimal to float
            results = []
            for row in cursor.fetchall():
                row_dict = dict(row)
                # Convert Decimal to float for JSON/graph database compatibility
                for key, value in row_dict.items():
                    if isinstance(value, Decimal):
                        row_dict[key] = float(value)
                results.append(row_dict)

            return results

    def __enter__(self):
        """Enter the context manager.

        Returns:
            PostgresConnection: Self for use in 'with' statements
        """
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        """Exit the context manager.

        Ensures the connection is properly closed when exiting the context.

        Args:
            exc_type: Exception type if an exception occurred
            exc_value: Exception value if an exception occurred
            exc_traceback: Exception traceback if an exception occurred
        """
        self.close()
        return False  # Don't suppress exceptions

    def close(self):
        """Close the PostgreSQL connection."""
        if hasattr(self, "conn") and self.conn:
            try:
                self.conn.close()
                logger.debug("PostgreSQL connection closed")
            except Exception as e:
                logger.warning(
                    f"Error closing PostgreSQL connection: {e}", exc_info=True
                )

    def _check_information_schema_reliable(self, schema_name: str) -> bool:
        """Check if information_schema is reliable for the given schema.

        Args:
            schema_name: Schema name to check

        Returns:
            True if information_schema appears reliable, False otherwise
        """
        try:
            # Try to query information_schema.tables
            query = """
                SELECT COUNT(*) as count
                FROM information_schema.tables
                WHERE table_schema = %s
                  AND table_type = 'BASE TABLE'
            """
            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query, (schema_name,))
                result = cursor.fetchone()
                # If query succeeds, check if we can also query constraints
                pk_query = """
                    SELECT COUNT(*) as count
                    FROM information_schema.table_constraints tc
                    JOIN information_schema.key_column_usage kcu
                        ON tc.constraint_name = kcu.constraint_name
                        AND tc.table_schema = kcu.table_schema
                    WHERE tc.constraint_type = 'PRIMARY KEY'
                      AND tc.table_schema = %s
                """
                cursor.execute(pk_query, (schema_name,))
                pk_result = cursor.fetchone()
                # If both queries work, information_schema seems reliable
                return result is not None and pk_result is not None
        except Exception as e:
            logger.debug(f"information_schema check failed: {e}")
            return False

    def _get_tables_pg_catalog(self, schema_name: str) -> list[dict[str, Any]]:
        """Get all tables using pg_catalog (fallback method).

        Args:
            schema_name: Schema name to query

        Returns:
            List of table information dictionaries with keys: table_name,
            table_schema, description (the table comment, empty when none)
        """
        query = """
            SELECT
                c.relname as table_name,
                n.nspname as table_schema,
                COALESCE(obj_description(c.oid, 'pg_class'), '') as description
            FROM pg_catalog.pg_class c
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            WHERE n.nspname = %s
              AND c.relkind = 'r'
              AND NOT c.relispartition
            ORDER BY c.relname;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name,))
            return [dict(row) for row in cursor.fetchall()]

    def get_tables(self, schema_name: str | None = None) -> list[dict[str, Any]]:
        """Get all tables in the specified schema.

        Tries information_schema first, falls back to pg_catalog if needed.

        Args:
            schema_name: Schema name to query. If None, uses 'public' or config schema_name.

        Returns:
            List of table information dictionaries with keys: table_name,
            table_schema, description (the table comment, empty when none)
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        # Try information_schema first
        try:
            query = """
                SELECT
                    table_name,
                    table_schema,
                    COALESCE(
                        obj_description(
                            (quote_ident(table_schema) || '.' || quote_ident(table_name))::regclass,
                            'pg_class'
                        ),
                        ''
                    ) AS description
                FROM information_schema.tables
                WHERE table_schema = %s
                  AND table_type = 'BASE TABLE'
                ORDER BY table_name;
            """

            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query, (schema_name,))
                results = [dict(row) for row in cursor.fetchall()]
                # If we got results, check if information_schema is reliable
                if results and self._check_information_schema_reliable(schema_name):
                    return results
                # If no results or unreliable, fall back to pg_catalog
                logger.debug(
                    f"information_schema returned no results or is unreliable, "
                    f"falling back to pg_catalog for schema '{schema_name}'"
                )
        except Exception as e:
            logger.debug(
                f"information_schema query failed: {e}, falling back to pg_catalog"
            )

        # Fallback to pg_catalog
        return self._get_tables_pg_catalog(schema_name)

    def _get_table_columns_pg_catalog(
        self, table_name: str, schema_name: str
    ) -> list[dict[str, Any]]:
        """Get columns using pg_catalog (fallback method).

        Args:
            table_name: Name of the table
            schema_name: Schema name

        Returns:
            List of column information dictionaries with keys:
            name, type, description, is_nullable, column_default
        """
        query = """
            SELECT
                a.attname as name,
                pg_catalog.format_type(a.atttypid, a.atttypmod) as type,
                CASE WHEN a.attnotnull THEN 'NO' ELSE 'YES' END as is_nullable,
                pg_catalog.pg_get_expr(d.adbin, d.adrelid) as column_default,
                COALESCE(dsc.description, '') as description,
                a.attnum as ordinal_position
            FROM pg_catalog.pg_attribute a
            JOIN pg_catalog.pg_class c ON c.oid = a.attrelid
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            LEFT JOIN pg_catalog.pg_attrdef d ON d.adrelid = a.attrelid AND d.adnum = a.attnum
            LEFT JOIN pg_catalog.pg_description dsc ON dsc.objoid = a.attrelid AND dsc.objsubid = a.attnum
            WHERE n.nspname = %s
              AND c.relname = %s
              AND a.attnum > 0
              AND NOT a.attisdropped
            ORDER BY a.attnum;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name, table_name))
            columns = []
            for row in cursor.fetchall():
                col_dict = dict(row)
                # Normalize type format
                if col_dict["type"]:
                    # Remove length info from type if present (e.g., "character varying(255)" -> "varchar")
                    type_str = col_dict["type"]
                    if "(" in type_str:
                        base_type = type_str.split("(")[0]
                        # Map common types
                        type_mapping = {
                            "character varying": "varchar",
                            "character": "char",
                            "double precision": "float8",
                            "real": "float4",
                            "integer": "int4",
                            "bigint": "int8",
                            "smallint": "int2",
                        }
                        col_dict["type"] = type_mapping.get(
                            base_type.lower(), base_type.lower()
                        )
                    else:
                        type_mapping = {
                            "character varying": "varchar",
                            "character": "char",
                            "double precision": "float8",
                            "real": "float4",
                            "integer": "int4",
                            "bigint": "int8",
                            "smallint": "int2",
                        }
                        col_dict["type"] = type_mapping.get(
                            type_str.lower(), type_str.lower()
                        )
                columns.append(col_dict)
            return columns

    def get_table_columns(
        self, table_name: str, schema_name: str | None = None
    ) -> list[dict[str, Any]]:
        """Get columns for a specific table with types and descriptions.

        Tries information_schema first, falls back to pg_catalog if needed.

        Args:
            table_name: Name of the table
            schema_name: Schema name. If None, uses 'public' or config schema_name.

        Returns:
            List of column information dictionaries with keys:
            name, type, description, is_nullable, column_default
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        # Try information_schema first
        try:
            query = """
                SELECT
                    c.column_name as name,
                    c.data_type as type,
                    c.udt_name as udt_name,
                    c.character_maximum_length,
                    c.is_nullable,
                    c.column_default,
                    COALESCE(d.description, '') as description,
                    c.ordinal_position as ordinal_position
                FROM information_schema.columns c
                LEFT JOIN pg_catalog.pg_statio_all_tables st
                    ON st.schemaname = c.table_schema
                    AND st.relname = c.table_name
                LEFT JOIN pg_catalog.pg_description d
                    ON d.objoid = st.relid
                    AND d.objsubid = c.ordinal_position
                WHERE c.table_schema = %s
                  AND c.table_name = %s
                ORDER BY c.ordinal_position;
            """

            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query, (schema_name, table_name))
                columns = []
                for row in cursor.fetchall():
                    col_dict = dict(row)
                    # Format type with length if applicable
                    if col_dict["character_maximum_length"]:
                        col_dict["type"] = (
                            f"{col_dict['type']}({col_dict['character_maximum_length']})"
                        )
                    # Use udt_name if it's more specific (e.g., varchar, int4)
                    if (
                        col_dict["udt_name"]
                        and col_dict["udt_name"] != col_dict["type"]
                    ):
                        col_dict["type"] = col_dict["udt_name"]
                    # Remove helper fields
                    col_dict.pop("character_maximum_length", None)
                    col_dict.pop("udt_name", None)
                    columns.append(col_dict)

                # If we got results and information_schema is reliable, return them
                if columns and self._check_information_schema_reliable(schema_name):
                    return columns
                # Otherwise fall back to pg_catalog
                logger.debug(
                    f"information_schema returned no results or is unreliable, "
                    f"falling back to pg_catalog for table '{schema_name}.{table_name}'"
                )
        except Exception as e:
            logger.debug(
                f"information_schema query failed: {e}, falling back to pg_catalog"
            )

        # Fallback to pg_catalog
        return self._get_table_columns_pg_catalog(table_name, schema_name)

    def _get_primary_keys_pg_catalog(
        self, table_name: str, schema_name: str
    ) -> list[str]:
        """Get primary key columns using pg_catalog (fallback method).

        Args:
            table_name: Name of the table
            schema_name: Schema name

        Returns:
            List of primary key column names
        """
        query = """
            SELECT a.attname
            FROM pg_catalog.pg_constraint con
            JOIN pg_catalog.pg_class c ON c.oid = con.conrelid
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            JOIN pg_catalog.pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = ANY(con.conkey)
            WHERE n.nspname = %s
              AND c.relname = %s
              AND con.contype = 'p'
            ORDER BY array_position(con.conkey, a.attnum);
        """

        with self.conn.cursor() as cursor:
            cursor.execute(query, (schema_name, table_name))
            return [row[0] for row in cursor.fetchall()]

    def get_primary_keys(
        self, table_name: str, schema_name: str | None = None
    ) -> list[str]:
        """Get primary key columns for a table.

        Tries information_schema first, falls back to pg_catalog if needed.

        Args:
            table_name: Name of the table
            schema_name: Schema name. If None, uses 'public' or config schema_name.

        Returns:
            List of primary key column names
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        # Try information_schema first
        try:
            query = """
                SELECT kcu.column_name
                FROM information_schema.table_constraints tc
                JOIN information_schema.key_column_usage kcu
                    ON tc.constraint_name = kcu.constraint_name
                    AND tc.table_schema = kcu.table_schema
                WHERE tc.constraint_type = 'PRIMARY KEY'
                  AND tc.table_schema = %s
                  AND tc.table_name = %s
                ORDER BY kcu.ordinal_position;
            """

            with self.conn.cursor() as cursor:
                cursor.execute(query, (schema_name, table_name))
                results = [row[0] for row in cursor.fetchall()]
                # If we got results and information_schema is reliable, return them
                if results and self._check_information_schema_reliable(schema_name):
                    return results
                # Otherwise fall back to pg_catalog
                logger.debug(
                    f"information_schema returned no results or is unreliable, "
                    f"falling back to pg_catalog for primary keys of '{schema_name}.{table_name}'"
                )
        except Exception as e:
            logger.debug(
                f"information_schema query failed: {e}, falling back to pg_catalog"
            )

        # Fallback to pg_catalog
        return self._get_primary_keys_pg_catalog(table_name, schema_name)

    def _get_unique_columns_pg_catalog(
        self, table_name: str, schema_name: str
    ) -> list[str]:
        """Get columns in UNIQUE constraints using pg_catalog (fallback method)."""
        query = """
            SELECT a.attname
            FROM pg_catalog.pg_constraint con
            JOIN pg_catalog.pg_class c ON c.oid = con.conrelid
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            JOIN pg_catalog.pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = ANY(con.conkey)
            WHERE n.nspname = %s
              AND c.relname = %s
              AND con.contype = 'u'
              AND a.attnum > 0
              AND NOT a.attisdropped
            ORDER BY array_position(con.conkey, a.attnum);
        """
        with self.conn.cursor() as cursor:
            cursor.execute(query, (schema_name, table_name))
            return list(dict.fromkeys(row[0] for row in cursor.fetchall()))

    def get_unique_columns(
        self, table_name: str, schema_name: str | None = None
    ) -> list[str]:
        """Get column names that participate in any UNIQUE constraint.

        Tries information_schema first, falls back to pg_catalog if needed.
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        try:
            query = """
                SELECT kcu.column_name
                FROM information_schema.table_constraints tc
                JOIN information_schema.key_column_usage kcu
                    ON tc.constraint_name = kcu.constraint_name
                    AND tc.table_schema = kcu.table_schema
                WHERE tc.constraint_type = 'UNIQUE'
                  AND tc.table_schema = %s
                  AND tc.table_name = %s
                ORDER BY kcu.ordinal_position;
            """
            with self.conn.cursor() as cursor:
                cursor.execute(query, (schema_name, table_name))
                results = list(dict.fromkeys(row[0] for row in cursor.fetchall()))
                if results and self._check_information_schema_reliable(schema_name):
                    return results
        except Exception as e:
            logger.debug(
                f"information_schema query failed for unique columns: {e}, "
                "falling back to pg_catalog"
            )
        return self._get_unique_columns_pg_catalog(table_name, schema_name)

    def _get_foreign_keys_pg_catalog(
        self, table_name: str, schema_name: str
    ) -> list[dict[str, Any]]:
        """Get foreign key relationships using pg_catalog (fallback method).

        Handles both single-column and multi-column foreign keys.
        For multi-column foreign keys, returns one row per column.

        Args:
            table_name: Name of the table
            schema_name: Schema name

        Returns:
            List of foreign key dictionaries with keys:
            column, references_table, references_column, constraint_name
        """
        # Use generate_subscripts for better compatibility with older PostgreSQL versions
        query = """
            SELECT
                a.attname as column,
                ref_c.relname as references_table,
                ref_a.attname as references_column,
                con.conname as constraint_name
            FROM pg_catalog.pg_constraint con
            JOIN pg_catalog.pg_class c ON c.oid = con.conrelid
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            JOIN pg_catalog.pg_class ref_c ON ref_c.oid = con.confrelid
            JOIN generate_subscripts(con.conkey, 1) AS i ON true
            JOIN pg_catalog.pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = con.conkey[i]
            JOIN pg_catalog.pg_attribute ref_a ON ref_a.attrelid = con.confrelid AND ref_a.attnum = con.confkey[i]
            WHERE n.nspname = %s
              AND c.relname = %s
              AND con.contype = 'f'
            ORDER BY con.conname, i;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name, table_name))
            return [dict(row) for row in cursor.fetchall()]

    def get_foreign_keys(
        self, table_name: str, schema_name: str | None = None
    ) -> list[dict[str, Any]]:
        """Get foreign key relationships for a table.

        Tries information_schema first, falls back to pg_catalog if needed.

        Args:
            table_name: Name of the table
            schema_name: Schema name. If None, uses 'public' or config schema_name.

        Returns:
            List of foreign key dictionaries with keys:
            column, references_table, references_column, constraint_name
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"

        # Try information_schema first
        try:
            query = """
                SELECT
                    kcu.column_name as column,
                    ccu.table_name as references_table,
                    ccu.column_name as references_column,
                    tc.constraint_name
                FROM information_schema.table_constraints tc
                JOIN information_schema.key_column_usage kcu
                    ON tc.constraint_name = kcu.constraint_name
                    AND tc.table_schema = kcu.table_schema
                JOIN information_schema.constraint_column_usage ccu
                    ON ccu.constraint_name = tc.constraint_name
                    AND ccu.table_schema = tc.table_schema
                WHERE tc.constraint_type = 'FOREIGN KEY'
                  AND tc.table_schema = %s
                  AND tc.table_name = %s
                ORDER BY kcu.ordinal_position;
            """

            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query, (schema_name, table_name))
                results = [dict(row) for row in cursor.fetchall()]
                # If we got results and information_schema is reliable, return them
                if self._check_information_schema_reliable(schema_name):
                    return results
                # Otherwise fall back to pg_catalog
                logger.debug(
                    f"information_schema returned no results or is unreliable, "
                    f"falling back to pg_catalog for foreign keys of '{schema_name}.{table_name}'"
                )
        except Exception as e:
            logger.debug(
                f"information_schema query failed: {e}, falling back to pg_catalog"
            )

        # Fallback to pg_catalog
        return self._get_foreign_keys_pg_catalog(table_name, schema_name)

    def get_table_row_count_estimate(
        self, table_name: str, schema_name: str | None = None
    ) -> int | None:
        """Return approximate row count from pg_class.reltuples (updated by ANALYZE).

        Avoids full table scan; may be stale until next ANALYZE.
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"
        query = """
            SELECT c.reltuples::bigint AS estimate
            FROM pg_catalog.pg_class c
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            WHERE n.nspname = %s AND c.relname = %s AND c.relkind = 'r';
        """
        with self.conn.cursor() as cursor:
            cursor.execute(query, (schema_name, table_name))
            row = cursor.fetchone()
            if row is None:
                return None
            val = row[0]
            return int(val) if val is not None else None

    def get_table_sample_rows(
        self,
        table_name: str,
        schema_name: str | None = None,
        limit: int = 5,
    ) -> list[dict[str, Any]]:
        """Return `limit` rows from the table, in a stable order.

        Ordered by the primary key, or by the whole row's text when the table
        has none, so sampling unchanged data returns the same rows every time.
        Without an order PostgreSQL returns rows in storage order, which an
        update or ``VACUUM FULL`` changes -- and a sample that feeds a prompt
        then changes the prompt, and with it any response cache keyed on it.
        Ordering a table without a key sorts it in full, which is the price of
        a repeatable sample there.
        """
        if schema_name is None:
            schema_name = self.config.schema_name or "public"
        primary_key = self.get_primary_keys(table_name, schema_name)
        order_by = (
            sql.SQL(", ").join(sql.Identifier("t", column) for column in primary_key)
            if primary_key
            else sql.SQL("t::text")
        )
        query = sql.SQL("SELECT * FROM {}.{} AS t ORDER BY {} LIMIT %s").format(
            sql.Identifier(schema_name),
            sql.Identifier(table_name),
            order_by,
        )
        try:
            with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
                cursor.execute(query, (limit,))
                return [dict(row) for row in cursor.fetchall()]
        except Exception as e:
            logger.debug(
                f"Could not fetch sample rows for '{schema_name}.{table_name}': {e}"
            )
            return []

    # ------------------------------------------------------------------
    # Schema introspection
    #
    # The classification below is dialect-neutral and lives in graflo.db.sql;
    # this class contributes the `pg_catalog` implementation of the metadata
    # primitives it reads. These wrappers stay because they are public API and
    # because `self` satisfies SqlMetadataProvider structurally.
    # ------------------------------------------------------------------

    def _default_schema(self) -> str:
        return self.config.schema_name or "public"

    def _is_edge_like_table(
        self, table_name: str, pk_columns: list[str], fk_columns: list[dict[str, Any]]
    ) -> bool:
        """Whether *table_name* looks like a relationship rather than an entity."""
        return sql_introspect.is_edge_like_table(table_name, pk_columns, fk_columns)

    def detect_vertex_tables(
        self, schema_name: str | None = None
    ) -> list[VertexTableInfo]:
        """Entity-like tables in the schema."""
        return sql_introspect.detect_vertex_tables(
            self, schema_name, default_schema=self._default_schema()
        )

    def detect_edge_tables(
        self,
        schema_name: str | None = None,
        vertex_table_names: list[str] | None = None,
    ) -> list[EdgeTableInfo]:
        """Relationship-like tables, with their recovered endpoints."""
        return sql_introspect.detect_edge_tables(
            self,
            schema_name,
            vertex_table_names,
            default_schema=self._default_schema(),
        )

    def _build_raw_tables(self, schema_name: str) -> list[RawTableInfo]:
        """Unclassified metadata for every table, including sampled values."""
        return sql_introspect.build_raw_tables(self, schema_name)

    def introspect_schema(
        self,
        schema_name: str | None = None,
        include_raw_tables: bool = False,
        entity_tables: Collection[str] | None = None,
    ) -> SchemaIntrospectionResult:
        """Classify the schema into vertex-like and edge-like tables."""
        return sql_introspect.introspect_schema(
            self,
            schema_name,
            include_raw_tables,
            default_schema=self._default_schema(),
            entity_tables=entity_tables,
        )

Attributes

config = config instance-attribute

conn = psycopg2.connect(**conn_params) instance-attribute

Methods:

__enter__()

Enter the context manager.

Returns:

Name	Type	Description
PostgresConnection		Self for use in 'with' statements

Source code in graflo/db/postgres/conn.py

def __enter__(self):
    """Enter the context manager.

    Returns:
        PostgresConnection: Self for use in 'with' statements
    """
    return self

__exit__(exc_type, exc_value, exc_traceback)

Exit the context manager.

Ensures the connection is properly closed when exiting the context.

Parameters:

Name	Type	Description	Default
exc_type		Exception type if an exception occurred	required
exc_value		Exception value if an exception occurred	required
exc_traceback		Exception traceback if an exception occurred	required

Source code in graflo/db/postgres/conn.py

def __exit__(self, exc_type, exc_value, exc_traceback):
    """Exit the context manager.

    Ensures the connection is properly closed when exiting the context.

    Args:
        exc_type: Exception type if an exception occurred
        exc_value: Exception value if an exception occurred
        exc_traceback: Exception traceback if an exception occurred
    """
    self.close()
    return False  # Don't suppress exceptions

__init__(config)

Initialize PostgreSQL connection.

Parameters:

Name	Type	Description	Default
config	PostgresConfig	PostgreSQL connection configuration containing URI and credentials	required

Source code in graflo/db/postgres/conn.py

def __init__(self, config: PostgresConfig):
    """Initialize PostgreSQL connection.

    Args:
        config: PostgreSQL connection configuration containing URI and credentials
    """
    super().__init__()
    self.config = config

    # Validate required config values
    if config.uri is None:
        raise ValueError("PostgreSQL connection requires a URI to be configured")
    if config.database is None:
        raise ValueError(
            "PostgreSQL connection requires a database name to be configured"
        )

    # Use config properties directly - all fallbacks are handled in PostgresConfig
    host = config.hostname or "localhost"
    port = int(config.port) if config.port else 5432
    database = config.database
    user = config.username or "postgres"
    password = config.password

    # Build connection parameters dict
    conn_params = {
        "host": host,
        "port": port,
        "database": database,
        "user": user,
    }

    if password:
        conn_params["password"] = password

    try:
        self.conn = psycopg2.connect(**conn_params)
        logger.info(f"Successfully connected to PostgreSQL database '{database}'")
    except Exception as e:
        logger.error(f"Failed to connect to PostgreSQL: {e}", exc_info=True)
        raise

close()

Close the PostgreSQL connection.

Source code in graflo/db/postgres/conn.py

def close(self):
    """Close the PostgreSQL connection."""
    if hasattr(self, "conn") and self.conn:
        try:
            self.conn.close()
            logger.debug("PostgreSQL connection closed")
        except Exception as e:
            logger.warning(
                f"Error closing PostgreSQL connection: {e}", exc_info=True
            )

detect_edge_tables(schema_name=None, vertex_table_names=None)

Relationship-like tables, with their recovered endpoints.

Source code in graflo/db/postgres/conn.py

def detect_edge_tables(
    self,
    schema_name: str | None = None,
    vertex_table_names: list[str] | None = None,
) -> list[EdgeTableInfo]:
    """Relationship-like tables, with their recovered endpoints."""
    return sql_introspect.detect_edge_tables(
        self,
        schema_name,
        vertex_table_names,
        default_schema=self._default_schema(),
    )

detect_vertex_tables(schema_name=None)

Entity-like tables in the schema.

Source code in graflo/db/postgres/conn.py

def detect_vertex_tables(
    self, schema_name: str | None = None
) -> list[VertexTableInfo]:
    """Entity-like tables in the schema."""
    return sql_introspect.detect_vertex_tables(
        self, schema_name, default_schema=self._default_schema()
    )

get_foreign_keys(table_name, schema_name=None)

Get foreign key relationships for a table.

Tries information_schema first, falls back to pg_catalog if needed.

Parameters:

Name	Type	Description	Default
table_name	str	Name of the table	required
schema_name	str | None	Schema name. If None, uses 'public' or config schema_name.	None

Returns:

Type	Description
list[dict[str, Any]]	List of foreign key dictionaries with keys:
list[dict[str, Any]]	column, references_table, references_column, constraint_name

Source code in graflo/db/postgres/conn.py

def get_foreign_keys(
    self, table_name: str, schema_name: str | None = None
) -> list[dict[str, Any]]:
    """Get foreign key relationships for a table.

    Tries information_schema first, falls back to pg_catalog if needed.

    Args:
        table_name: Name of the table
        schema_name: Schema name. If None, uses 'public' or config schema_name.

    Returns:
        List of foreign key dictionaries with keys:
        column, references_table, references_column, constraint_name
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    # Try information_schema first
    try:
        query = """
            SELECT
                kcu.column_name as column,
                ccu.table_name as references_table,
                ccu.column_name as references_column,
                tc.constraint_name
            FROM information_schema.table_constraints tc
            JOIN information_schema.key_column_usage kcu
                ON tc.constraint_name = kcu.constraint_name
                AND tc.table_schema = kcu.table_schema
            JOIN information_schema.constraint_column_usage ccu
                ON ccu.constraint_name = tc.constraint_name
                AND ccu.table_schema = tc.table_schema
            WHERE tc.constraint_type = 'FOREIGN KEY'
              AND tc.table_schema = %s
              AND tc.table_name = %s
            ORDER BY kcu.ordinal_position;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name, table_name))
            results = [dict(row) for row in cursor.fetchall()]
            # If we got results and information_schema is reliable, return them
            if self._check_information_schema_reliable(schema_name):
                return results
            # Otherwise fall back to pg_catalog
            logger.debug(
                f"information_schema returned no results or is unreliable, "
                f"falling back to pg_catalog for foreign keys of '{schema_name}.{table_name}'"
            )
    except Exception as e:
        logger.debug(
            f"information_schema query failed: {e}, falling back to pg_catalog"
        )

    # Fallback to pg_catalog
    return self._get_foreign_keys_pg_catalog(table_name, schema_name)

get_primary_keys(table_name, schema_name=None)

Get primary key columns for a table.

Tries information_schema first, falls back to pg_catalog if needed.

Parameters:

Name	Type	Description	Default
table_name	str	Name of the table	required
schema_name	str | None	Schema name. If None, uses 'public' or config schema_name.	None

Returns:

Type	Description
list[str]	List of primary key column names

Source code in graflo/db/postgres/conn.py

def get_primary_keys(
    self, table_name: str, schema_name: str | None = None
) -> list[str]:
    """Get primary key columns for a table.

    Tries information_schema first, falls back to pg_catalog if needed.

    Args:
        table_name: Name of the table
        schema_name: Schema name. If None, uses 'public' or config schema_name.

    Returns:
        List of primary key column names
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    # Try information_schema first
    try:
        query = """
            SELECT kcu.column_name
            FROM information_schema.table_constraints tc
            JOIN information_schema.key_column_usage kcu
                ON tc.constraint_name = kcu.constraint_name
                AND tc.table_schema = kcu.table_schema
            WHERE tc.constraint_type = 'PRIMARY KEY'
              AND tc.table_schema = %s
              AND tc.table_name = %s
            ORDER BY kcu.ordinal_position;
        """

        with self.conn.cursor() as cursor:
            cursor.execute(query, (schema_name, table_name))
            results = [row[0] for row in cursor.fetchall()]
            # If we got results and information_schema is reliable, return them
            if results and self._check_information_schema_reliable(schema_name):
                return results
            # Otherwise fall back to pg_catalog
            logger.debug(
                f"information_schema returned no results or is unreliable, "
                f"falling back to pg_catalog for primary keys of '{schema_name}.{table_name}'"
            )
    except Exception as e:
        logger.debug(
            f"information_schema query failed: {e}, falling back to pg_catalog"
        )

    # Fallback to pg_catalog
    return self._get_primary_keys_pg_catalog(table_name, schema_name)

get_table_columns(table_name, schema_name=None)

Get columns for a specific table with types and descriptions.

Tries information_schema first, falls back to pg_catalog if needed.

Parameters:

Name	Type	Description	Default
table_name	str	Name of the table	required
schema_name	str | None	Schema name. If None, uses 'public' or config schema_name.	None

Returns:

Type	Description
list[dict[str, Any]]	List of column information dictionaries with keys:
list[dict[str, Any]]	name, type, description, is_nullable, column_default

Source code in graflo/db/postgres/conn.py

def get_table_columns(
    self, table_name: str, schema_name: str | None = None
) -> list[dict[str, Any]]:
    """Get columns for a specific table with types and descriptions.

    Tries information_schema first, falls back to pg_catalog if needed.

    Args:
        table_name: Name of the table
        schema_name: Schema name. If None, uses 'public' or config schema_name.

    Returns:
        List of column information dictionaries with keys:
        name, type, description, is_nullable, column_default
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    # Try information_schema first
    try:
        query = """
            SELECT
                c.column_name as name,
                c.data_type as type,
                c.udt_name as udt_name,
                c.character_maximum_length,
                c.is_nullable,
                c.column_default,
                COALESCE(d.description, '') as description,
                c.ordinal_position as ordinal_position
            FROM information_schema.columns c
            LEFT JOIN pg_catalog.pg_statio_all_tables st
                ON st.schemaname = c.table_schema
                AND st.relname = c.table_name
            LEFT JOIN pg_catalog.pg_description d
                ON d.objoid = st.relid
                AND d.objsubid = c.ordinal_position
            WHERE c.table_schema = %s
              AND c.table_name = %s
            ORDER BY c.ordinal_position;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name, table_name))
            columns = []
            for row in cursor.fetchall():
                col_dict = dict(row)
                # Format type with length if applicable
                if col_dict["character_maximum_length"]:
                    col_dict["type"] = (
                        f"{col_dict['type']}({col_dict['character_maximum_length']})"
                    )
                # Use udt_name if it's more specific (e.g., varchar, int4)
                if (
                    col_dict["udt_name"]
                    and col_dict["udt_name"] != col_dict["type"]
                ):
                    col_dict["type"] = col_dict["udt_name"]
                # Remove helper fields
                col_dict.pop("character_maximum_length", None)
                col_dict.pop("udt_name", None)
                columns.append(col_dict)

            # If we got results and information_schema is reliable, return them
            if columns and self._check_information_schema_reliable(schema_name):
                return columns
            # Otherwise fall back to pg_catalog
            logger.debug(
                f"information_schema returned no results or is unreliable, "
                f"falling back to pg_catalog for table '{schema_name}.{table_name}'"
            )
    except Exception as e:
        logger.debug(
            f"information_schema query failed: {e}, falling back to pg_catalog"
        )

    # Fallback to pg_catalog
    return self._get_table_columns_pg_catalog(table_name, schema_name)

get_table_row_count_estimate(table_name, schema_name=None)

Return approximate row count from pg_class.reltuples (updated by ANALYZE).

Avoids full table scan; may be stale until next ANALYZE.

Source code in graflo/db/postgres/conn.py

def get_table_row_count_estimate(
    self, table_name: str, schema_name: str | None = None
) -> int | None:
    """Return approximate row count from pg_class.reltuples (updated by ANALYZE).

    Avoids full table scan; may be stale until next ANALYZE.
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"
    query = """
        SELECT c.reltuples::bigint AS estimate
        FROM pg_catalog.pg_class c
        JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
        WHERE n.nspname = %s AND c.relname = %s AND c.relkind = 'r';
    """
    with self.conn.cursor() as cursor:
        cursor.execute(query, (schema_name, table_name))
        row = cursor.fetchone()
        if row is None:
            return None
        val = row[0]
        return int(val) if val is not None else None

get_table_sample_rows(table_name, schema_name=None, limit=5)

Return limit rows from the table, in a stable order.

Ordered by the primary key, or by the whole row's text when the table has none, so sampling unchanged data returns the same rows every time. Without an order PostgreSQL returns rows in storage order, which an update or VACUUM FULL changes -- and a sample that feeds a prompt then changes the prompt, and with it any response cache keyed on it. Ordering a table without a key sorts it in full, which is the price of a repeatable sample there.

Source code in graflo/db/postgres/conn.py

def get_table_sample_rows(
    self,
    table_name: str,
    schema_name: str | None = None,
    limit: int = 5,
) -> list[dict[str, Any]]:
    """Return `limit` rows from the table, in a stable order.

    Ordered by the primary key, or by the whole row's text when the table
    has none, so sampling unchanged data returns the same rows every time.
    Without an order PostgreSQL returns rows in storage order, which an
    update or ``VACUUM FULL`` changes -- and a sample that feeds a prompt
    then changes the prompt, and with it any response cache keyed on it.
    Ordering a table without a key sorts it in full, which is the price of
    a repeatable sample there.
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"
    primary_key = self.get_primary_keys(table_name, schema_name)
    order_by = (
        sql.SQL(", ").join(sql.Identifier("t", column) for column in primary_key)
        if primary_key
        else sql.SQL("t::text")
    )
    query = sql.SQL("SELECT * FROM {}.{} AS t ORDER BY {} LIMIT %s").format(
        sql.Identifier(schema_name),
        sql.Identifier(table_name),
        order_by,
    )
    try:
        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (limit,))
            return [dict(row) for row in cursor.fetchall()]
    except Exception as e:
        logger.debug(
            f"Could not fetch sample rows for '{schema_name}.{table_name}': {e}"
        )
        return []

get_tables(schema_name=None)

Get all tables in the specified schema.

Tries information_schema first, falls back to pg_catalog if needed.

Parameters:

Name	Type	Description	Default
schema_name	str | None	Schema name to query. If None, uses 'public' or config schema_name.	None

Returns:

Type	Description
list[dict[str, Any]]	List of table information dictionaries with keys: table_name,
list[dict[str, Any]]	table_schema, description (the table comment, empty when none)

Source code in graflo/db/postgres/conn.py

def get_tables(self, schema_name: str | None = None) -> list[dict[str, Any]]:
    """Get all tables in the specified schema.

    Tries information_schema first, falls back to pg_catalog if needed.

    Args:
        schema_name: Schema name to query. If None, uses 'public' or config schema_name.

    Returns:
        List of table information dictionaries with keys: table_name,
        table_schema, description (the table comment, empty when none)
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    # Try information_schema first
    try:
        query = """
            SELECT
                table_name,
                table_schema,
                COALESCE(
                    obj_description(
                        (quote_ident(table_schema) || '.' || quote_ident(table_name))::regclass,
                        'pg_class'
                    ),
                    ''
                ) AS description
            FROM information_schema.tables
            WHERE table_schema = %s
              AND table_type = 'BASE TABLE'
            ORDER BY table_name;
        """

        with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(query, (schema_name,))
            results = [dict(row) for row in cursor.fetchall()]
            # If we got results, check if information_schema is reliable
            if results and self._check_information_schema_reliable(schema_name):
                return results
            # If no results or unreliable, fall back to pg_catalog
            logger.debug(
                f"information_schema returned no results or is unreliable, "
                f"falling back to pg_catalog for schema '{schema_name}'"
            )
    except Exception as e:
        logger.debug(
            f"information_schema query failed: {e}, falling back to pg_catalog"
        )

    # Fallback to pg_catalog
    return self._get_tables_pg_catalog(schema_name)

get_unique_columns(table_name, schema_name=None)

Get column names that participate in any UNIQUE constraint.

Tries information_schema first, falls back to pg_catalog if needed.

Source code in graflo/db/postgres/conn.py

def get_unique_columns(
    self, table_name: str, schema_name: str | None = None
) -> list[str]:
    """Get column names that participate in any UNIQUE constraint.

    Tries information_schema first, falls back to pg_catalog if needed.
    """
    if schema_name is None:
        schema_name = self.config.schema_name or "public"

    try:
        query = """
            SELECT kcu.column_name
            FROM information_schema.table_constraints tc
            JOIN information_schema.key_column_usage kcu
                ON tc.constraint_name = kcu.constraint_name
                AND tc.table_schema = kcu.table_schema
            WHERE tc.constraint_type = 'UNIQUE'
              AND tc.table_schema = %s
              AND tc.table_name = %s
            ORDER BY kcu.ordinal_position;
        """
        with self.conn.cursor() as cursor:
            cursor.execute(query, (schema_name, table_name))
            results = list(dict.fromkeys(row[0] for row in cursor.fetchall()))
            if results and self._check_information_schema_reliable(schema_name):
                return results
    except Exception as e:
        logger.debug(
            f"information_schema query failed for unique columns: {e}, "
            "falling back to pg_catalog"
        )
    return self._get_unique_columns_pg_catalog(table_name, schema_name)

introspect_schema(schema_name=None, include_raw_tables=False, entity_tables=None)

Classify the schema into vertex-like and edge-like tables.

Source code in graflo/db/postgres/conn.py

def introspect_schema(
    self,
    schema_name: str | None = None,
    include_raw_tables: bool = False,
    entity_tables: Collection[str] | None = None,
) -> SchemaIntrospectionResult:
    """Classify the schema into vertex-like and edge-like tables."""
    return sql_introspect.introspect_schema(
        self,
        schema_name,
        include_raw_tables,
        default_schema=self._default_schema(),
        entity_tables=entity_tables,
    )

read(query, params=None)

Execute a SELECT query and return results as a list of dictionaries.

Parameters:

Name	Type	Description	Default
query	str	SQL SELECT query to execute	required
params	tuple | dict[str, Any] | None	Optional parameters — a tuple for %s placeholders, or a mapping for named %(name)s ones.	None

Returns:

Type	Description
list[dict[str, Any]]	List of dictionaries, where each dictionary represents a row with column names as keys.
list[dict[str, Any]]	Decimal values are converted to float for compatibility with graph databases.

Source code in graflo/db/postgres/conn.py

def read(
    self, query: str, params: tuple | dict[str, Any] | None = None
) -> list[dict[str, Any]]:
    """Execute a SELECT query and return results as a list of dictionaries.

    Args:
        query: SQL SELECT query to execute
        params: Optional parameters — a tuple for ``%s`` placeholders, or a
            mapping for named ``%(name)s`` ones.

    Returns:
        List of dictionaries, where each dictionary represents a row with column names as keys.
        Decimal values are converted to float for compatibility with graph databases.
    """
    from decimal import Decimal

    with self.conn.cursor(cursor_factory=RealDictCursor) as cursor:
        if params:
            cursor.execute(query, params)
        else:
            cursor.execute(query)

        # Convert rows to dictionaries and convert Decimal to float
        results = []
        for row in cursor.fetchall():
            row_dict = dict(row)
            # Convert Decimal to float for JSON/graph database compatibility
            for key, value in row_dict.items():
                if isinstance(value, Decimal):
                    row_dict[key] = float(value)
            results.append(row_dict)

        return results

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.

Source code in graflo/db/postgres/resource_mapping.py

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 __init__(self, fuzzy_threshold: float = 0.8):
        """Initialize the resource mapper.

        Args:
            fuzzy_threshold: Similarity threshold for fuzzy matching (0.0 to 1.0, default 0.8)
        """
        self.fuzzy_threshold = fuzzy_threshold

    def create_vertex_resource(
        self,
        table_name: str,
        vertex_name: str,
        reference_edges: Sequence[ReferenceEdgeInfo] = (),
    ) -> Resource:
        """Create a Resource for a vertex table.

        The pipeline is a single ``{"vertex": vertex_name}`` step, followed for
        each of *reference_edges* by the vertex the row refers to and the edge
        to it. Field renames driven by reserved-word sanitization are added
        later by :class:`graflo.hq.sanitizer.Sanitizer`.

        Args:
            table_name: Name of the PostgreSQL table
            vertex_name: Name of the vertex type (typically same as table_name)
            reference_edges: Edges stated by the table's own foreign keys

        Returns:
            Resource: Resource configured to ingest vertex data
        """
        references = [
            (edge.target_table, edge.relation, edge.columns) for edge in reference_edges
        ]
        resource = Resource(
            name=table_name,
            pipeline=subject_with_references(vertex_name, references),
            infer_edges=not references,
        )

        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,
        matcher: FuzzyMatcher,
    ) -> Resource:
        """Create a Resource for an edge table.

        Field renames driven by reserved-word sanitization are added later by
        :class:`graflo.hq.sanitizer.Sanitizer`.

        Args:
            edge_table_info: Edge table information from introspection
            vertex_config: Vertex configuration for source/target validation
            matcher: Optional fuzzy matcher for better performance (with caching enabled)

        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

        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"
            )

        source_vertex_obj = vertex_config[source_table]
        target_vertex_obj = vertex_config[target_table]

        source_pk_fields = list(source_vertex_obj.identity)
        target_pk_fields = list(target_vertex_obj.identity)

        source_pk_field = self._infer_pk_field_from_column(
            source_column,
            source_table,
            source_pk_fields,
            matcher,
        )
        target_pk_field = self._infer_pk_field_from_column(
            target_column,
            target_table,
            target_pk_fields,
            matcher,
        )

        apply: list[dict[str, Any]] = []

        if source_column:
            apply.append(
                {
                    "vertex": source_table,
                    "from": {source_pk_field: source_column},
                }
            )

        if target_column:
            apply.append(
                {
                    "vertex": target_table,
                    "from": {target_pk_field: target_column},
                }
            )

        resource = Resource(
            name=table_name,
            pipeline=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

    def _infer_pk_field_from_column(
        self,
        column_name: str,
        vertex_name: str,
        pk_fields: list[str],
        matcher: FuzzyMatcher,
    ) -> 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
            matcher: Optional fuzzy matcher for better performance (with caching enabled)

        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
            matched_vertex = matcher.get_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 create_resources_from_tables(
        self,
        introspection_result: SchemaIntrospectionResult,
        vertex_config: VertexConfig,
        edge_config: EdgeConfig,
        fuzzy_threshold: float | None = None,
    ) -> list[Resource]:
        """Create Resources from PostgreSQL tables.

        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
            fuzzy_threshold: Similarity threshold for fuzzy matching (0.0 to 1.0)

        Returns:
            list[Resource]: List of Resources for all tables
        """
        resources = []

        vertex_names = list(vertex_config.vertex_set)
        threshold = (
            fuzzy_threshold if fuzzy_threshold is not None else self.fuzzy_threshold
        )
        from graflo.util.fuzzy_matcher import FuzzyMatcher

        matcher = FuzzyMatcher(vertex_names, threshold, enable_cache=True)

        vertex_tables = introspection_result.vertex_tables
        for table_info in vertex_tables:
            table_name = table_info.name
            vertex_name = table_name
            # Only the references the schema declares an edge for.
            reference_edges = [
                edge
                for edge in introspection_result.reference_edges
                if edge.source_table == table_name
                and (edge.source_table, edge.target_table, edge.relation) in edge_config
            ]
            resource = self.create_vertex_resource(
                table_name, vertex_name, reference_edges
            )
            resources.append(resource)

        edge_tables = introspection_result.edge_tables
        for edge_table_info in edge_tables:
            try:
                resource = self.create_edge_resource(
                    edge_table_info, vertex_config, matcher
                )
                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

Attributes

fuzzy_threshold = fuzzy_threshold instance-attribute

Methods:

__init__(fuzzy_threshold=0.8)

Initialize the resource mapper.

Parameters:

Name	Type	Description	Default
fuzzy_threshold	float	Similarity threshold for fuzzy matching (0.0 to 1.0, default 0.8)	0.8

Source code in graflo/db/postgres/resource_mapping.py

def __init__(self, fuzzy_threshold: float = 0.8):
    """Initialize the resource mapper.

    Args:
        fuzzy_threshold: Similarity threshold for fuzzy matching (0.0 to 1.0, default 0.8)
    """
    self.fuzzy_threshold = fuzzy_threshold

create_edge_resource(edge_table_info, vertex_config, matcher)

Create a Resource for an edge table.

Field renames driven by reserved-word sanitization are added later by :class:graflo.hq.sanitizer.Sanitizer.

Parameters:

Name	Type	Description	Default
edge_table_info	EdgeTableInfo	Edge table information from introspection	required
vertex_config	VertexConfig	Vertex configuration for source/target validation	required
matcher	FuzzyMatcher	Optional fuzzy matcher for better performance (with caching enabled)	required

Returns:

Name	Type	Description
Resource	Resource	Resource configured to ingest edge data

Source code in graflo/db/postgres/resource_mapping.py

def create_edge_resource(
    self,
    edge_table_info: EdgeTableInfo,
    vertex_config: VertexConfig,
    matcher: FuzzyMatcher,
) -> Resource:
    """Create a Resource for an edge table.

    Field renames driven by reserved-word sanitization are added later by
    :class:`graflo.hq.sanitizer.Sanitizer`.

    Args:
        edge_table_info: Edge table information from introspection
        vertex_config: Vertex configuration for source/target validation
        matcher: Optional fuzzy matcher for better performance (with caching enabled)

    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

    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"
        )

    source_vertex_obj = vertex_config[source_table]
    target_vertex_obj = vertex_config[target_table]

    source_pk_fields = list(source_vertex_obj.identity)
    target_pk_fields = list(target_vertex_obj.identity)

    source_pk_field = self._infer_pk_field_from_column(
        source_column,
        source_table,
        source_pk_fields,
        matcher,
    )
    target_pk_field = self._infer_pk_field_from_column(
        target_column,
        target_table,
        target_pk_fields,
        matcher,
    )

    apply: list[dict[str, Any]] = []

    if source_column:
        apply.append(
            {
                "vertex": source_table,
                "from": {source_pk_field: source_column},
            }
        )

    if target_column:
        apply.append(
            {
                "vertex": target_table,
                "from": {target_pk_field: target_column},
            }
        )

    resource = Resource(
        name=table_name,
        pipeline=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

create_resources_from_tables(introspection_result, vertex_config, edge_config, fuzzy_threshold=None)

Create Resources from PostgreSQL tables.

Creates Resources for both vertex and edge tables, enabling ingestion of the entire database schema.

Parameters:

Name	Type	Description	Default
introspection_result	SchemaIntrospectionResult	Result from PostgresConnection.introspect_schema()	required
vertex_config	VertexConfig	Inferred vertex configuration	required
edge_config	EdgeConfig	Inferred edge configuration	required
fuzzy_threshold	float | None	Similarity threshold for fuzzy matching (0.0 to 1.0)	None

Returns:

Type	Description
list[Resource]	list[Resource]: List of Resources for all tables

Source code in graflo/db/postgres/resource_mapping.py

def create_resources_from_tables(
    self,
    introspection_result: SchemaIntrospectionResult,
    vertex_config: VertexConfig,
    edge_config: EdgeConfig,
    fuzzy_threshold: float | None = None,
) -> list[Resource]:
    """Create Resources from PostgreSQL tables.

    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
        fuzzy_threshold: Similarity threshold for fuzzy matching (0.0 to 1.0)

    Returns:
        list[Resource]: List of Resources for all tables
    """
    resources = []

    vertex_names = list(vertex_config.vertex_set)
    threshold = (
        fuzzy_threshold if fuzzy_threshold is not None else self.fuzzy_threshold
    )
    from graflo.util.fuzzy_matcher import FuzzyMatcher

    matcher = FuzzyMatcher(vertex_names, threshold, enable_cache=True)

    vertex_tables = introspection_result.vertex_tables
    for table_info in vertex_tables:
        table_name = table_info.name
        vertex_name = table_name
        # Only the references the schema declares an edge for.
        reference_edges = [
            edge
            for edge in introspection_result.reference_edges
            if edge.source_table == table_name
            and (edge.source_table, edge.target_table, edge.relation) in edge_config
        ]
        resource = self.create_vertex_resource(
            table_name, vertex_name, reference_edges
        )
        resources.append(resource)

    edge_tables = introspection_result.edge_tables
    for edge_table_info in edge_tables:
        try:
            resource = self.create_edge_resource(
                edge_table_info, vertex_config, matcher
            )
            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

create_vertex_resource(table_name, vertex_name, reference_edges=())

Create a Resource for a vertex table.

The pipeline is a single {"vertex": vertex_name} step, followed for each of reference_edges by the vertex the row refers to and the edge to it. Field renames driven by reserved-word sanitization are added later by :class:graflo.hq.sanitizer.Sanitizer.

Parameters:

Name	Type	Description	Default
table_name	str	Name of the PostgreSQL table	required
vertex_name	str	Name of the vertex type (typically same as table_name)	required
reference_edges	Sequence[ReferenceEdgeInfo]	Edges stated by the table's own foreign keys	()

Returns:

Name	Type	Description
Resource	Resource	Resource configured to ingest vertex data

Source code in graflo/db/postgres/resource_mapping.py

def create_vertex_resource(
    self,
    table_name: str,
    vertex_name: str,
    reference_edges: Sequence[ReferenceEdgeInfo] = (),
) -> Resource:
    """Create a Resource for a vertex table.

    The pipeline is a single ``{"vertex": vertex_name}`` step, followed for
    each of *reference_edges* by the vertex the row refers to and the edge
    to it. Field renames driven by reserved-word sanitization are added
    later by :class:`graflo.hq.sanitizer.Sanitizer`.

    Args:
        table_name: Name of the PostgreSQL table
        vertex_name: Name of the vertex type (typically same as table_name)
        reference_edges: Edges stated by the table's own foreign keys

    Returns:
        Resource: Resource configured to ingest vertex data
    """
    references = [
        (edge.target_table, edge.relation, edge.columns) for edge in reference_edges
    ]
    resource = Resource(
        name=table_name,
        pipeline=subject_with_references(vertex_name, references),
        infer_edges=not references,
    )

    logger.debug(
        f"Created vertex resource '{table_name}' for vertex '{vertex_name}'"
    )

    return resource

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.

Source code in graflo/db/postgres/schema_inference.py

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: DBType = DBType.ARANGO,
        provider: SqlMetadataProvider | None = None,
        type_mapper: SqlTypeMapper | None = None,
    ):
        """Initialize the schema inferencer.

        Args:
            db_flavor: Target database flavor for the inferred schema
            provider: Metadata source, used to sample rows when refining a
                declared column type. Optional — without it the declared type
                is taken as-is, which is what a catalogue-only inference does.
            type_mapper: Type-name mapping; defaults to the dialect-neutral one.
        """
        self.db_flavor = db_flavor
        self.type_mapper = type_mapper or SqlTypeMapper()
        self.provider = provider

    def infer_vertex_config(
        self, introspection_result: SchemaIntrospectionResult
    ) -> VertexConfig:
        """Infer VertexConfig from vertex tables.

        Args:
            introspection_result: Result from PostgresConnection.introspect_schema()

        Returns:
            VertexConfig: Inferred vertex configuration
        """
        vertex_tables = introspection_result.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
                raw_type, item_type = self.type_mapper.map_field(col.type)
                if raw_type == FieldType.LIST.value:
                    if item_type is None:
                        # Avoid inventing a wrong scalar; leave untyped rather than LIST without item
                        fields.append(Field(name=field_name, type=None))
                    else:
                        fields.append(
                            Field(
                                name=field_name,
                                type=FieldType.LIST,
                                item_type=FieldType(item_type),
                            )
                        )
                else:
                    field_type = FieldType(raw_type) if raw_type else None
                    fields.append(Field(name=field_name, type=field_type))

            # Create vertex
            vertex = Vertex(
                name=table_name,
                properties=fields,
                identity=list(pk_columns),
            )

            vertices.append(vertex)
            logger.debug(
                f"Inferred vertex '{table_name}' with {len(fields)} fields and "
                f"identity {list(pk_columns)}"
            )

        return VertexConfig(vertices=vertices)

    def _infer_type_from_samples(
        self, table_name: str, schema_name: str, column_name: str, pg_type: str
    ) -> str:
        """Infer field type by sampling 5 rows from the table.

        Uses heuristics to determine if a column contains integers, floats, datetimes, etc.
        Falls back to PostgreSQL type mapping if sampling fails or is unavailable.

        Args:
            table_name: Name of the table
            schema_name: Schema name
            column_name: Name of the column to sample
            pg_type: PostgreSQL type from schema introspection

        Returns:
            str: FieldType value (INT, FLOAT, DATETIME, STRING, etc.)
        """
        # First try PostgreSQL type mapping
        mapped_type = self.type_mapper.map_type(pg_type)
        if mapped_type == FieldType.LIST.value:
            # Sampling does not refine LIST; caller should use map_field for item_type
            return mapped_type

        # With a metadata source we can sample rows and refine the declared type
        if self.provider is None:
            logger.debug(
                f"No connection available for sampling, using mapped type '{mapped_type}' "
                f"for column '{column_name}' in table '{table_name}'"
            )
            return mapped_type

        try:
            samples = self.provider.get_table_sample_rows(
                table_name, schema_name or None, limit=5
            )

            if not samples:
                logger.debug(
                    f"No samples found for column '{column_name}' in table '{table_name}', "
                    f"using mapped type '{mapped_type}'"
                )
                return mapped_type

            # Extract non-None values
            values = [
                row[column_name] for row in samples if row[column_name] is not None
            ]

            if not values:
                logger.debug(
                    f"All samples are NULL for column '{column_name}' in table '{table_name}', "
                    f"using mapped type '{mapped_type}'"
                )
                return mapped_type

            # Heuristics to infer type from values
            # Check for integers (all values are integers)
            if all(isinstance(v, int) for v in values):
                logger.debug(
                    f"Inferred INT type for column '{column_name}' in table '{table_name}' "
                    f"from samples"
                )
                return FieldType.INT.value

            # Check for floats (all values are floats or ints that could be floats)
            if all(isinstance(v, (int, float)) for v in values):
                # If any value has decimal part, it's a float
                if any(isinstance(v, float) and v != float(int(v)) for v in values):
                    logger.debug(
                        f"Inferred FLOAT type for column '{column_name}' in table '{table_name}' "
                        f"from samples"
                    )
                    return FieldType.FLOAT.value
                # All integers, but might be stored as float - check PostgreSQL type
                if mapped_type == FieldType.FLOAT.value:
                    return FieldType.FLOAT.value
                return FieldType.INT.value

            # Check for datetime/date objects
            from datetime import date, datetime, time

            if all(isinstance(v, (datetime, date, time)) for v in values):
                logger.debug(
                    f"Inferred DATETIME type for column '{column_name}' in table '{table_name}' "
                    f"from samples"
                )
                return FieldType.DATETIME.value

            # Check for ISO format datetime strings
            if all(isinstance(v, str) for v in values):
                # Try to parse as ISO datetime
                iso_datetime_count = 0
                for v in values:
                    try:
                        # Try ISO format (with or without timezone)
                        datetime.fromisoformat(v)
                        iso_datetime_count += 1
                    except (ValueError, AttributeError):
                        # Try other common formats
                        try:
                            datetime.strptime(v, "%Y-%m-%d %H:%M:%S")
                            iso_datetime_count += 1
                        except ValueError:
                            try:
                                datetime.strptime(v, "%Y-%m-%d")
                                iso_datetime_count += 1
                            except ValueError:
                                pass

                # If most values look like datetimes, infer DATETIME
                if iso_datetime_count >= len(values) * 0.8:  # 80% threshold
                    logger.debug(
                        f"Inferred DATETIME type for column '{column_name}' in table '{table_name}' "
                        f"from ISO format strings"
                    )
                    return FieldType.DATETIME.value

            # Default to mapped type
            logger.debug(
                f"Using mapped type '{mapped_type}' for column '{column_name}' in table '{table_name}' "
                f"(could not infer from samples)"
            )
            return mapped_type

        except Exception as e:
            logger.warning(
                f"Error sampling data for column '{column_name}' in table '{table_name}': {e}. "
                f"Using mapped type '{mapped_type}'"
            )
            return mapped_type

    def infer_edge_weights(self, edge_table_info: EdgeTableInfo) -> list[Field] | None:
        """Infer edge attributes from edge table columns with types.

        Uses PostgreSQL column types and optionally samples data to infer accurate types.

        Args:
            edge_table_info: Edge table information from introspection

        Returns:
            List of attribute fields if there are non-key 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

        # Create Field objects with types for each weight column
        direct_weights = []
        for col in weight_columns:
            raw_type, item_type = self.type_mapper.map_field(col.type)
            if raw_type == FieldType.LIST.value:
                if item_type is None:
                    direct_weights.append(Field(name=col.name, type=None))
                else:
                    direct_weights.append(
                        Field(
                            name=col.name,
                            type=FieldType.LIST,
                            item_type=FieldType(item_type),
                        )
                    )
                continue
            # Refine non-list types from samples when useful
            refined = self._infer_type_from_samples(
                edge_table_info.name,
                edge_table_info.schema_name,
                col.name,
                col.type,
            )
            direct_weights.append(
                Field(name=col.name, type=FieldType(refined) if refined else None)
            )

        logger.debug(
            f"Inferred {len(direct_weights)} weights for edge table "
            f"'{edge_table_info.name}': {[f.name for f in direct_weights]}"
        )

        return direct_weights

    def infer_edge_config(
        self,
        introspection_result: SchemaIntrospectionResult,
        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.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

            # 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

            attrs = self.infer_edge_weights(edge_table_info) or []
            edge = Edge(
                source=source_table,
                target=target_table,
                properties=attrs,
                relation=edge_table_info.relation,
            )

            edges.append(edge)
            logger.debug(
                f"Inferred edge '{table_name}' from {source_table} to {target_table}"
            )

        # Foreign keys inside entity tables: a row refers to a row.
        declared = {edge.edge_id for edge in edges}
        for reference in introspection_result.reference_edges:
            edge = Edge(
                source=reference.source_table,
                target=reference.target_table,
                relation=reference.relation,
            )
            if (
                edge.edge_id in declared
                or reference.source_table not in vertex_names
                or reference.target_table not in vertex_names
            ):
                continue
            declared.add(edge.edge_id)
            edges.append(edge)

        return EdgeConfig(edges=edges)

    def infer_schema(
        self,
        introspection_result: SchemaIntrospectionResult,
        schema_name: str | None = None,
    ) -> Schema:
        """Infer complete Schema from PostgreSQL introspection.

        Args:
            introspection_result: Result from PostgresConnection.introspect_schema()
            schema_name: Schema name (defaults to schema_name from introspection if None)

        Returns:
            Schema: Complete inferred schema with vertices, edges, and metadata
        """
        if schema_name is None:
            schema_name = introspection_result.schema_name

        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.values()))
        logger.info(f"Inferred {edges_count} edges")

        # Create schema metadata
        metadata = GraphMetadata(name=schema_name)

        # Create schema (ingestion model is created separately)
        schema = Schema(
            metadata=metadata,
            core_schema=CoreSchema(
                vertex_config=vertex_config, edge_config=edge_config
            ),
            db_profile=DatabaseProfile(
                db_flavor=self.db_flavor,
                vertex_storage_names={v.name: v.name for v in vertex_config.vertices},
            ),
        )

        logger.info(
            f"Successfully inferred schema '{schema_name}' with "
            f"{len(vertex_config.vertices)} vertices and "
            f"{len(list(edge_config.values()))} edges"
        )

        return schema

Attributes

db_flavor = db_flavor instance-attribute

provider = provider instance-attribute

type_mapper = type_mapper or SqlTypeMapper() instance-attribute

Methods:

__init__(db_flavor=DBType.ARANGO, provider=None, type_mapper=None)

Initialize the schema inferencer.

Parameters:

Name	Type	Description	Default
db_flavor	DBType	Target database flavor for the inferred schema	ARANGO
provider	SqlMetadataProvider | None	Metadata source, used to sample rows when refining a declared column type. Optional — without it the declared type is taken as-is, which is what a catalogue-only inference does.	None
type_mapper	SqlTypeMapper | None	Type-name mapping; defaults to the dialect-neutral one.	None

Source code in graflo/db/postgres/schema_inference.py

def __init__(
    self,
    db_flavor: DBType = DBType.ARANGO,
    provider: SqlMetadataProvider | None = None,
    type_mapper: SqlTypeMapper | None = None,
):
    """Initialize the schema inferencer.

    Args:
        db_flavor: Target database flavor for the inferred schema
        provider: Metadata source, used to sample rows when refining a
            declared column type. Optional — without it the declared type
            is taken as-is, which is what a catalogue-only inference does.
        type_mapper: Type-name mapping; defaults to the dialect-neutral one.
    """
    self.db_flavor = db_flavor
    self.type_mapper = type_mapper or SqlTypeMapper()
    self.provider = provider

infer_edge_config(introspection_result, vertex_config)

Infer EdgeConfig from edge tables.

Parameters:

Name	Type	Description	Default
introspection_result	SchemaIntrospectionResult	Result from PostgresConnection.introspect_schema()	required
vertex_config	VertexConfig	Inferred vertex configuration	required

Returns:

Name	Type	Description
EdgeConfig	EdgeConfig	Inferred edge configuration

Source code in graflo/db/postgres/schema_inference.py

def infer_edge_config(
    self,
    introspection_result: SchemaIntrospectionResult,
    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.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

        # 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

        attrs = self.infer_edge_weights(edge_table_info) or []
        edge = Edge(
            source=source_table,
            target=target_table,
            properties=attrs,
            relation=edge_table_info.relation,
        )

        edges.append(edge)
        logger.debug(
            f"Inferred edge '{table_name}' from {source_table} to {target_table}"
        )

    # Foreign keys inside entity tables: a row refers to a row.
    declared = {edge.edge_id for edge in edges}
    for reference in introspection_result.reference_edges:
        edge = Edge(
            source=reference.source_table,
            target=reference.target_table,
            relation=reference.relation,
        )
        if (
            edge.edge_id in declared
            or reference.source_table not in vertex_names
            or reference.target_table not in vertex_names
        ):
            continue
        declared.add(edge.edge_id)
        edges.append(edge)

    return EdgeConfig(edges=edges)

infer_edge_weights(edge_table_info)

Infer edge attributes from edge table columns with types.

Uses PostgreSQL column types and optionally samples data to infer accurate types.

Parameters:

Name	Type	Description	Default
edge_table_info	EdgeTableInfo	Edge table information from introspection	required

Returns:

Type	Description
list[Field] | None	List of attribute fields if there are non-key columns, None otherwise.

Source code in graflo/db/postgres/schema_inference.py

def infer_edge_weights(self, edge_table_info: EdgeTableInfo) -> list[Field] | None:
    """Infer edge attributes from edge table columns with types.

    Uses PostgreSQL column types and optionally samples data to infer accurate types.

    Args:
        edge_table_info: Edge table information from introspection

    Returns:
        List of attribute fields if there are non-key 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

    # Create Field objects with types for each weight column
    direct_weights = []
    for col in weight_columns:
        raw_type, item_type = self.type_mapper.map_field(col.type)
        if raw_type == FieldType.LIST.value:
            if item_type is None:
                direct_weights.append(Field(name=col.name, type=None))
            else:
                direct_weights.append(
                    Field(
                        name=col.name,
                        type=FieldType.LIST,
                        item_type=FieldType(item_type),
                    )
                )
            continue
        # Refine non-list types from samples when useful
        refined = self._infer_type_from_samples(
            edge_table_info.name,
            edge_table_info.schema_name,
            col.name,
            col.type,
        )
        direct_weights.append(
            Field(name=col.name, type=FieldType(refined) if refined else None)
        )

    logger.debug(
        f"Inferred {len(direct_weights)} weights for edge table "
        f"'{edge_table_info.name}': {[f.name for f in direct_weights]}"
    )

    return direct_weights

infer_schema(introspection_result, schema_name=None)

Infer complete Schema from PostgreSQL introspection.

Parameters:

Name	Type	Description	Default
introspection_result	SchemaIntrospectionResult	Result from PostgresConnection.introspect_schema()	required
schema_name	str | None	Schema name (defaults to schema_name from introspection if None)	None

Returns:

Name	Type	Description
Schema	Schema	Complete inferred schema with vertices, edges, and metadata

Source code in graflo/db/postgres/schema_inference.py

def infer_schema(
    self,
    introspection_result: SchemaIntrospectionResult,
    schema_name: str | None = None,
) -> Schema:
    """Infer complete Schema from PostgreSQL introspection.

    Args:
        introspection_result: Result from PostgresConnection.introspect_schema()
        schema_name: Schema name (defaults to schema_name from introspection if None)

    Returns:
        Schema: Complete inferred schema with vertices, edges, and metadata
    """
    if schema_name is None:
        schema_name = introspection_result.schema_name

    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.values()))
    logger.info(f"Inferred {edges_count} edges")

    # Create schema metadata
    metadata = GraphMetadata(name=schema_name)

    # Create schema (ingestion model is created separately)
    schema = Schema(
        metadata=metadata,
        core_schema=CoreSchema(
            vertex_config=vertex_config, edge_config=edge_config
        ),
        db_profile=DatabaseProfile(
            db_flavor=self.db_flavor,
            vertex_storage_names={v.name: v.name for v in vertex_config.vertices},
        ),
    )

    logger.info(
        f"Successfully inferred schema '{schema_name}' with "
        f"{len(vertex_config.vertices)} vertices and "
        f"{len(list(edge_config.values()))} edges"
    )

    return schema

infer_vertex_config(introspection_result)

Infer VertexConfig from vertex tables.

Parameters:

Name	Type	Description	Default
introspection_result	SchemaIntrospectionResult	Result from PostgresConnection.introspect_schema()	required

Returns:

Name	Type	Description
VertexConfig	VertexConfig	Inferred vertex configuration

Source code in graflo/db/postgres/schema_inference.py

def infer_vertex_config(
    self, introspection_result: SchemaIntrospectionResult
) -> VertexConfig:
    """Infer VertexConfig from vertex tables.

    Args:
        introspection_result: Result from PostgresConnection.introspect_schema()

    Returns:
        VertexConfig: Inferred vertex configuration
    """
    vertex_tables = introspection_result.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
            raw_type, item_type = self.type_mapper.map_field(col.type)
            if raw_type == FieldType.LIST.value:
                if item_type is None:
                    # Avoid inventing a wrong scalar; leave untyped rather than LIST without item
                    fields.append(Field(name=field_name, type=None))
                else:
                    fields.append(
                        Field(
                            name=field_name,
                            type=FieldType.LIST,
                            item_type=FieldType(item_type),
                        )
                    )
            else:
                field_type = FieldType(raw_type) if raw_type else None
                fields.append(Field(name=field_name, type=field_type))

        # Create vertex
        vertex = Vertex(
            name=table_name,
            properties=fields,
            identity=list(pk_columns),
        )

        vertices.append(vertex)
        logger.debug(
            f"Inferred vertex '{table_name}' with {len(fields)} fields and "
            f"identity {list(pk_columns)}"
        )

    return VertexConfig(vertices=vertices)