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

@@ -2838,33 +2838,45 @@ 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
                             pg_type pt
                             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

tests/test_edgeql_vector.py

@@ -861,3 +861,22 @@ async def test_edgeql_vector_config(self):
             'select <int64>_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 <ext::pgvector::vector>[4.2];
+                    '''
+                )
+        finally:
+            # this should drop the cache function of the query above as well
+            await self.con.execute("drop extension pgvector")

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 := <uuid>'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 := <uuid>'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 := <uuid>'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 := <uuid>
+                                    '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;