Make age extension usable from shared_preload_libraries

When AGE is loaded via shared_preload_libraries, its hooks
(post_parse_analyze, set_rel_pathlist, object_access) are active
before CREATE EXTENSION age is run. This causes errors when
non-Cypher queries trigger those hooks and ag_catalog does not
yet exist.

Changes:

- Add is_age_extension_exist() with a relcache callback cache so that
  checking pg_extension is not repeated on every hook invocation.
- Guard post_parse_analyze, set_rel_pathlist, and object_access hooks
  with is_age_extension_exist() so they become no-ops when the
  extension is not installed.
- Refactor ag_ProcessUtility_hook to detect CREATE/DROP EXTENSION age
  and broadcast a relcache invalidation via
  CacheInvalidateRelcacheByRelid(ExtensionRelationId) so other
  backends update their cached extension state.
- Wrap DROP EXTENSION processing in PG_TRY/PG_CATCH to restore
  object_access_hook if the drop fails (e.g. dependent objects).
- Skip _PG_init during pg_upgrade (IsBinaryUpgrade) to avoid hook
  registration when the binary-upgrade machinery is running.
- Add regression tests that verify hooks do not error when ag_catalog
  schema is absent.

Author: serdarmumcu

regress/expected/drop.out

@@ -40,6 +40,12 @@ SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname = 'ag_catalog';
 -----------
 (0 rows)
 
+-- When ag_catalog is missing extension hooks shouldn't fail with the
+-- ERROR  schema "ag_catalog" does not exist.
+-- It might happen when 'age' is loaded but extension isn't created yet.
+DROP SCHEMA ag_catalog;
+CREATE SCHEMA _regress_drop;
+DROP SCHEMA _regress_drop; -- shouldn't produce the ERROR
 -- Recreate the extension and validate we can recreate a graph
 CREATE EXTENSION age;
 SELECT create_graph('drop');

regress/sql/drop.sql

@@ -28,6 +28,13 @@ SELECT nspname FROM pg_catalog.pg_namespace WHERE nspname = 'drop';
 
 SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname = 'ag_catalog';
 
+-- When ag_catalog is missing extension hooks shouldn't fail with the
+-- ERROR  schema "ag_catalog" does not exist.
+-- It might happen when 'age' is loaded but extension isn't created yet.
+DROP SCHEMA ag_catalog;
+CREATE SCHEMA _regress_drop;
+DROP SCHEMA _regress_drop; -- shouldn't produce the ERROR
+
 -- Recreate the extension and validate we can recreate a graph
 CREATE EXTENSION age;
 

src/backend/age.c

@@ -17,6 +17,9 @@
  * under the License.
  */
 
+#include "postgres.h"
+#include "miscadmin.h"
+
 #include "catalog/ag_catalog.h"
 #include "nodes/ag_nodes.h"
 #include "optimizer/cypher_paths.h"
@@ -25,7 +28,6 @@
 #include "utils/age_global_graph.h"
 
 #if PG_VERSION_NUM < 170000
-#include "miscadmin.h"
 
 /* saved hook pointers for PG < 17 shmem path */
 static shmem_request_hook_type prev_shmem_request_hook = NULL;
@@ -56,6 +58,9 @@ void _PG_init(void);
 
 void _PG_init(void)
 {
+    if (IsBinaryUpgrade)
+        return;
+
     register_ag_nodes();
     set_rel_pathlist_init();
     object_access_hook_init();

src/backend/catalog/ag_catalog.c

@@ -19,21 +19,27 @@
 
 #include "postgres.h"
 
+#include "access/xact.h"
 #include "catalog/dependency.h"
 #include "catalog/namespace.h"
 #include "catalog/objectaccess.h"
 #include "catalog/pg_class_d.h"
+#include "catalog/pg_extension_d.h"
 #include "catalog/pg_namespace_d.h"
 #include "commands/defrem.h"
+#include "commands/extension.h"
 #include "nodes/parsenodes.h"
 #include "tcop/utility.h"
+#include "utils/inval.h"
 #include "utils/lsyscache.h"
 
 #include "catalog/ag_graph.h"
 #include "catalog/ag_label.h"
 #include "utils/ag_cache.h"
 #include "utils/age_global_graph.h"
 
+static bool extension_cache_is_valid = false;
+static bool age_extension_exists = false;
 static object_access_hook_type prev_object_access_hook;
 static ProcessUtility_hook_type prev_process_utility_hook;
 static bool prev_object_hook_is_set;
@@ -45,8 +51,44 @@ void ag_ProcessUtility_hook(PlannedStmt *pstmt, const char *queryString, bool re
                             QueryEnvironment *queryEnv, DestReceiver *dest,
                             QueryCompletion *qc);
 
-static bool is_age_drop(PlannedStmt *pstmt);
-static void drop_age_extension(DropStmt *stmt);
+static bool is_age_drop(DropStmt *drop_stmt);
+
+static void
+invalidate_extension_cache_callback(Datum argument, Oid relationId)
+{
+    if (relationId == ExtensionRelationId)
+        extension_cache_is_valid = false;
+}
+
+/*
+ * We don't want most of hooks to do anything if the "age" extension isn't
+ * created. However, scanning pg_extension is a costly operation, therefore we
+ * implement a caching mechanism and reset it with the help of the relcache
+ * callback mechanism.
+ *
+ * Please also see ag_ProcessUtility_hook() function for more details.
+ */
+bool
+is_age_extension_exist(void)
+{
+    static bool callback_registered = false;
+
+    if (extension_cache_is_valid)
+        return age_extension_exists;
+
+    if (!callback_registered)
+    {
+        CacheRegisterRelcacheCallback(invalidate_extension_cache_callback,
+                                      (Datum) 0);
+        callback_registered = true;
+    }
+
+    age_extension_exists = OidIsValid(get_extension_oid("age", true));
+
+    extension_cache_is_valid = true;
+
+    return age_extension_exists;
+}
 
 void object_access_hook_init(void)
 {
@@ -86,50 +128,97 @@ void process_utility_hook_fini(void)
  * information in the indexes and tables being dropped. To prevent an error
  * from being thrown, we need to disable the object_access_hook before dropping
  * the extension.
+ *
+ * Besides that, we want to notify other backends about the fact that "age"
+ * extension was probably created/dropped so that they can enable/disable
+ * hooks.
  */
 void ag_ProcessUtility_hook(PlannedStmt *pstmt, const char *queryString,
                             bool readOnlyTree, ProcessUtilityContext context,
                             ParamListInfo params, QueryEnvironment *queryEnv,
                             DestReceiver *dest, QueryCompletion *qc)
 {
-    if (is_age_drop(pstmt))
-    {
-        drop_age_extension((DropStmt *)pstmt->utilityStmt);
-    }
-    else
+    bool creating_age = false;
+    bool dropping_age = false;
+
+    if (!IsAbortedTransactionBlockState())
     {
-        /*
-         * Check for TRUNCATE on graph label tables. If any truncated
-         * table is a graph label table, increment the version counter
-         * for that graph to invalidate VLE caches. We do this before
-         * the truncate executes so the cache is invalidated regardless.
-         */
-        if (IsA(pstmt->utilityStmt, TruncateStmt))
+        Node *parsetree = pstmt->utilityStmt;
+
+        switch (nodeTag(parsetree))
         {
-            TruncateStmt *tstmt = (TruncateStmt *) pstmt->utilityStmt;
-            ListCell *lc;
+            case T_CreateExtensionStmt:
+                {
+                    CreateExtensionStmt *stmt =
+                        (CreateExtensionStmt *) parsetree;
+                    creating_age = strcmp(stmt->extname, "age") == 0;
+                }
+                break;
+            case T_DropStmt:
+                {
+                    DropStmt *stmt = (DropStmt *) parsetree;
 
-            foreach(lc, tstmt->relations)
-            {
-                RangeVar *rv = (RangeVar *) lfirst(lc);
-                Oid rel_oid = RangeVarGetRelid(rv, AccessShareLock, true);
+                    if (stmt->removeType != OBJECT_EXTENSION)
+                        break;
 
-                if (OidIsValid(rel_oid))
-                {
-                    Oid graph_oid = get_graph_oid_for_table(rel_oid);
+                    if (!is_age_drop(stmt))
+                        break;
 
-                    if (OidIsValid(graph_oid))
+                    dropping_age = true;
+                }
+                break;
+            case T_TruncateStmt:
+                {
+                    /*
+                     * Check for TRUNCATE on graph label tables. If any
+                     * truncated table is a graph label table, increment the
+                     * version counter for that graph to invalidate VLE caches.
+                     * We do this before the truncate executes so the cache is
+                     * invalidated regardless.
+                     */
+                    TruncateStmt *tstmt = (TruncateStmt *) parsetree;
+                    ListCell *lc;
+
+                    foreach(lc, tstmt->relations)
                     {
-                        increment_graph_version(graph_oid);
+                        RangeVar *rv = (RangeVar *) lfirst(lc);
+                        Oid rel_oid = RangeVarGetRelid(rv, AccessShareLock,
+                                                       true);
+
+                        if (OidIsValid(rel_oid))
+                        {
+                            Oid graph_oid =
+                                get_graph_oid_for_table(rel_oid);
+
+                            if (OidIsValid(graph_oid))
+                            {
+                                increment_graph_version(graph_oid);
+                            }
+                        }
                     }
                 }
-            }
+                break;
+            default:
+                break;
         }
+    }
 
+    if (dropping_age)
+    {
+        /* Remove all graphs */
+        drop_graphs(get_graphnames());
+
+        /* Remove the object access hook */
+        object_access_hook_fini();
+    }
+
+    PG_TRY();
+    {
         if (prev_process_utility_hook)
         {
             (*prev_process_utility_hook) (pstmt, queryString, readOnlyTree,
-                                          context, params, queryEnv, dest, qc);
+                                          context, params, queryEnv, dest,
+                                          qc);
         }
         else
         {
@@ -141,39 +230,46 @@ void ag_ProcessUtility_hook(PlannedStmt *pstmt, const char *queryString,
                                     params, queryEnv, dest, qc);
         }
     }
-}
-
-static void drop_age_extension(DropStmt *stmt)
-{
-    /* Remove all graphs */
-    drop_graphs(get_graphnames());
+    PG_CATCH();
+    {
+        if (dropping_age)
+        {
+            /*
+             * We have to restore the disabled object_access_hook if
+             * DROP EXTENSION age failed.
+             */
+            object_access_hook_init();
+        }
+        PG_RE_THROW();
+    }
+    PG_END_TRY();
 
-    /* Remove the object access hook */
-    object_access_hook_fini();
+    if (dropping_age)
+    {
+        /* reset global variables for OIDs */
+        clear_global_Oids_AGTYPE();
+        clear_global_Oids_GRAPHID();
+        clear_global_Oids_VERTEX_EDGE();
 
-    /*
-     * Run Postgres' logic to perform the remaining work to drop the
-     * extension.
-     */
-    RemoveObjects(stmt);
+        /* Restore the object access hook */
+        object_access_hook_init();
+    }
 
-    /* reset global variables for OIDs */
-    clear_global_Oids_AGTYPE();
-    clear_global_Oids_GRAPHID();
-    clear_global_Oids_VERTEX_EDGE();
+    if (creating_age || dropping_age)
+    {
+        /* Notify all backends that pg_extension was modified. */
+        CacheInvalidateRelcacheByRelid(ExtensionRelationId);
+    }
 }
 
 /* Check to see if the Utility Command is to drop the AGE Extension. */
-static bool is_age_drop(PlannedStmt *pstmt)
+static bool is_age_drop(DropStmt *drop_stmt)
 {
     ListCell *lc;
-    DropStmt *drop_stmt;
 
-    if (!IsA(pstmt->utilityStmt, DropStmt))
+    if (!is_age_extension_exist())
         return false;
 
-    drop_stmt = (DropStmt *)pstmt->utilityStmt;
-
     foreach(lc, drop_stmt->objects)
     {
         Node *obj = lfirst(lc);
@@ -183,7 +279,7 @@ static bool is_age_drop(PlannedStmt *pstmt)
             String *val = (String *)obj;
             char *str = val->sval;
 
-            if (!pg_strcasecmp(str, "age"))
+            if (strcmp(str, "age") == 0)
                 return true;
         }
     }
@@ -205,15 +301,11 @@ static void object_access(ObjectAccessType access, Oid class_id, Oid object_id,
     if (prev_object_access_hook)
         prev_object_access_hook(access, class_id, object_id, sub_id, arg);
 
-    /* We are interested in DROP SCHEMA and DROP TABLE commands. */
-    if (access != OAT_DROP)
+    if (!is_age_extension_exist())
         return;
 
-    /*
-     * Age might be installed into shared_preload_libraries before extension is
-     * created. In this case we must bail out from this hook.
-     */
-    if (!OidIsValid(get_namespace_oid("ag_catalog", true)))
+    /* We are interested in DROP SCHEMA and DROP TABLE commands. */
+    if (access != OAT_DROP)
         return;
 
     drop_arg = arg;

src/backend/optimizer/cypher_paths.c

@@ -19,6 +19,7 @@
 
 #include "postgres.h"
 
+#include "catalog/ag_catalog.h"
 #include "optimizer/pathnode.h"
 #include "optimizer/paths.h"
 
@@ -66,6 +67,9 @@ static void set_rel_pathlist(PlannerInfo *root, RelOptInfo *rel, Index rti,
     if (prev_set_rel_pathlist_hook)
         prev_set_rel_pathlist_hook(root, rel, rti, rte);
 
+    if (!is_age_extension_exist())
+        return;
+
     switch (get_cypher_clause_kind(rte))
     {
     case CYPHER_CLAUSE_CREATE:

src/backend/parser/cypher_analyze.c

@@ -19,6 +19,7 @@
 
 #include "postgres.h"
 
+#include "catalog/ag_catalog.h"
 #include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
 #include "parser/analyze.h"
@@ -86,6 +87,9 @@ static void post_parse_analyze(ParseState *pstate, Query *query, JumbleState *js
         prev_post_parse_analyze_hook(pstate, query, jstate);
     }
 
+    if (!is_age_extension_exist())
+        return;
+
     /*
      * extra_node is set in the parsing stage to keep track of EXPLAIN.
      * So it needs to be set to NULL prior to any cypher parsing.

src/include/catalog/ag_catalog.h

@@ -24,6 +24,8 @@
 
 #include "utils/agtype.h"
 
+bool is_age_extension_exist(void);
+
 void object_access_hook_init(void);
 void object_access_hook_fini(void);