Merge lp:~ack/landscape-charm/tools-handle-partitioned into lp:~landscape/landscape-charm/tools

[Alberto Donato]

Update the db import script with Free's changes to skip partitioned resource tables.

[Free Ekanayaka (free.ekanayaka)]

+1

review: Approve

=== modified file 'db-migrator.py'
--- db-migrator.py	2016-09-13 15:04:24 +0000
+++ db-migrator.py	2016-09-14 08:39:22 +0000
@@ -6,6 +6,7 @@
 import subprocess
 import tempfile
 import shutil
+import multiprocessing
 
 import psycopg2
 
@@ -13,6 +14,25 @@
 
 STORES = ["main", "package", "session", "knowledge", "account-1", "resource-1"]
 
+DEAD_TABLES = {
+    "main": [
+        "ec2_region", "ec2_region_id_seq", "lds_trial", "lds_trial_id_seq"],
+    "account-1": [
+        "change_package_lock_request", "computer_rule_state_id_seq",
+        "computer_rule_state_pkey", "computer_rule_state_rule_id_key",
+        "remove_package_lock_request", "create_package_lock_request"],
+}
+
+INCONSISTENT_SEQUENCES = {
+    "account-1": [
+        ("computer_rule_state_id_seq", "rule_state_id_seq"),
+        ("event_log_id_seq", "event_log_entry_id_seq"),
+    ],
+    "resource-1": [
+        ("metadata_id_seq", "annotation_id_seq"),
+    ],
+}
+
 
 def parse_args(args):
     parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
@@ -117,7 +137,23 @@
         cursor.execute(
             "SELECT tablename FROM pg_tables WHERE schemaname='public'")
         for row in cursor.fetchall():
-            cursor.execute("DROP TABLE IF EXISTS {} CASCADE".format(row[0]))
+            table = row[0]
+            if store_name.startswith("resource-"):
+                # Check if this is a partioned table
+                if "_" in table:
+                    suffix = table.split("_")[-1]
+                    if len(suffix) == 8:
+                        try:
+                            int(suffix)
+                        except ValueError:
+                            # Not a partioned table
+                            pass
+                        else:
+                            # This is a partioned table, it will
+                            # be dropped by CASCADE on the parent,
+                            # so let's skip it.
+                            continue
+            cursor.execute("DROP TABLE {} CASCADE".format(table))
 
         # Drop all extensions
         cursor.execute(
@@ -148,26 +184,35 @@
         connection.close()
 
 
-def fresh_schema(schema_script):
+def fresh_schema(schema_script, stores_info):
     subprocess.check_call(schema_script)
 
 
+def rename_incosistent_sequences(stores_info):
+    for store_name, sequences in INCONSISTENT_SEQUENCES.iteritems():
+        connection = connect(stores_info, store_name, as_admin=True)
+        cursor = connection.cursor()
+        for old_name, new_name in sequences:
+            cursor.execute(
+                "ALTER TABLE IF EXISTS {} RENAME TO {}".format(
+                    old_name, new_name))
+        connection.commit()
+        connection.close()
+
+
 def pg_dump(stores_info, store_name, dump_dir):
     dump_file = os.path.join(dump_dir, store_name + ".sql")
     options = [
            "--host={}".format(stores_info["postgres"]["host"]),
            "--port={}".format(stores_info["postgres"]["port"]),
            "--format=custom",
+           "--data-only",
            "--no-acl",
            "--exclude-table=patch",
            "--file={}".format(dump_file),
     ]
-    if store_name == "main":
-        options.append("--exclude-table=ec2_region*")
-    if store_name.startswith("account-"):
-        options.append("--exclude-table=create_package_lock_request*")
-        options.append("--exclude-table=remove_package_lock_request*")
-        options.append("--exclude-table=change_package_lock_request*")
+    for table in DEAD_TABLES.get(store_name, []):
+        options.append("--exclude-table={}".format(table))
     user = stores_info["admin"]["user"]
     if user:
         options.append("--user={}".format(user))
@@ -181,9 +226,91 @@
     subprocess.check_call(cmd, env=env)
 
 
+def pg_dump_partitioned_schemas(stores_info, dump_dir):
+    store_name = "resource-1"
+    dump_file = os.path.join(dump_dir, store_name + "-partioned.sql")
+
+    options = [
+           "--host={}".format(stores_info["postgres"]["host"]),
+           "--port={}".format(stores_info["postgres"]["port"]),
+           "--schema-only",
+           "--no-acl",
+           "--file={}".format(dump_file),
+    ]
+
+    parents = {}
+
+    connection = connect(stores_info, store_name)
+    cursor = connection.cursor()
+    cursor.execute(
+        "SELECT table_name FROM information_schema.tables "
+        "WHERE table_schema='public' AND table_name ~ E'[a-z_]+_[0-9]{8}'")
+
+    for row in cursor.fetchall():
+        table = row[0]
+        options.append("--table={}".format(table))
+        parent = "_".join(table.split("_")[:-1])
+        if parent not in parents:
+            triggers = {}
+            trigger_before = parent + "_insert_trigger"
+            trigger_after = parent + "_insert_trigger_after"
+            for trigger in [trigger_before, trigger_after]:
+                cursor.execute(
+                    "SELECT pg_get_functiondef('{}'::regproc)".format(trigger))
+                triggers[trigger] = cursor.fetchone()[0]
+            parents[parent] = triggers
+    connection.close()
+
+    user = stores_info["admin"]["user"]
+    if user:
+        options.append("--user={}".format(user))
+    cmd = ["pg_dump"] + options + [stores_info["stores"][store_name]]
+    password = stores_info["admin"]["password"]
+    env = os.environ.copy()
+    if password:
+        env["PGPASSWORD"] = password
+    subprocess.check_call(cmd, env=env)
+
+    with open(dump_file, "a") as fd:
+        for parent, triggers in parents.iteritems():
+            for trigger, body in triggers.iteritems():
+                fd.write(body.strip() + ";\n")
+            trigger_before = parent + "_insert_trigger"
+            trigger_after = parent + "_insert_trigger_after"
+            fd.write(
+                "CREATE TRIGGER {} BEFORE INSERT ON {} "
+                "FOR EACH ROW EXECUTE PROCEDURE {}();\n".format(
+                    trigger_before, parent, trigger_before))
+            fd.write(
+                "CREATE TRIGGER {} AFTER INSERT ON {} "
+                "FOR EACH ROW EXECUTE PROCEDURE {}();\n".format(
+                    trigger_after, parent, trigger_after))
+
+
+def psql_restore_partitioned_schema(stores_info, dump_dir):
+    store_name = "resource-1"
+    dump_file = os.path.join(dump_dir, store_name + "-partioned.sql")
+
+    options = [
+           "--host={}".format(stores_info["postgres"]["host"]),
+           "--port={}".format(stores_info["postgres"]["port"]),
+           "--file={}".format(dump_file),
+    ]
+    user = stores_info["admin"]["user"]
+    if user:
+        options.append("--user={}".format(user))
+    cmd = ["psql"] + options + [stores_info["stores"][store_name]]
+    password = stores_info["admin"]["password"]
+    env = os.environ.copy()
+    if password:
+        env["PGPASSWORD"] = password
+    subprocess.check_call(cmd, env=env)
+
+
 def pg_restore(stores_info, store_name, dump_dir):
     dump_file = os.path.join(dump_dir, store_name + ".sql")
     options = [
+           "--jobs={}".format(multiprocessing.cpu_count()),
            "--disable-triggers",
            "--no-owner",
            "--data-only",
@@ -209,10 +336,13 @@
 def action_export(stores_info, tarball_file):
     try:
         dump_dir = tempfile.mkdtemp()
+        rename_incosistent_sequences(stores_info)
         dump_files = []
         for store_name in STORES:
             pg_dump(stores_info, store_name, dump_dir)
             dump_files.append(store_name + ".sql")
+        pg_dump_partitioned_schemas(stores_info, dump_dir)
+        dump_files.append("resource-1-partioned.sql")
         cmd = ["tar", "cfz", tarball_file, "-C", dump_dir] + dump_files
         subprocess.check_call(" ".join(cmd), shell=True)
     finally:
@@ -225,7 +355,8 @@
         cmd = ["tar", "xfz", tarball_file, "-C", dump_dir]
         subprocess.check_call(" ".join(cmd), shell=True)
         drop_data(stores_info)
-        fresh_schema(schema_script)
+        fresh_schema(schema_script, stores_info)
+        psql_restore_partitioned_schema(stores_info, dump_dir)
         for store_name in STORES:
             pg_restore(stores_info, store_name, dump_dir)
         set_root_url(stores_info)