refactor(db)!: finalize pooled WNDB v2 migration
This commit is contained in:
@@ -0,0 +1,284 @@
|
||||
from psycopg import Connection, sql
|
||||
|
||||
from .connection import project_connection, project_transaction
|
||||
from .database import acquire_model_mutation_lock
|
||||
|
||||
_MODEL_SCHEMAS = ("network", "gis")
|
||||
_ALLOWED_EXTERNAL_REFERENCES = {
|
||||
("analysis", "results", "network", "nodes"),
|
||||
("analysis", "results", "network", "links"),
|
||||
("asset", "scada_devices", "network", "nodes"),
|
||||
("asset", "scada_devices", "network", "links"),
|
||||
}
|
||||
|
||||
|
||||
def _model_tables(conn: Connection) -> list[tuple[str, str]]:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
select n.nspname as schema_name, c.relname as table_name
|
||||
from pg_class c
|
||||
join pg_namespace n on n.oid = c.relnamespace
|
||||
where n.nspname = any(%s) and c.relkind in ('r', 'p')
|
||||
order by n.nspname, c.relname
|
||||
""",
|
||||
(list(_MODEL_SCHEMAS),),
|
||||
)
|
||||
return [(row["schema_name"], row["table_name"]) for row in cur.fetchall()]
|
||||
|
||||
|
||||
def _copy_order(
|
||||
conn: Connection, tables: list[tuple[str, str]]
|
||||
) -> list[tuple[str, str]]:
|
||||
table_set = set(tables)
|
||||
dependencies: dict[tuple[str, str], set[tuple[str, str]]] = {
|
||||
table: set() for table in tables
|
||||
}
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
select source_ns.nspname as source_schema,
|
||||
source.relname as source_table,
|
||||
target_ns.nspname as target_schema,
|
||||
target.relname as target_table
|
||||
from pg_constraint constraint_row
|
||||
join pg_class source on source.oid = constraint_row.conrelid
|
||||
join pg_namespace source_ns on source_ns.oid = source.relnamespace
|
||||
join pg_class target on target.oid = constraint_row.confrelid
|
||||
join pg_namespace target_ns on target_ns.oid = target.relnamespace
|
||||
where constraint_row.contype = 'f'
|
||||
and source_ns.nspname = any(%s)
|
||||
and target_ns.nspname = any(%s)
|
||||
""",
|
||||
(list(_MODEL_SCHEMAS), list(_MODEL_SCHEMAS)),
|
||||
)
|
||||
for row in cur.fetchall():
|
||||
source = (row["source_schema"], row["source_table"])
|
||||
target = (row["target_schema"], row["target_table"])
|
||||
if source in table_set and target in table_set and source != target:
|
||||
dependencies[source].add(target)
|
||||
|
||||
ordered: list[tuple[str, str]] = []
|
||||
remaining = set(tables)
|
||||
while remaining:
|
||||
ready = sorted(
|
||||
table for table in remaining if not (dependencies[table] & remaining)
|
||||
)
|
||||
if not ready:
|
||||
cycle = ", ".join(f"{schema}.{table}" for schema, table in sorted(remaining))
|
||||
raise RuntimeError(f"Model table foreign-key cycle detected: {cycle}")
|
||||
ordered.extend(ready)
|
||||
remaining.difference_update(ready)
|
||||
return ordered
|
||||
|
||||
|
||||
def _table_columns(
|
||||
conn: Connection, schema_name: str, table_name: str
|
||||
) -> list[str]:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
select column_name
|
||||
from information_schema.columns
|
||||
where table_schema = %s and table_name = %s
|
||||
and is_generated = 'NEVER'
|
||||
order by ordinal_position
|
||||
""",
|
||||
(schema_name, table_name),
|
||||
)
|
||||
return [row["column_name"] for row in cur.fetchall()]
|
||||
|
||||
|
||||
def _external_references(conn: Connection) -> set[tuple[str, str, str, str]]:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
select source_ns.nspname as source_schema,
|
||||
source.relname as source_table,
|
||||
target_ns.nspname as target_schema,
|
||||
target.relname as target_table
|
||||
from pg_constraint constraint_row
|
||||
join pg_class source on source.oid = constraint_row.conrelid
|
||||
join pg_namespace source_ns on source_ns.oid = source.relnamespace
|
||||
join pg_class target on target.oid = constraint_row.confrelid
|
||||
join pg_namespace target_ns on target_ns.oid = target.relnamespace
|
||||
where constraint_row.contype = 'f'
|
||||
and target_ns.nspname = any(%s)
|
||||
and source_ns.nspname <> all(%s)
|
||||
""",
|
||||
(list(_MODEL_SCHEMAS), list(_MODEL_SCHEMAS)),
|
||||
)
|
||||
return {
|
||||
(
|
||||
row["source_schema"],
|
||||
row["source_table"],
|
||||
row["target_schema"],
|
||||
row["target_table"],
|
||||
)
|
||||
for row in cur.fetchall()
|
||||
}
|
||||
|
||||
|
||||
def _copy_table(
|
||||
source_conn: Connection,
|
||||
target_conn: Connection,
|
||||
schema_name: str,
|
||||
table_name: str,
|
||||
columns: list[str],
|
||||
) -> None:
|
||||
relation = sql.Identifier(schema_name, table_name)
|
||||
column_list = sql.SQL(", ").join(map(sql.Identifier, columns))
|
||||
copy_out = sql.SQL("copy {} ({}) to stdout").format(relation, column_list)
|
||||
copy_in = sql.SQL("copy {} ({}) from stdin").format(relation, column_list)
|
||||
with source_conn.cursor().copy(copy_out) as source_copy:
|
||||
with target_conn.cursor().copy(copy_in) as target_copy:
|
||||
for chunk in source_copy:
|
||||
target_copy.write(chunk)
|
||||
|
||||
|
||||
def replace_project_model(
|
||||
target_project: str,
|
||||
source_project: str,
|
||||
*,
|
||||
copy_source_scada: bool = False,
|
||||
) -> None:
|
||||
"""Atomically replace WNDB/GIS model tables from a validated staging DB.
|
||||
|
||||
Business and analysis tables remain in the target database. Historical
|
||||
analysis rows keep references that still exist and become element-neutral
|
||||
when an element disappeared. SCADA devices are retained only when their
|
||||
bound node or link still exists in the replacement model.
|
||||
|
||||
``copy_source_scada`` is used only for temporary project clones. Normal INP
|
||||
replacement keeps the target project's existing device mappings and drops
|
||||
mappings whose model element no longer exists.
|
||||
"""
|
||||
with project_connection(source_project) as source_conn, source_conn.transaction():
|
||||
with source_conn.cursor() as cur:
|
||||
cur.execute("set transaction isolation level repeatable read, read only")
|
||||
|
||||
source_tables = _model_tables(source_conn)
|
||||
source_columns = {
|
||||
table: _table_columns(source_conn, *table) for table in source_tables
|
||||
}
|
||||
source_scada_columns = (
|
||||
_table_columns(source_conn, "asset", "scada_devices")
|
||||
if copy_source_scada
|
||||
else []
|
||||
)
|
||||
copy_order = _copy_order(source_conn, source_tables)
|
||||
|
||||
with project_transaction(target_project) as target_conn:
|
||||
acquire_model_mutation_lock(target_conn, target_project)
|
||||
|
||||
target_tables = _model_tables(target_conn)
|
||||
if set(target_tables) != set(source_tables):
|
||||
missing = sorted(set(source_tables) - set(target_tables))
|
||||
extra = sorted(set(target_tables) - set(source_tables))
|
||||
raise RuntimeError(
|
||||
f"Staging/target model schema mismatch; missing={missing}, extra={extra}"
|
||||
)
|
||||
for table, columns in source_columns.items():
|
||||
if _table_columns(target_conn, *table) != columns:
|
||||
raise RuntimeError(
|
||||
f"Staging/target columns differ for {table[0]}.{table[1]}"
|
||||
)
|
||||
if copy_source_scada and _table_columns(
|
||||
target_conn, "asset", "scada_devices"
|
||||
) != source_scada_columns:
|
||||
raise RuntimeError(
|
||||
"Source/target columns differ for asset.scada_devices"
|
||||
)
|
||||
|
||||
unexpected = _external_references(target_conn) - _ALLOWED_EXTERNAL_REFERENCES
|
||||
if unexpected:
|
||||
formatted = ", ".join(
|
||||
f"{source_schema}.{source_table}->{target_schema}.{target_table}"
|
||||
for source_schema, source_table, target_schema, target_table in sorted(
|
||||
unexpected
|
||||
)
|
||||
)
|
||||
raise RuntimeError(
|
||||
f"Model replacement has unsupported external references: {formatted}"
|
||||
)
|
||||
|
||||
with target_conn.cursor() as cur:
|
||||
if not copy_source_scada:
|
||||
cur.execute(
|
||||
"create temporary table model_scada_snapshot on commit drop "
|
||||
"as table asset.scada_devices"
|
||||
)
|
||||
cur.execute("delete from asset.scada_devices")
|
||||
cur.execute(
|
||||
"create temporary table model_result_refs on commit drop as "
|
||||
"select result_id, node_id, link_id from analysis.results"
|
||||
)
|
||||
cur.execute("update analysis.results set node_id = null, link_id = null")
|
||||
for schema_name, table_name in reversed(copy_order):
|
||||
cur.execute(
|
||||
sql.SQL("delete from {}").format(
|
||||
sql.Identifier(schema_name, table_name)
|
||||
)
|
||||
)
|
||||
|
||||
for schema_name, table_name in copy_order:
|
||||
_copy_table(
|
||||
source_conn,
|
||||
target_conn,
|
||||
schema_name,
|
||||
table_name,
|
||||
source_columns[(schema_name, table_name)],
|
||||
)
|
||||
|
||||
if copy_source_scada:
|
||||
_copy_table(
|
||||
source_conn,
|
||||
target_conn,
|
||||
"asset",
|
||||
"scada_devices",
|
||||
source_scada_columns,
|
||||
)
|
||||
|
||||
with target_conn.cursor() as cur:
|
||||
if not copy_source_scada:
|
||||
cur.execute(
|
||||
"""
|
||||
insert into asset.scada_devices
|
||||
select snapshot.*
|
||||
from model_scada_snapshot snapshot
|
||||
where (
|
||||
snapshot.node_id is not null
|
||||
and exists (
|
||||
select 1 from network.nodes node
|
||||
where node.id = snapshot.node_id
|
||||
)
|
||||
) or (
|
||||
snapshot.link_id is not null
|
||||
and exists (
|
||||
select 1 from network.links link
|
||||
where link.id = snapshot.link_id
|
||||
)
|
||||
)
|
||||
"""
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
update analysis.results result
|
||||
set node_id = case
|
||||
when exists (
|
||||
select 1 from network.nodes node
|
||||
where node.id = refs.node_id
|
||||
) then refs.node_id
|
||||
else null
|
||||
end,
|
||||
link_id = case
|
||||
when exists (
|
||||
select 1 from network.links link
|
||||
where link.id = refs.link_id
|
||||
) then refs.link_id
|
||||
else null
|
||||
end
|
||||
from model_result_refs refs
|
||||
where refs.result_id = result.result_id
|
||||
"""
|
||||
)
|
||||
Reference in New Issue
Block a user