Skip to content

Instantly share code, notes, and snippets.

@MarkusH
Last active September 12, 2024 16:40
Show Gist options
  • Save MarkusH/4c03d0867ba37b6b89379bced12cb3ae to your computer and use it in GitHub Desktop.
Save MarkusH/4c03d0867ba37b6b89379bced12cb3ae to your computer and use it in GitHub Desktop.
A script that outputs the raw SQL to convert Django's serial columns to identity columns.
WITH tab AS (
SELECT
a.attrelid::regclass::text AS t,
a.attname AS c,
pg_get_serial_sequence(a.attrelid::regclass::text, a.attname) AS s,
nextval(pg_get_serial_sequence(a.attrelid::regclass::text, a.attname)) AS v
FROM pg_attribute a
JOIN pg_class c ON c.oid = a.attrelid
WHERE
a.attrelid::regclass::text LIKE '%'
AND c.relkind IN ('r', 'p') /* regular and partitioned tables */
AND a.attnum > 0
AND NOT a.attisdropped
AND a.atttypid = ANY ('{int,int8,int2}'::regtype[])
AND EXISTS (
SELECT FROM pg_attrdef ad
WHERE
ad.adrelid = a.attrelid
AND ad.adnum = a.attnum
AND (
pg_get_expr(ad.adbin, ad.adrelid)
=
'nextval('''
|| (
pg_get_serial_sequence(a.attrelid::regclass::text, a.attname)
)::regclass
|| '''::regclass)'
)
)
ORDER BY a.attnum
)
SELECT
'BEGIN; '
|| 'LOCK TABLE ' || quote_ident(t) || ' IN ACCESS EXCLUSIVE MODE; '
|| 'ALTER TABLE ' || quote_ident(t) || ' ALTER COLUMN ' || quote_ident(c) || ' DROP DEFAULT; '
|| 'DROP SEQUENCE ' || s || '; '
|| 'ALTER TABLE ' || quote_ident(t) || ' ALTER ' || quote_ident(c) || ' ADD GENERATED BY DEFAULT AS IDENTITY (RESTART ' || v || '); '
|| 'COMMIT;'
FROM tab;
@crgwbr
Copy link

crgwbr commented Sep 12, 2024

Unless I'm mistaken, this appears to get the nextval of the old sequence at the the SQL is generated, rather than at the time the generated SQL itself it run. That means that if any rows are inserted in-between the time the SQL is generated and when it's run, the new sequence value will be incorrect. Shouldn't the call to nextval be part of the generated SQL (and inside the exclusive table lock)? That'd also mean re-ordering commands so that the DROP SEQUENCE statement is after the last ALTER TABLE.

Update: Here's an updated version of this script that fixes the issue w/ nextval timing.

WITH tab AS (
    SELECT
        a.attrelid::regclass::text AS table_name,
        a.attname AS pk_col,
        pg_get_serial_sequence(a.attrelid::regclass::text, a.attname) AS seq_name
    FROM pg_attribute a
    JOIN pg_class c ON c.oid = a.attrelid
    WHERE
        a.attrelid::regclass::text LIKE '%'
        AND c.relkind IN ('r', 'p')  /* regular and partitioned tables */
        AND a.attnum > 0
        AND NOT a.attisdropped
        AND a.atttypid = ANY ('{int,int8,int2}'::regtype[])
        AND EXISTS (
            SELECT FROM pg_attrdef ad
            WHERE
                ad.adrelid = a.attrelid
                AND ad.adnum = a.attnum
                AND (
                    pg_get_expr(ad.adbin, ad.adrelid)
                    =
                    'nextval('''
                    || (
                        pg_get_serial_sequence(a.attrelid::regclass::text, a.attname)
                    )::regclass
                    || '''::regclass)'
                )
        )
    ORDER BY a.attnum
)
SELECT
    'BEGIN; '
    || 'DO $$ DECLARE nextSeqVal bigint; '
    || 'BEGIN '
    || 'LOCK TABLE ' || quote_ident(table_name) || ' IN ACCESS EXCLUSIVE MODE; '
    || 'ALTER TABLE ' || quote_ident(table_name) || ' ALTER COLUMN ' || quote_ident(pk_col) || ' DROP DEFAULT; '
    || 'nextSeqVal := nextval(''' || seq_name || '''); '
    || 'DROP SEQUENCE ' || seq_name || '; '
    || 'EXECUTE format(''ALTER TABLE ' || quote_ident(table_name) || ' ALTER ' || quote_ident(pk_col) || ' ADD GENERATED BY DEFAULT AS IDENTITY (RESTART %s);'', nextSeqVal); '
    || 'END $$; '
    || 'COMMIT;'
FROM tab;

As well, here's an example bash script of how you can generate and run the migration code on a batch of DBs:

#!/usr/bin/env bash

set -euxo pipefail

DB_URL="postgres://foo:[email protected]:5432"

for DB in foo bar baz
do
    psql \
        -d "${DB_URL}/${DB}" \
        -f serial-to-identity.sql \
        | grep 'BEGIN;' \
        | awk '{$1=$1};1' \
            > "migrate-${DB}.sql";
    psql -d "${DB_URL}/${DB}" \
        -f "migrate-${DB}.sql";
    echo "Finished ${DB}";
done

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment