fix(api): skip uuidv7() creation when PostgreSQL 18 provides it natively (#36998)

Co-authored-by: Yunlu Wen <yunlu.wen@dify.ai>
This commit is contained in:
Muhammad Zeeshan 2026-06-30 07:34:27 +05:00 committed by GitHub
parent 49a92f096f
commit 995ba6b00e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 159 additions and 27 deletions

View File

@ -45,34 +45,52 @@ def upgrade():
# PostgreSQL 18's `uuidv7` function. This capability is rarely needed in practice, as IDs can be
# generated and controlled within the application layer.
conn = op.get_bind()
if _is_pg(conn):
# PostgreSQL: Create uuidv7 functions
op.execute(sa.text(r"""
/* Main function to generate a uuidv7 value with millisecond precision */
CREATE FUNCTION uuidv7() RETURNS uuid
AS
$$
-- Replace the first 48 bits of a uuidv4 with the current
-- number of milliseconds since 1970-01-01 UTC
-- and set the "ver" field to 7 by setting additional bits
SELECT encode(
set_bit(
set_bit(
overlay(uuid_send(gen_random_uuid()) placing
substring(int8send((extract(epoch from clock_timestamp()) * 1000)::bigint) from
3)
from 1 for 6),
52, 1),
53, 1), 'hex')::uuid;
$$ LANGUAGE SQL VOLATILE PARALLEL SAFE;
COMMENT ON FUNCTION uuidv7 IS
'Generate a uuid-v7 value with a 48-bit timestamp (millisecond precision) and 74 bits of randomness';
if _is_pg(conn):
# PostgreSQL: Create uuidv7 functions.
# PostgreSQL 18 ships a native pg_catalog.uuidv7(), so only create our own
# implementation when the server does not already provide one. Otherwise the
# CREATE FUNCTION below and the unqualified COMMENT statement collide with the
# built-in and the migration fails.
#
# The existence check is done server-side via a DO block rather than
# conn.execute().scalar() because the latter returns None in offline
# migration mode (no real database connection), causing an AttributeError.
op.execute(sa.text(r"""
DO $do$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_proc p
JOIN pg_namespace n ON p.pronamespace = n.oid
WHERE p.proname = 'uuidv7' AND n.nspname = 'pg_catalog'
) THEN
/* Main function to generate a uuidv7 value with millisecond precision */
CREATE FUNCTION public.uuidv7() RETURNS uuid
AS
$func$
-- Replace the first 48 bits of a uuidv4 with the current
-- number of milliseconds since 1970-01-01 UTC
-- and set the "ver" field to 7 by setting additional bits
SELECT encode(
set_bit(
set_bit(
overlay(uuid_send(gen_random_uuid()) placing
substring(int8send((extract(epoch from clock_timestamp()) * 1000)::bigint) from
3)
from 1 for 6),
52, 1),
53, 1), 'hex')::uuid;
$func$ LANGUAGE SQL VOLATILE PARALLEL SAFE;
COMMENT ON FUNCTION public.uuidv7 IS
'Generate a uuid-v7 value with a 48-bit timestamp (millisecond precision) and 74 bits of randomness';
END IF;
END
$do$;
"""))
op.execute(sa.text(r"""
CREATE FUNCTION uuidv7_boundary(timestamptz) RETURNS uuid
CREATE FUNCTION public.uuidv7_boundary(timestamptz) RETURNS uuid
AS
$$
/* uuid fields: version=0b0111, variant=0b10 */
@ -83,7 +101,7 @@ SELECT encode(
'hex')::uuid;
$$ LANGUAGE SQL STABLE STRICT PARALLEL SAFE;
COMMENT ON FUNCTION uuidv7_boundary(timestamptz) IS
COMMENT ON FUNCTION public.uuidv7_boundary(timestamptz) IS
'Generate a non-random uuidv7 with the given timestamp (first 48 bits) and all random bits to 0. As the smallest possible uuidv7 for that timestamp, it may be used as a boundary for partitions.';
"""
))
@ -95,7 +113,10 @@ def downgrade():
conn = op.get_bind()
if _is_pg(conn):
op.execute(sa.text("DROP FUNCTION uuidv7"))
op.execute(sa.text("DROP FUNCTION uuidv7_boundary"))
# IF EXISTS keeps the downgrade a no-op on PostgreSQL 18, where the native
# pg_catalog.uuidv7() was kept and no public.uuidv7() was created. Scoping the
# drop to the public schema avoids touching the built-in.
op.execute(sa.text("DROP FUNCTION IF EXISTS public.uuidv7()"))
op.execute(sa.text("DROP FUNCTION IF EXISTS public.uuidv7_boundary(timestamptz)"))
else:
pass

View File

@ -0,0 +1,111 @@
"""Tests for the uuidv7 SQL migration's PostgreSQL 18 compatibility guard.
The migration file name is not a valid Python identifier (it starts with a date and
contains hyphens), so it is loaded directly from its path. The ``models`` import at the
top of the migration is stubbed because the migration never uses it during
``upgrade()``/``downgrade()`` and pulling in the real package would require a full app
context.
"""
import importlib.util
import sys
import types
from pathlib import Path
from unittest import mock
import pytest
MIGRATION_PATH = (
Path(__file__).resolve().parents[3]
/ "migrations"
/ "versions"
/ "2025_07_02_2332-1c9ba48be8e4_add_uuidv7_function_in_sql.py"
)
def _load_migration():
# The migration does `import models as models` but never references it, so a stub is
# enough and keeps the test free of any database/app configuration.
sys.modules.setdefault("models", types.ModuleType("models"))
spec = importlib.util.spec_from_file_location("uuidv7_pg18_migration_under_test", MIGRATION_PATH)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
def _make_bind(dialect_name):
bind = mock.MagicMock()
bind.dialect.name = dialect_name
return bind
def _executed_sql(fake_op):
return [str(call.args[0]) for call in fake_op.execute.call_args_list]
@pytest.fixture
def migration():
return _load_migration()
def test_upgrade_creates_both_functions_when_native_uuidv7_absent(migration):
# PostgreSQL 13 to 17: no native pg_catalog.uuidv7(), so both functions are created.
# The DO block contains the CREATE FUNCTION guarded by IF NOT EXISTS, and
# uuidv7_boundary is created unconditionally.
bind = _make_bind("postgresql")
with mock.patch.object(migration, "op") as fake_op:
fake_op.get_bind.return_value = bind
migration.upgrade()
sql = _executed_sql(fake_op)
assert any("CREATE FUNCTION public.uuidv7()" in stmt for stmt in sql)
assert any("CREATE FUNCTION public.uuidv7_boundary(timestamptz)" in stmt for stmt in sql)
def test_upgrade_skips_uuidv7_but_keeps_boundary_when_native_present(migration):
# PostgreSQL 18: native pg_catalog.uuidv7() exists, so the DO block must guard
# the CREATE FUNCTION with an IF NOT EXISTS check against pg_catalog.
# uuidv7_boundary is still missing and has to be created unconditionally.
bind = _make_bind("postgresql")
with mock.patch.object(migration, "op") as fake_op:
fake_op.get_bind.return_value = bind
migration.upgrade()
sql = _executed_sql(fake_op)
# The DO block must contain the pg_catalog existence check.
do_block = next((stmt for stmt in sql if "DO $do$" in stmt), None)
assert do_block is not None
assert "pg_catalog" in do_block
assert "uuidv7" in do_block
assert "IF NOT EXISTS" in do_block
# uuidv7_boundary is always created (not guarded by the DO block).
assert any("CREATE FUNCTION public.uuidv7_boundary(timestamptz)" in stmt for stmt in sql)
def test_upgrade_is_noop_on_non_postgres(migration):
bind = _make_bind("sqlite")
with mock.patch.object(migration, "op") as fake_op:
fake_op.get_bind.return_value = bind
migration.upgrade()
fake_op.execute.assert_not_called()
def test_downgrade_uses_if_exists_and_public_schema(migration):
bind = _make_bind("postgresql")
with mock.patch.object(migration, "op") as fake_op:
fake_op.get_bind.return_value = bind
migration.downgrade()
sql = _executed_sql(fake_op)
assert "DROP FUNCTION IF EXISTS public.uuidv7()" in sql
assert "DROP FUNCTION IF EXISTS public.uuidv7_boundary(timestamptz)" in sql
def test_downgrade_is_noop_on_non_postgres(migration):
bind = _make_bind("sqlite")
with mock.patch.object(migration, "op") as fake_op:
fake_op.get_bind.return_value = bind
migration.downgrade()
fake_op.execute.assert_not_called()