Skip to content

Commit 03b5b2e

Browse files
jaytaphjaytaph
committed
Add one-off copy of beheer data into admin schema
1 parent aeac990 commit 03b5b2e

3 files changed

Lines changed: 129 additions & 0 deletions

File tree

‎docker/Dockerfile‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ COPY --chown=${APP_USER}:${APP_GROUP} app.conf.example /src/
9696
COPY --chown=${APP_USER}:${APP_GROUP} static /src/static
9797
COPY --chown=${APP_USER}:${APP_GROUP} sql /src/sql
9898
COPY --chown=${APP_USER}:${APP_GROUP} tools/migrate_db.sh /src/tools/migrate_db.sh
99+
COPY --chown=${APP_USER}:${APP_GROUP} tools/migrate_beheer_data.py /src/tools/migrate_beheer_data.py
99100

100101
# version.json is generated from build args so CI workspace artifacts never leak in
101102
ARG RELEASE_VERSION=dev

‎docker/init.sh‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ fi
1414
echo "Migrating"
1515

1616
DSN=$(grep dsn $APP_PATH | sed -r 's/dsn=postgresql\+psycopg/postgresql/' ) tools/./migrate_db.sh
17+
DSN=$(grep dsn $APP_PATH | sed -r 's/dsn=postgresql\+psycopg/postgresql/' ) python tools/migrate_beheer_data.py
1718

1819
echo "Start main process"
1920
python -m app.main

‎tools/migrate_beheer_data.py‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
"""
2+
One-off copy of the nvi-beheer-api data (separate database, public schema) into the admin schema of the NVI database.
3+
4+
Runs after migrate_db.sh. Requires DSN (target, NVI database) and SOURCE_DSN (nvi_beheer_api_db). Without SOURCE_DSN
5+
nothing happens. The copy runs in a single transaction and is recorded in the migrations table, so it runs only once.
6+
"""
7+
8+
import os
9+
import sys
10+
11+
import psycopg
12+
13+
MARKER = "beheer-data-copy"
14+
15+
# In foreign key order
16+
TABLES = [
17+
"organizations",
18+
"scopes",
19+
"clients",
20+
"certificates",
21+
"sources",
22+
"organizations_scopes",
23+
"clients_scopes",
24+
"clients_certificates",
25+
"clients_sources",
26+
]
27+
28+
29+
class CopyError(Exception):
30+
pass
31+
32+
33+
def columns(conn: psycopg.Connection, table: str) -> list[str]:
34+
rows = conn.execute(
35+
"SELECT column_name FROM information_schema.columns "
36+
"WHERE table_schema = 'admin' AND table_name = %s ORDER BY ordinal_position",
37+
(table,),
38+
).fetchall()
39+
return [row[0] for row in rows]
40+
41+
42+
def count(conn: psycopg.Connection, schema: str, table: str) -> int:
43+
row = conn.execute(f'SELECT count(*) FROM {schema}."{table}"').fetchone()
44+
return row[0] if row else 0
45+
46+
47+
def copy_table(source: psycopg.Connection, target: psycopg.Connection, table: str) -> None:
48+
cols = ", ".join(f'"{c}"' for c in columns(target, table))
49+
with (
50+
source.cursor().copy(f'COPY (SELECT {cols} FROM public."{table}") TO STDOUT') as src,
51+
target.cursor().copy(f'COPY admin."{table}" ({cols}) FROM STDIN') as dst,
52+
):
53+
for block in src:
54+
dst.write(block)
55+
56+
57+
def report_referrals(target: psycopg.Connection) -> None:
58+
without_org, only_deleted = target.execute(
59+
"""
60+
SELECT
61+
count(*) FILTER (WHERE NOT EXISTS (
62+
SELECT 1 FROM admin.organizations o WHERE o.external_id = r.ura_number)),
63+
count(*) FILTER (WHERE EXISTS (
64+
SELECT 1 FROM admin.organizations o WHERE o.external_id = r.ura_number)
65+
AND NOT EXISTS (
66+
SELECT 1 FROM admin.organizations o WHERE o.external_id = r.ura_number AND o.deleted_at IS NULL))
67+
FROM public.referrals r
68+
"""
69+
).fetchone() or (0, 0)
70+
print(f"Referrals without any organization: {without_org}")
71+
print(f"Referrals with only deleted organizations: {only_deleted}")
72+
73+
74+
def main() -> int:
75+
target_dsn = os.environ.get("DSN")
76+
source_dsn = os.environ.get("SOURCE_DSN")
77+
if not target_dsn:
78+
print("DSN is not set", file=sys.stderr)
79+
return 1
80+
if not source_dsn:
81+
print("SOURCE_DSN is not set, skipping beheer data copy")
82+
return 0
83+
84+
try:
85+
copy(target_dsn, source_dsn)
86+
except CopyError as e:
87+
print(e, file=sys.stderr)
88+
return 1
89+
return 0
90+
91+
92+
def copy(target_dsn: str, source_dsn: str) -> None:
93+
# Raising inside the connection blocks rolls back the transaction
94+
with psycopg.connect(target_dsn) as target, psycopg.connect(source_dsn) as source:
95+
target.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (MARKER,))
96+
if target.execute("SELECT 1 FROM migrations WHERE name = %s", (MARKER,)).fetchone():
97+
print("Beheer data already copied, skipping")
98+
return
99+
100+
not_empty = [t for t in TABLES if t != "scopes" and count(target, "admin", t) > 0]
101+
if not_empty:
102+
raise CopyError(f"Target admin tables are not empty: {', '.join(not_empty)}")
103+
104+
source.execute("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY")
105+
106+
# Scopes are seeded by 026, replace them with the source rows so ids match
107+
target.execute("DELETE FROM admin.scopes")
108+
109+
for table in TABLES:
110+
copy_table(source, target, table)
111+
expected, actual = count(source, "public", table), count(target, "admin", table)
112+
if expected != actual:
113+
raise CopyError(f"{table}: copied {actual} rows, expected {expected}")
114+
print(f"{table}: {actual} rows")
115+
116+
target.execute(
117+
"SELECT setval(pg_get_serial_sequence('admin.scopes', 'id'), coalesce(max(id), 0) + 1, false) "
118+
"FROM admin.scopes"
119+
)
120+
report_referrals(target)
121+
target.execute("INSERT INTO migrations (name, checksum) VALUES (%s, %s)", (MARKER, "-"))
122+
123+
print("Beheer data copied")
124+
125+
126+
if __name__ == "__main__":
127+
sys.exit(main())

0 commit comments

Comments
 (0)