Find all sequences for light dump including for IDENTITY columns

okfde · Apr 25, 2024 · f596e20 · f596e20
1 parent de9034d
commit f596e20
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/scripts/create_light_sql.py b/scripts/create_light_sql.py
@@ -12,6 +12,9 @@
 SEQUENCE_RE = re.compile(
     r"ALTER TABLE ONLY (?P<table>public\.\w+) ALTER COLUMN (?P<column>\w+) SET DEFAULT nextval\('(?P<sequence>public\.\w+)'::regclass\);"
 )
+SEQUENCE_IDENTITY_RE = re.compile(
+    r"ALTER TABLE (?P<table>public\.\w+) ALTER COLUMN (?P<column>\w+) ADD GENERATED BY DEFAULT AS IDENTITY \(\s*SEQUENCE NAME (?P<sequence>public\.\w+)"
+)
 
 SQL_TEMPLATE = """
 ALTER TABLE ONLY {table}
@@ -233,7 +236,7 @@ def generate_copy_script(
 
     outfile.write(f"psql {target_connection} {target_db} < constraints.sql\n")
 
-    matches = SEQUENCE_RE.findall(schema)
+    matches = SEQUENCE_RE.findall(schema) + SEQUENCE_IDENTITY_RE.findall(schema)
     for match in matches:
         table, column, sequence = match
         sql = f"""SELECT SETVAL('{sequence}', COALESCE(MAX({column}), 1) ) FROM {table};"""