class PostgresTargetWriteMixin:
"""Mixin implementing :class:`~graflo.db.conn.Connection` target operations."""
flavor = DBType.POSTGRES
supports_schema_introspection = True
# The graph shape lives in the table layout, which `information_schema`
# reports in full; nothing here is sampled.
schema_introspection_is_sampled = False
config: Any
conn: _Psycopg2Conn
# Supplied by Connection, which follows this mixin in the MRO: annotate
# rather than stub, so the real implementation is not shadowed.
define_indexes: Any
report_edge_direction_support: Any
def read(
self, query: str, params: tuple | dict[str, Any] | None = None
) -> list[dict[str, Any]]:
raise NotImplementedError
def get_tables(self, schema_name: str | None = None) -> list[dict[str, Any]]:
raise NotImplementedError
def get_table_columns(
self, table_name: str, schema_name: str | None = None
) -> list[dict[str, Any]]:
raise NotImplementedError
def _execute_write(self, query: str, params: tuple | list | None = None) -> None:
with self.conn.cursor() as cursor:
if params is not None:
cursor.execute(query, params)
else:
cursor.execute(query)
self.conn.commit()
def create_database(self, name: str) -> None:
schema_name = name
q = sql.SQL("CREATE SCHEMA IF NOT EXISTS {}").format(
sql.Identifier(schema_name)
)
with self.conn.cursor() as cursor:
cursor.execute(q)
self.conn.commit()
def delete_database(self, name: str) -> None:
q = sql.SQL("DROP SCHEMA IF EXISTS {} CASCADE").format(sql.Identifier(name))
with self.conn.cursor() as cursor:
cursor.execute(q)
self.conn.commit()
def execute(self, query: str | Any, **kwargs: Any) -> Any:
params = kwargs.get("params")
if isinstance(query, str) and query.strip().upper().startswith("SELECT"):
return self.read(query, params)
self._execute_write(str(query), params)
return None
def define_schema(self, schema: Schema) -> None:
self._target_schema = schema
from graflo.db.field_type_support import assert_schema_field_types_supported
assert_schema_field_types_supported(DBType.POSTGRES, schema)
self._define_postgres_tables(schema)
def define_vertex_classes(self, schema: Schema) -> None:
self._define_vertex_tables(schema)
def define_edge_classes(self, edges: list[Edge]) -> None:
for edge in edges:
self._create_edge_table(edge)
def delete_graph_structure(
self,
vertex_types: tuple[str, ...] | list[str] = (),
graph_names: tuple[str, ...] | list[str] = (),
delete_all: bool = False,
) -> None:
pg_schema = _pg_schema_name(self.config)
present = [row["table_name"] for row in self.get_tables(schema_name=pg_schema)]
tables: list[str] = []
if delete_all:
tables = list(present)
else:
requested = [vertex_table_name(v) for v in vertex_types]
tables.extend(requested)
# Dropping a vertex type must drop the edges incident to it, the way
# every graph backend does (Neo4j DETACH DELETE, Arango dropping the
# graph's edge collections). PostgreSQL stores edges in freestanding
# tables that no foreign key ties to the vertex table, so without
# this they survive every drop and accumulate in the namespace.
dropped = set(requested)
vertex_universe = dropped | {
t for t in present if not t.endswith(EDGE_TABLE_SUFFIX)
}
for table in present:
parts = split_edge_table_name(table, vertex_universe)
if parts is None:
continue
source, target, _ = parts
if source in dropped or target in dropped:
tables.append(table)
for table in tables:
q = sql.SQL("DROP TABLE IF EXISTS {}.{} CASCADE").format(
sql.Identifier(pg_schema),
sql.Identifier(table),
)
with self.conn.cursor() as cursor:
cursor.execute(q)
self.conn.commit()
def _pg_schema_exists(self, schema_name: str) -> bool:
rows = self.read(
"SELECT schema_name FROM information_schema.schemata WHERE schema_name = %s",
(schema_name,),
)
return bool(rows)
def ensure_target_namespace(self, schema: Schema, *, create: bool) -> None:
"""Ensure the PostgreSQL schema namespace exists."""
pg_schema = _pg_schema_name(self.config)
if self._pg_schema_exists(pg_schema):
return
if not create:
raise NamespaceNotFoundError(
f"PostgreSQL schema '{pg_schema}' does not exist. "
"Create it manually or call with create_namespace=True."
)
self.create_database(pg_schema)
def apply_target_schema(
self,
schema: Schema,
*,
recreate: bool,
create_namespace: bool = True,
) -> None:
"""Create vertex/edge tables for the schema."""
self.report_edge_direction_support(schema)
pg_schema = _pg_schema_name(self.config)
existing = {row["table_name"] for row in self.get_tables(schema_name=pg_schema)}
expected_vertices = {
vertex_table_name(v.name) for v in schema.core_schema.vertex_config.vertices
}
expected_edges = {
edge_table_name(e.source, e.target, e.relation)
for e in schema.core_schema.edge_config.values()
}
expected = expected_vertices | expected_edges
overlap = existing & expected
if overlap and not recreate:
raise SchemaExistsError(
f"PostgreSQL tables already exist in schema '{pg_schema}': "
f"{sorted(overlap)}"
)
if recreate and overlap:
self.delete_graph_structure(vertex_types=tuple(expected), delete_all=False)
if create_namespace and not self._pg_schema_exists(pg_schema):
self.create_database(pg_schema)
self.define_schema(schema)
self.define_indexes(schema)
def init_db(
self,
schema: Schema,
recreate_schema: bool = False,
*,
create_namespace: bool = True,
) -> None:
"""Convenience wrapper: ensure schema namespace then apply tables."""
self.ensure_target_namespace(schema, create=create_namespace)
self.apply_target_schema(
schema, recreate=recreate_schema, create_namespace=create_namespace
)
def clear_data(self, schema: Schema) -> None:
pg_schema = _pg_schema_name(self.config)
table_names = [
vertex_table_name(v.name) for v in schema.core_schema.vertex_config.vertices
]
table_names.extend(
edge_table_name(e.source, e.target, e.relation)
for e in schema.core_schema.edge_config.values()
)
with self.conn.cursor() as cursor:
for table in table_names:
q = sql.SQL("TRUNCATE TABLE {}.{} CASCADE").format(
sql.Identifier(pg_schema),
sql.Identifier(table),
)
try:
cursor.execute(q)
except Exception:
logger.debug("Skipping truncate for missing table %s", table)
self.conn.commit()
def _define_postgres_tables(self, schema: Schema) -> None:
self._define_vertex_tables(schema)
self.define_edge_classes(list(schema.core_schema.edge_config.values()))
def _define_vertex_tables(self, schema: Schema) -> None:
pg_schema = _pg_schema_name(self.config)
for vertex in schema.core_schema.vertex_config.vertices:
columns = {f.name: _pg_column_type_for_field(f) for f in vertex.properties}
for ident in vertex.identity:
columns.setdefault(ident, _PG_TEXT)
if not columns:
columns["id"] = _PG_TEXT
identity = vertex.identity or ["id"]
col_defs = [
sql.SQL("{} {}").format(sql.Identifier(name), sql.SQL(col_type))
for name, col_type in columns.items()
]
pk = sql.SQL(", ").join(sql.Identifier(i) for i in identity)
create_q = sql.SQL(
"CREATE TABLE IF NOT EXISTS {}.{} ({}, PRIMARY KEY ({}))"
).format(
sql.Identifier(pg_schema),
sql.Identifier(vertex_table_name(vertex.name)),
sql.SQL(", ").join(col_defs),
pk,
)
with self.conn.cursor() as cursor:
cursor.execute(create_q)
self.conn.commit()
def _create_edge_table(self, edge: Edge) -> None:
pg_schema = _pg_schema_name(self.config)
table = edge_table_name(edge.source, edge.target, edge.relation)
source_table = vertex_table_name(edge.source)
target_table = vertex_table_name(edge.target)
src_pk = "id"
tgt_pk = "id"
schema = getattr(self, "_target_schema", None)
if schema is not None:
vc = schema.core_schema.vertex_config
src_fields = vc.identity_fields(edge.source)
tgt_fields = vc.identity_fields(edge.target)
if src_fields:
src_pk = src_fields[0]
if tgt_fields:
tgt_pk = tgt_fields[0]
weight_cols = list(edge.properties) if edge.properties else []
col_defs: list[sql.Composable] = [
sql.SQL("{} BIGSERIAL PRIMARY KEY").format(sql.Identifier("id")),
sql.SQL("{} {} NOT NULL").format(
sql.Identifier("source_id"), sql.SQL(_PG_TEXT)
),
sql.SQL("{} {} NOT NULL").format(
sql.Identifier("target_id"), sql.SQL(_PG_TEXT)
),
]
for field in weight_cols:
col_defs.append(
sql.SQL("{} {}").format(
sql.Identifier(field.name),
sql.SQL(_pg_column_type_for_field(field)),
)
)
fk_clauses: list[sql.Composable] = []
fk_source = sql.SQL("FOREIGN KEY (source_id) REFERENCES {}.{} ({})").format(
sql.Identifier(pg_schema),
sql.Identifier(source_table),
sql.Identifier(src_pk),
)
fk_target = sql.SQL("FOREIGN KEY (target_id) REFERENCES {}.{} ({})").format(
sql.Identifier(pg_schema),
sql.Identifier(target_table),
sql.Identifier(tgt_pk),
)
fk_clauses = [fk_source, fk_target]
create_q = sql.SQL("CREATE TABLE IF NOT EXISTS {}.{} ({})").format(
sql.Identifier(pg_schema),
sql.Identifier(table),
sql.SQL(", ").join([*col_defs, *fk_clauses]),
)
with self.conn.cursor() as cursor:
try:
cursor.execute(create_q)
except Exception as exc:
logger.warning(
"Edge table %s creation with FK failed: %s; creating without FK",
table,
exc,
)
create_q_no_fk = sql.SQL(
"CREATE TABLE IF NOT EXISTS {}.{} ({})"
).format(
sql.Identifier(pg_schema),
sql.Identifier(table),
sql.SQL(", ").join(col_defs),
)
cursor.execute(create_q_no_fk)
if weight_cols:
# weight_cols holds Field objects; the index needs column names.
unique_cols = sql.SQL(", ").join(
sql.Identifier(column)
for column in (
"source_id",
"target_id",
*(field.name for field in weight_cols),
)
)
idx_q = sql.SQL(
"CREATE UNIQUE INDEX IF NOT EXISTS {} ON {}.{} ({})"
).format(
sql.Identifier(_edge_unique_index_name(table)),
sql.Identifier(pg_schema),
sql.Identifier(table),
unique_cols,
)
cursor.execute(idx_q)
self.conn.commit()
def upsert_docs_batch(
self,
docs: list[dict[str, Any]],
class_name: str,
match_keys: list[str] | tuple[str, ...],
**kwargs: Any,
) -> None:
if kwargs.get("dry") or not docs:
return
pg_schema = _pg_schema_name(self.config)
table = vertex_table_name(class_name)
match_keys = tuple(match_keys) or ("id",)
all_keys: list[str] = []
for doc in docs:
all_keys.extend(doc.keys())
columns = sorted({k for k in all_keys if not k.startswith("_")})
if not columns:
return
update_cols = [c for c in columns if c not in match_keys]
col_idents = sql.SQL(", ").join(sql.Identifier(c) for c in columns)
conflict = sql.SQL(", ").join(sql.Identifier(k) for k in match_keys)
if update_cols:
set_clause = sql.SQL(", ").join(
sql.SQL("{} = EXCLUDED.{}").format(sql.Identifier(c), sql.Identifier(c))
for c in update_cols
)
upsert_q = sql.SQL(
"INSERT INTO {}.{} ({}) VALUES %s ON CONFLICT ({}) DO UPDATE SET {}"
).format(
sql.Identifier(pg_schema),
sql.Identifier(table),
col_idents,
conflict,
set_clause,
)
else:
upsert_q = sql.SQL(
"INSERT INTO {}.{} ({}) VALUES %s ON CONFLICT ({}) DO NOTHING"
).format(
sql.Identifier(pg_schema),
sql.Identifier(table),
col_idents,
conflict,
)
values = [tuple(doc.get(c) for c in columns) for doc in docs]
with self.conn.cursor() as cursor:
execute_values(cursor, upsert_q, values)
self.conn.commit()
def insert_edges_batch(
self,
docs_edges: list[list[dict[str, Any]]] | list[Any] | None,
source_class: str,
target_class: str,
relation_name: str | None,
match_keys_source: tuple[str, ...],
match_keys_target: tuple[str, ...],
filter_uniques: bool = True,
head: int | None = None,
**kwargs: Any,
) -> None:
if kwargs.get("dry") or not docs_edges:
return
if head is not None:
docs_edges = docs_edges[:head]
pg_schema = _pg_schema_name(self.config)
table = edge_table_name(source_class, target_class, relation_name)
match_keys_source = match_keys_source or ("id",)
match_keys_target = match_keys_target or ("id",)
src_key = match_keys_source[0]
tgt_key = match_keys_target[0]
rows: list[tuple] = []
weight_keys: set[str] = set()
for item in docs_edges:
if not isinstance(item, (list, tuple)) or len(item) < 2:
continue
source_doc, target_doc = item[0], item[1]
weight = item[2] if len(item) > 2 and isinstance(item[2], dict) else {}
weight_keys.update(weight.keys())
rows.append(
(
source_doc.get(src_key),
target_doc.get(tgt_key),
weight,
)
)
if not rows:
return
columns = ["source_id", "target_id", *sorted(weight_keys)]
col_idents = sql.SQL(", ").join(sql.Identifier(c) for c in columns)
# No conflict target: the edge table's unique index covers
# (source_id, target_id) plus any weight columns, so naming a fixed pair
# fails with "no unique or exclusion constraint matching" as soon as the
# edge carries properties. A bare DO NOTHING matches whichever index exists.
upsert_q = sql.SQL(
"INSERT INTO {}.{} ({}) VALUES %s ON CONFLICT DO NOTHING"
).format(
sql.Identifier(pg_schema),
sql.Identifier(table),
col_idents,
)
values = [
(
source_id,
target_id,
*[weight.get(k) for k in sorted(weight_keys)],
)
for source_id, target_id, weight in rows
if source_id is not None and target_id is not None
]
if not values:
return
with self.conn.cursor() as cursor:
execute_values(cursor, upsert_q, values)
self.conn.commit()
def insert_return_batch(
self, docs: list[dict[str, Any]], class_name: str
) -> list[dict[str, Any]] | str:
raise NotImplementedError(
"insert_return_batch is not implemented for PostgreSQL"
)
def fetch_docs(
self,
class_name: str,
filters: list[Any] | dict[str, Any] | None = None,
limit: int | None = None,
return_keys: list[str] | None = None,
unset_keys: list[str] | None = None,
**kwargs: Any,
) -> list[dict[str, Any]]:
pg_schema = _pg_schema_name(self.config)
table = vertex_table_name(class_name)
if return_keys:
keep = [k for k in return_keys if not unset_keys or k not in unset_keys]
select_clause = ", ".join(_quote_ident(k) for k in keep) if keep else "*"
else:
select_clause = "*"
where_clause = ""
if filters is not None:
expr = parse_filter_expression(filters)
rendered = str(expr(kind=ExpressionFlavor.SQL))
if rendered:
where_clause = f" WHERE {rendered}"
limit_clause = f" LIMIT {int(limit)}" if limit is not None else ""
q = (
f"SELECT {select_clause} FROM "
f"{_quote_ident(pg_schema)}.{_quote_ident(table)}"
f"{where_clause}{limit_clause}"
)
return self.read(q)
def fetch_edges(
self,
from_type: str,
from_id: str,
edge_type: str | None = None,
to_type: str | None = None,
to_id: str | None = None,
filters: list | dict | None = None,
limit: int | None = None,
return_keys: list | None = None,
unset_keys: list | None = None,
direction: EdgeDirection = EdgeDirection.OUT,
**kwargs,
) -> list[dict[str, Any]]:
"""Edges incident to one vertex, read from the edge table.
``edge_type`` names the edge table (the storage name), matching
``fetch_all_edges``'s ``collection_name``. Endpoints live in
``source_id`` / ``target_id``; ``define_edge_indexes`` indexes the latter,
so the inbound branch is not a sequential scan.
"""
if edge_type is None:
raise ValueError(
"PostgreSQL fetch_edges requires edge_type (the edge table name)"
)
pg_schema = _pg_schema_name(self.config)
qualified = f"{_quote_ident(pg_schema)}.{_quote_ident(edge_type)}"
extra = ""
if filters is not None:
rendered = str(parse_filter_expression(filters)(kind=ExpressionFlavor.SQL))
if rendered:
extra = f" AND ({rendered})"
far_clause = ""
if to_id is not None:
far_clause = " AND {far} = %(to_id)s"
def branch(anchor_column: str, far_column: str) -> str:
clause = far_clause.format(far=_quote_ident(far_column))
return (
f"SELECT * FROM {qualified} "
f"WHERE {_quote_ident(anchor_column)} = %(from_id)s{clause}{extra}"
)
if direction is EdgeDirection.OUT:
sql = branch("source_id", "target_id")
elif direction is EdgeDirection.IN:
sql = branch("target_id", "source_id")
else:
# No edge is both outgoing and incoming for the same anchor unless it
# is a self-loop, so UNION (not UNION ALL) also dedupes that case.
sql = f"{branch('source_id', 'target_id')} UNION {branch('target_id', 'source_id')}"
if limit is not None:
sql = f"{sql} LIMIT {int(limit)}"
params: dict[str, Any] = {"from_id": from_id}
if to_id is not None:
params["to_id"] = to_id
rows = self.read(sql, params)
if return_keys or unset_keys:
keep = set(return_keys) if return_keys else None
drop = set(unset_keys) if unset_keys else set()
rows = [
{
k: v
for k, v in row.items()
if (keep is None or k in keep) and k not in drop
}
for row in rows
]
return rows
def fetch_present_documents(
self,
batch: list[dict[str, Any]],
class_name: str,
match_keys: list[str] | tuple[str, ...],
keep_keys: list[str] | tuple[str, ...] | None = None,
flatten: bool = False,
filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]] | dict[int, list[dict[str, Any]]]:
raise NotImplementedError(
"fetch_present_documents is not implemented for PostgreSQL"
)
def aggregate(
self,
class_name: str,
aggregation_function: AggregationType,
discriminant: str | None = None,
aggregated_field: str | None = None,
filters: FilterExpression | list | dict | None = None,
) -> int | float | list[dict[str, Any]] | dict[str, int | float] | None:
"""Aggregate over a vertex table, optionally grouped by *discriminant*.
Mirrors the shape the other backends return: a list of
``{discriminant, _value}`` rows when grouping, otherwise a single
``{_value}`` row.
"""
pg_schema = _pg_schema_name(self.config)
table = vertex_table_name(class_name)
qualified = f"{_quote_ident(pg_schema)}.{_quote_ident(table)}"
sql_function = _PG_AGGREGATIONS.get(aggregation_function)
if sql_function is None:
raise ValueError(
f"Aggregation {aggregation_function!r} is not supported on PostgreSQL; "
f"supported: {sorted(a.value for a in _PG_AGGREGATIONS)}"
)
if aggregation_function == AggregationType.COUNT and aggregated_field is None:
expression = "COUNT(*)"
elif aggregated_field is None:
raise ValueError(
f"Aggregation {aggregation_function!r} requires aggregated_field"
)
else:
expression = f"{sql_function}({_quote_ident(aggregated_field)})"
where_clause = ""
if filters is not None:
rendered = str(parse_filter_expression(filters)(kind=ExpressionFlavor.SQL))
if rendered:
where_clause = f" WHERE {rendered}"
if discriminant is None:
q = f"SELECT {expression} AS _value FROM {qualified}{where_clause}"
else:
column = _quote_ident(discriminant)
q = (
f"SELECT {column} AS {_quote_ident(discriminant)}, "
f"{expression} AS _value FROM {qualified}{where_clause} "
f"GROUP BY {column}"
)
return self.read(q)
def keep_absent_documents(
self,
batch: list[dict[str, Any]],
class_name: str,
match_keys: list[str] | tuple[str, ...],
keep_keys: list[str] | tuple[str, ...] | None = None,
filters: list[Any] | dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
raise NotImplementedError(
"keep_absent_documents is not implemented for PostgreSQL"
)
def define_vertex_indexes(
self, vertex_config: VertexConfig, schema: Schema | None = None
) -> None:
"""Create the secondary indexes declared in the database profile.
The primary identity is already covered by the table's PRIMARY KEY, so
only profile-declared indexes (which include secondary identities) are
created here.
"""
if schema is None:
logger.warning(
"Schema is None: vertex secondary indexes cannot be ensured without schema"
)
return
pg_schema = _pg_schema_name(self.config)
for vertex_name in vertex_config.vertex_set:
table = vertex_table_name(vertex_name)
for index in schema.db_profile.vertex_secondary_indexes(vertex_name):
fields = [str(f) for f in index.fields]
if not fields:
continue
index_name = f"ix_{table}_{'_'.join(fields)}"
unique_clause = sql.SQL("UNIQUE ") if index.unique else sql.SQL("")
q = sql.SQL("CREATE {}INDEX IF NOT EXISTS {} ON {}.{} ({})").format(
unique_clause,
sql.Identifier(index_name),
sql.Identifier(pg_schema),
sql.Identifier(table),
sql.SQL(", ").join(sql.Identifier(f) for f in fields),
)
try:
with self.conn.cursor() as cursor:
cursor.execute(q)
self.conn.commit()
except Exception as error:
self.conn.rollback()
logger.warning(
"Failed to create index %s on %s.%s: %s",
index_name,
pg_schema,
table,
error,
)
def define_edge_indexes(
self, edges: list[Edge], schema: Schema | None = None
) -> None:
"""Index ``target_id`` on every edge table, making reverse lookup viable.
The only pre-existing edge index is the composite uniqueness constraint,
whose leading column is ``source_id`` — it cannot serve a lookup keyed on
the target, so reaching an edge from its target end meant a sequential
scan. That is the whole cost of an undirected edge on PostgreSQL, and it
is one index per table.
"""
pg_schema = _pg_schema_name(self.config)
for edge in edges:
table = edge_table_name(edge.source, edge.target, edge.relation)
index_name = f"ix_{table}_target_id"
q = sql.SQL("CREATE INDEX IF NOT EXISTS {} ON {}.{} ({})").format(
sql.Identifier(index_name),
sql.Identifier(pg_schema),
sql.Identifier(table),
sql.Identifier("target_id"),
)
try:
with self.conn.cursor() as cursor:
cursor.execute(q)
self.conn.commit()
except Exception as error:
self.conn.rollback()
logger.warning(
"Failed to create reverse-lookup index %s on %s.%s: %s",
index_name,
pg_schema,
table,
error,
)
def fetch_all_docs(
self,
class_name: str,
*,
limit: int | None = None,
) -> list[dict[str, Any]]:
return self.fetch_docs(class_name, limit=limit)
def introspect_graph_schema(
self,
schema_name: str | None = None,
*,
sample_limit: int = 100,
) -> Schema:
"""Recover a graflo Schema from a graph-shaped PostgreSQL namespace.
Reads the catalogue rather than sampling rows: the graph shape lives in
the table layout graflo writes -- one table per vertex type, and
``{source}_{target}_{relation}_edges`` with ``source_id`` / ``target_id``
for each edge type -- so ``information_schema`` answers the whole
question and ``sample_limit`` is accepted only for interface symmetry.
Distinct from :meth:`introspect_schema`, which infers a graph from an
*arbitrary* relational database by following foreign keys. This one
assumes the graflo layout and recovers exactly what was written.
"""
from graflo.db.graph_introspection import (
GraphEdgeIntrospection,
GraphIntrospectionResult,
GraphSchemaInferencer,
GraphVertexIntrospection,
infer_identity_fields,
)
pg_schema = _pg_schema_name(self.config)
present = [row["table_name"] for row in self.get_tables(schema_name=pg_schema)]
vertex_tables = [t for t in present if not t.endswith(EDGE_TABLE_SUFFIX)]
def columns(table: str) -> tuple[list[str], dict[str, FieldType]]:
names: list[str] = []
types: dict[str, FieldType] = {}
for column in self.get_table_columns(table, schema_name=pg_schema):
name = column.get("name")
if not name:
continue
names.append(name)
declared = field_type_from_postgres(column.get("type"))
if declared is not None:
types[name] = declared
return names, types
vertices: list[GraphVertexIntrospection] = []
for table in vertex_tables:
properties, types = columns(table)
vertices.append(
GraphVertexIntrospection(
name=table,
properties=properties,
identity=infer_identity_fields(properties),
property_types=types,
)
)
edges: list[GraphEdgeIntrospection] = []
for table in present:
parts = split_edge_table_name(table, vertex_tables)
if parts is None:
continue
source, target, relation = parts
names, types = columns(table)
weights = [c for c in names if c not in EDGE_ENDPOINT_COLUMNS]
edges.append(
GraphEdgeIntrospection(
source=source,
target=target,
relation=relation,
properties=weights,
property_types={
k: v for k, v in types.items() if k in set(weights)
},
collection_name=table,
)
)
introspection = GraphIntrospectionResult(
name=schema_name or pg_schema, vertices=vertices, edges=edges
)
return GraphSchemaInferencer(db_flavor=DBType.POSTGRES).infer_schema(
introspection, schema_name=schema_name or pg_schema
)
def fetch_all_edges(
self,
source_class: str,
target_class: str,
relation_name: str | None,
*,
match_keys_source: tuple[str, ...] | None = None,
match_keys_target: tuple[str, ...] | None = None,
limit: int | None = None,
collection_name: str | None = None,
) -> list[list[dict[str, Any]]]:
pg_schema = _pg_schema_name(self.config)
table = collection_name or edge_table_name(
source_class, target_class, relation_name
)
limit_clause = f" LIMIT {int(limit)}" if limit is not None else ""
q = (
f"SELECT * FROM {_quote_ident(pg_schema)}.{_quote_ident(table)}"
f"{limit_clause}"
)
rows = self.read(q)
result: list[list[dict[str, Any]]] = []
for row in rows:
source_doc = {"id": row.get("source_id")}
target_doc = {"id": row.get("target_id")}
weight = {
k: v for k, v in row.items() if k not in ("source_id", "target_id")
}
result.append([source_doc, target_doc, weight])
return result