From 1664daa536889c9990f54def9bdeb1520d0ce66f Mon Sep 17 00:00:00 2001 From: Fantix King Date: Mon, 22 Jul 2024 12:49:19 -0400 Subject: [PATCH] func cache: fix dropping extension with scalar type (#7564) Scalar types in extensions may defined `sql_type` thus not using `common.get_backend_name()`. Instead, here I think we should use `types.pg_type_from_scalar()` to cover more cases. Also fixed a flaky test. --- edb/pgsql/delta.py | 42 ++++++++++++++++++++++++------------- tests/test_edgeql_vector.py | 19 +++++++++++++++++ tests/test_server_config.py | 39 ++++++++++++++++++++++------------ 3 files changed, 73 insertions(+), 27 deletions(-) diff --git a/edb/pgsql/delta.py b/edb/pgsql/delta.py index 03031d6b5c0..4ff2a2da2a5 100644 --- a/edb/pgsql/delta.py +++ b/edb/pgsql/delta.py @@ -2838,21 +2838,24 @@ def delete_scalar( cls, scalar: s_scalars.ScalarType, orig_schema: s_schema.Schema ) -> dbops.Command: ops = dbops.CommandGroup() - old_domain_name = common.get_backend_name( - orig_schema, scalar, catenate=False) # The custom scalar types are sometimes included in the function # signatures of query cache functions under QueryCacheMode.PgFunc. # We need to find such functions through pg_depend and evict the cache # before dropping the custom scalar type. - drop_func_cache_sql = textwrap.dedent(f''' - DO $$ - DECLARE - qc RECORD; - BEGIN - FOR qc IN - WITH - types AS ( + pg_type = types.pg_type_from_scalar(orig_schema, scalar) + if len(pg_type) == 1: + types_cte = f''' + SELECT + pt.oid AS oid + FROM + pg_type pt + WHERE + pt.typname = {ql(pg_type[0])} + OR pt.typname = {ql('_' + pg_type[0])}\ + ''' + else: + types_cte = f''' SELECT pt.oid AS oid FROM @@ -2860,11 +2863,20 @@ def delete_scalar( JOIN pg_namespace pn ON pt.typnamespace = pn.oid WHERE - pn.nspname = {ql(old_domain_name[0])} + pn.nspname = {ql(pg_type[0])} AND ( - pt.typname = {ql(old_domain_name[1])} - OR pt.typname = {ql('_' + old_domain_name[1])} - ) + pt.typname = {ql(pg_type[1])} + OR pt.typname = {ql('_' + pg_type[1])} + )\ + ''' + drop_func_cache_sql = textwrap.dedent(f''' + DO $$ + DECLARE + qc RECORD; + BEGIN + FOR qc IN + WITH + types AS ({types_cte} ), class AS ( SELECT @@ -2896,6 +2908,8 @@ class AS ( ''') ops.add_command(dbops.Query(drop_func_cache_sql)) + old_domain_name = common.get_backend_name( + orig_schema, scalar, catenate=False) cond: dbops.Condition if scalar.is_concrete_enum(orig_schema): old_enum_name = old_domain_name diff --git a/tests/test_edgeql_vector.py b/tests/test_edgeql_vector.py index 5d2dd0925f4..775b68c0da0 100644 --- a/tests/test_edgeql_vector.py +++ b/tests/test_edgeql_vector.py @@ -861,3 +861,22 @@ async def test_edgeql_vector_config(self): 'select _current_setting("hnsw.ef_search")', [40], ) + + +class TestEdgeQLVectorExtension(tb.QueryTestCase): + BACKEND_SUPERUSER = True + TRANSACTION_ISOLATION = False + + async def test_edgeql_vector_drop_extension_with_func_cache(self): + await self.con.execute("create extension pgvector") + try: + # Run many times to wait for the func cache creation + for _i in range(64): + await self.con.query( + ''' + select [4.2]; + ''' + ) + finally: + # this should drop the cache function of the query above as well + await self.con.execute("drop extension pgvector") diff --git a/tests/test_server_config.py b/tests/test_server_config.py index 50a89875ae4..8d6a3990272 100644 --- a/tests/test_server_config.py +++ b/tests/test_server_config.py @@ -1434,25 +1434,38 @@ async def test_server_proto_non_transactional_pg_14_7(self): async def test_server_proto_recompile_on_db_config(self): await self.con.execute("create type RecompileOnDBConfig;") try: + # We need the retries here because the 2 `configure database` may + # race with each other and cause temporary inconsistency await self.con.execute(''' configure current database set allow_user_specified_id := true; ''') - await self.con.execute(''' - insert RecompileOnDBConfig { - id := '8c425e34-d1c3-11ee-8c78-8f34556d1111' - }; - ''') + async for tr in self.try_until_succeeds(ignore=edgedb.QueryError): + async with tr: + await self.con.execute(''' + insert RecompileOnDBConfig { + id := '8c425e34-d1c3-11ee-8c78-8f34556d1111' + }; + ''') + await self.con.execute(''' configure current database set allow_user_specified_id := false; ''') - with self.assertRaisesRegex( - edgedb.QueryError, "cannot assign to property 'id'" - ): - await self.con.execute(''' - insert RecompileOnDBConfig { - id := '8c425e34-d1c3-11ee-8c78-8f34556d2222' - }; - ''') + async for tr in self.try_until_succeeds(ignore=AssertionError): + async with tr: + try: + with self.assertRaisesRegex( + edgedb.QueryError, "cannot assign to property 'id'" + ): + await self.con.execute(''' + insert RecompileOnDBConfig { + id := + '8c425e34-d1c3-11ee-8c78-8f34556d2222' + }; + ''') + finally: + await self.con.execute(''' + delete RecompileOnDBConfig; + ''') finally: await self.con.execute(''' configure current database reset allow_user_specified_id;