graflo.db.postgres.schema_inference

Schema inference from PostgreSQL database introspection.

This module provides functionality to infer graflo Schema objects from PostgreSQL 3NF database schemas by analyzing table structures, relationships, and column types.

Attributes

logger = logging.getLogger(__name__) module-attribute

Classes

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 = []
            list_columns: set[str] = set()
            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:
                    list_columns.add(field_name)
                    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),
                secondary_identities=secondary_identities_from_unique(
                    table_info.unique_constraints,
                    pk_columns,
                    [f.name for f in fields],
                    list_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,
        endpoint_fields: Collection[str] = (),
    ) -> 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
            endpoint_fields: Identity fields of the edge's endpoints; a column of
                the same name is renamed (:func:`edge_property_names`).

        A primary-key column that does not name an endpoint is kept: it is what
        the row is known by (see :meth:`edge_key`).

        Returns:
            List of attribute fields if there are non-key columns, None otherwise.
        """
        columns = edge_table_info.columns
        endpoint_columns = {fk.column for fk in edge_table_info.foreign_keys} | {
            edge_table_info.source_column,
            edge_table_info.target_column,
        }

        # Every column but the endpoints' (these become weights)
        weight_columns = [col for col in columns if col.name not in endpoint_columns]
        names = edge_property_names(
            edge_table_info.name, [col.name for col in weight_columns], endpoint_fields
        )

        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)
            name = names[col.name]
            if raw_type == FieldType.LIST.value:
                if item_type is None:
                    direct_weights.append(Field(name=name, type=None))
                else:
                    direct_weights.append(
                        Field(
                            name=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=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

    @staticmethod
    def edge_key(edge_table_info: EdgeTableInfo) -> list[str]:
        """Primary-key columns of an edge table that do not name an endpoint.

        An edge is otherwise keyed by its endpoints and relation, so two rows
        between the same pair -- the same client ordering the same item twice --
        would fold into one edge. Keyed by these columns as well, each row stays
        its own edge.
        """
        endpoint_columns = {fk.column for fk in edge_table_info.foreign_keys} | {
            edge_table_info.source_column,
            edge_table_info.target_column,
        }
        return [
            column
            for column in edge_table_info.primary_key
            if column not in endpoint_columns
        ]

    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

            endpoint_fields = edge_endpoint_fields(
                vertex_config, source_table, target_table
            )
            attrs = self.infer_edge_weights(edge_table_info, endpoint_fields) or []
            names = edge_property_names(
                table_name, self.edge_key(edge_table_info), endpoint_fields
            )
            key = list(names.values())
            edge = Edge(
                source=source_table,
                target=target_table,
                properties=attrs,
                relation=edge_table_info.relation,
                identities=[key] if key else [],
            )

            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

edge_key(edge_table_info) staticmethod

Primary-key columns of an edge table that do not name an endpoint.

An edge is otherwise keyed by its endpoints and relation, so two rows between the same pair -- the same client ordering the same item twice -- would fold into one edge. Keyed by these columns as well, each row stays its own edge.

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

@staticmethod
def edge_key(edge_table_info: EdgeTableInfo) -> list[str]:
    """Primary-key columns of an edge table that do not name an endpoint.

    An edge is otherwise keyed by its endpoints and relation, so two rows
    between the same pair -- the same client ordering the same item twice --
    would fold into one edge. Keyed by these columns as well, each row stays
    its own edge.
    """
    endpoint_columns = {fk.column for fk in edge_table_info.foreign_keys} | {
        edge_table_info.source_column,
        edge_table_info.target_column,
    }
    return [
        column
        for column in edge_table_info.primary_key
        if column not in endpoint_columns
    ]

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

        endpoint_fields = edge_endpoint_fields(
            vertex_config, source_table, target_table
        )
        attrs = self.infer_edge_weights(edge_table_info, endpoint_fields) or []
        names = edge_property_names(
            table_name, self.edge_key(edge_table_info), endpoint_fields
        )
        key = list(names.values())
        edge = Edge(
            source=source_table,
            target=target_table,
            properties=attrs,
            relation=edge_table_info.relation,
            identities=[key] if key else [],
        )

        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, endpoint_fields=())

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
endpoint_fields	Collection[str]	Identity fields of the edge's endpoints; a column of the same name is renamed (:func:edge_property_names).	()

A primary-key column that does not name an endpoint is kept: it is what the row is known by (see :meth:edge_key).

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,
    endpoint_fields: Collection[str] = (),
) -> 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
        endpoint_fields: Identity fields of the edge's endpoints; a column of
            the same name is renamed (:func:`edge_property_names`).

    A primary-key column that does not name an endpoint is kept: it is what
    the row is known by (see :meth:`edge_key`).

    Returns:
        List of attribute fields if there are non-key columns, None otherwise.
    """
    columns = edge_table_info.columns
    endpoint_columns = {fk.column for fk in edge_table_info.foreign_keys} | {
        edge_table_info.source_column,
        edge_table_info.target_column,
    }

    # Every column but the endpoints' (these become weights)
    weight_columns = [col for col in columns if col.name not in endpoint_columns]
    names = edge_property_names(
        edge_table_info.name, [col.name for col in weight_columns], endpoint_fields
    )

    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)
        name = names[col.name]
        if raw_type == FieldType.LIST.value:
            if item_type is None:
                direct_weights.append(Field(name=name, type=None))
            else:
                direct_weights.append(
                    Field(
                        name=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=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 = []
        list_columns: set[str] = set()
        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:
                list_columns.add(field_name)
                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),
            secondary_identities=secondary_identities_from_unique(
                table_info.unique_constraints,
                pk_columns,
                [f.name for f in fields],
                list_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)

Functions:

edge_endpoint_fields(vertex_config, source, target)

Fields an edge reads from its endpoint vertices: their identities.

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

def edge_endpoint_fields(
    vertex_config: VertexConfig, source: str, target: str
) -> set[str]:
    """Fields an edge reads from its endpoint vertices: their identities."""
    return set(vertex_config.identity_fields(source)) | set(
        vertex_config.identity_fields(target)
    )

secondary_identities_from_unique(unique_constraints, primary_key, columns, list_columns=())

One secondary identity per UNIQUE constraint usable as a lookup key.

A constraint is skipped when its column set equals the primary key or an earlier constraint, or when it names an array (LIST) or unknown column. Each is named unique_<col1>_<col2>… in constraint column order.

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

def secondary_identities_from_unique(
    unique_constraints: list[list[str]],
    primary_key: list[str],
    columns: Collection[str],
    list_columns: Collection[str] = (),
) -> list[SecondaryIdentity]:
    """One secondary identity per UNIQUE constraint usable as a lookup key.

    A constraint is skipped when its column set equals the primary key or an
    earlier constraint, or when it names an array (LIST) or unknown column.
    Each is named ``unique_<col1>_<col2>…`` in constraint column order.
    """
    known = set(columns) - set(list_columns)
    seen: set[frozenset[str]] = {frozenset(primary_key)}
    names: set[str] = set()
    out: list[SecondaryIdentity] = []
    for constraint in unique_constraints:
        field_set = frozenset(constraint)
        if not constraint or field_set in seen or not field_set <= known:
            continue
        seen.add(field_set)
        base = "unique_" + "_".join(constraint)
        name, n = base, 2
        while name in names:
            name, n = f"{base}_{n}", n + 1
        names.add(name)
        out.append(SecondaryIdentity(name=name, fields=list(constraint)))
    return out