From c4b625392a0af2d306fd1be2042324e7a583165a Mon Sep 17 00:00:00 2001
From: Federico Ceratto <federico@openobservatory.org>
Date: Fri, 20 Nov 2020 13:52:31 +0000
Subject: [PATCH 1/2] Install database_upgrade_schema

---
 af/fastpath/setup.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/af/fastpath/setup.py b/af/fastpath/setup.py
index 0641b7cc..b8d34b14 100644
--- a/af/fastpath/setup.py
+++ b/af/fastpath/setup.py
@@ -16,6 +16,7 @@
         "fastpath=fastpath.core:main",
         "reprocessor=fastpath.reprocessor:main",
         "domain_input_updater=fastpath.domain_input:main",
+        "database_upgrade_schema=fastpath.database_upgrade_schema:main",
     ]},
     install_requires=REQUIRED,
     include_package_data=True,

From 0ad44ff1343993675889bee56a605c37886b6ae3 Mon Sep 17 00:00:00 2001
From: Federico Ceratto <federico@openobservatory.org>
Date: Fri, 23 Jul 2021 16:02:48 +0100
Subject: [PATCH 2/2] Add save/diff functions

---
 af/fastpath/database_upgrade_schema.py | 54 ++++++++++++++++++++++++--
 1 file changed, 51 insertions(+), 3 deletions(-)

diff --git a/af/fastpath/database_upgrade_schema.py b/af/fastpath/database_upgrade_schema.py
index 4db612ed..709c8a2d 100755
--- a/af/fastpath/database_upgrade_schema.py
+++ b/af/fastpath/database_upgrade_schema.py
@@ -6,6 +6,9 @@
 """
 
 from argparse import ArgumentParser
+from pathlib import Path
+from difflib import context_diff
+import subprocess
 
 # Refresh / compare using
 # ssh ams-pg.ooni.org -t pg_dump metadb -U shovel --schema-only | \
@@ -13,6 +16,43 @@
 #
 # meld database_upgrade_schema.py database_upgrade_schema.sql
 
+DBNAME = "metadb"
+DBUSER = "shovel"
+
+
+def pg_dump(dbname, dbuser):
+    cmd = ["pg_dump", dbname, "-U", dbuser, "--schema-only"]
+    return subprocess.check_output(cmd).decode()
+
+
+def pg_dump_sorted(dbname, dbuser):
+    lines = pg_dump(dbname, dbuser).splitlines()
+    lines = [i for i in lines if not i.startswith("--")]
+    lines = [i.rstrip() for i in lines]
+    chunks = "\n".join(lines).split(";")
+    chunks = [c.strip() for c in chunks if c.strip()]
+    return sorted(chunks)
+
+
+def diff_local_db_vs_file(conf):
+    fname = conf.diff
+    file_chunks = Path(fname).read_text().split(";\n")
+    if file_chunks[-1] == "":
+        file_chunks.pop()
+
+    db_chunks = pg_dump_sorted(conf.dbname, conf.dbuser)
+    g = context_diff(db_chunks, file_chunks, fromfile="DB", tofile=fname)
+    g = list(g)
+    for i in g:
+        print(i)
+
+
+def save_local_db_schema(conf):
+    f = Path(conf.save)
+    db_chunks = pg_dump_sorted(conf.dbname, conf.dbuser)
+    out = ";\n".join(db_chunks) + ";\n"
+    f.write_text(out)
+
 
 def db_setup():
     """Setup fastpath database from scratch"""
@@ -263,13 +303,15 @@ def db_drop():
     raise NotImplementedError
 
 
-steps = (
-    (db_setup, db_drop),
-)
+steps = ((db_setup, db_drop),)
 
 if __name__ == "__main__":
     ap = ArgumentParser()
     ap.add_argument("--list", action="store_true", help="list steps")
+    ap.add_argument("--diff", help="diff local DB schema and sql file")
+    ap.add_argument("--save", help="save local DB schema to sql file")
+    ap.add_argument("--dbname", default=DBNAME)
+    ap.add_argument("--dbuser", default=DBUSER)
     conf = ap.parse_args()
 
     if conf.list:
@@ -277,3 +319,9 @@ def db_drop():
             print(f"{rf.__name__} - {rf.__doc__}")
             print(f"  {rollback.__name__} - {rollback.__doc__}")
             print()
+
+    elif conf.diff:
+        diff_local_db_vs_file(conf)
+
+    elif conf.save:
+        save_local_db_schema(conf)