feat: run boostrap queries in one transaction and queries for extension in separate transaction.

Author: casperr0

mamonsu/tools/bootstrap/sql.py

@@ -4,7 +4,7 @@
 
 """
 
-CreateSchemaSQL = """
+CreateSchemaDefaultSQL = """
 CREATE TABLE IF NOT EXISTS public.mamonsu_config (
   version text,
   inserted_at timestamp DEFAULT NOW()
@@ -90,17 +90,6 @@
 WHERE filename similar to '{2}'
 $$ LANGUAGE SQL SECURITY DEFINER;
 
-CREATE EXTENSION IF NOT EXISTS pg_buffercache;
-
-CREATE OR REPLACE FUNCTION public.mamonsu_buffer_cache()
-RETURNS TABLE(SIZE BIGINT, TWICE_USED BIGINT, DIRTY BIGINT) AS $$
-SELECT
-   SUM(1) * (current_setting('block_size')::int8),
-   SUM(CASE WHEN usagecount > 1 THEN 1 ELSE 0 END) * (current_setting('block_size')::int8),
-   SUM(CASE isdirty WHEN true THEN 1 ELSE 0 END) * (current_setting('block_size')::int8)
-FROM public.pg_buffercache
-$$ LANGUAGE SQL SECURITY DEFINER;
-
 CREATE OR REPLACE FUNCTION public.mamonsu_archive_command_files()
 RETURNS TABLE(COUNT_FILES BIGINT, SIZE_FILES NUMERIC) AS $$
 SELECT count(name) AS COUNT_FILES ,
@@ -133,7 +122,18 @@
 $$ LANGUAGE SQL SECURITY DEFINER;
 """
 
-GrantsOnSchemaSQL = """
+CreateSchemaExtensionSQL = """
+CREATE OR REPLACE FUNCTION public.mamonsu_buffer_cache()
+RETURNS TABLE(SIZE BIGINT, TWICE_USED BIGINT, DIRTY BIGINT) AS $$
+SELECT
+   SUM(1) * (current_setting('block_size')::int8),
+   SUM(CASE WHEN usagecount > 1 THEN 1 ELSE 0 END) * (current_setting('block_size')::int8),
+   SUM(CASE isdirty WHEN true THEN 1 ELSE 0 END) * (current_setting('block_size')::int8)
+FROM public.pg_buffercache
+$$ LANGUAGE SQL SECURITY DEFINER;
+"""
+
+GrantsOnDefaultSchemaSQL = """
 ALTER TABLE public.mamonsu_config OWNER TO {1};
 
 ALTER TABLE public.mamonsu_timestamp_master_{0} OWNER TO {1};
@@ -150,8 +150,6 @@
 
 GRANT EXECUTE ON FUNCTION public.mamonsu_count_{2}_files() TO {1};
 
-GRANT EXECUTE ON FUNCTION public.mamonsu_buffer_cache() TO {1};
-
 GRANT EXECUTE ON FUNCTION public.mamonsu_archive_command_files() TO {1};
 
 GRANT EXECUTE ON FUNCTION public.mamonsu_archive_stat() TO {1};
@@ -164,3 +162,7 @@
 
 GRANT EXECUTE ON FUNCTION public.mamonsu_count_{2}_lag_lsn() TO {1};
 """
+
+GrantsOnExtensionSchemaSQL = """
+GRANT EXECUTE ON FUNCTION public.mamonsu_buffer_cache() TO {1};
+"""

mamonsu/tools/bootstrap/start.py

@@ -10,7 +10,8 @@
 from mamonsu import __version__ as mamonsu_version
 from mamonsu.lib.default_config import DefaultConfig
 from mamonsu.plugins.pgsql.pool import Pooler
-from mamonsu.tools.bootstrap.sql import CreateSchemaSQL, GrantsOnSchemaSQL, QuerySplit
+from mamonsu.tools.bootstrap.sql import CreateSchemaExtensionSQL, CreateSchemaDefaultSQL, \
+    GrantsOnDefaultSchemaSQL, GrantsOnExtensionSchemaSQL, QuerySplit
 
 
 class Args(DefaultConfig):
@@ -145,6 +146,37 @@ def test_db(self, host_pre):
         return True
 
 
+def fill_query_params(queries):
+    formatted_queries = ""
+    for sql in queries.format(
+            mamonsu_version,
+            mamonsu_version.replace('.', '_'),
+            '[0-9A-F]{24}',
+            'wal' if Pooler.server_version_greater('10.0') else 'xlog',
+            'wal_lsn' if Pooler.server_version_greater('10.0') else 'xlog_location',
+            'waiting' if Pooler.server_version_less('9.6.0') else 'case when wait_event_type is null then false '
+                                                                  ' else true end  as waiting',
+            'flush_lag, replay_lag, write_lag,' if Pooler.server_version_greater('10.0') else '',
+            'wal_lsn' if Pooler.server_version_greater('10.0') else 'xlog_location',
+            'flush_lag INTERVAL, replay_lag INTERVAL, write_lag INTERVAL,' if Pooler.server_version_greater('10.0')
+            else '',
+            'lsn' if Pooler.server_version_greater('10.0') else 'location'
+    ).split(QuerySplit):
+        formatted_queries += sql
+    return formatted_queries
+
+
+def fill_grant_params(queries, args):
+    formatted_grants_queries = ""
+    for sql in queries.format(
+            mamonsu_version.replace('.', '_'),
+            args.args.mamonsu_username,
+            'wal' if Pooler.server_version_greater('10.0') else 'xlog'
+    ).split(QuerySplit):
+        formatted_grants_queries += sql
+    return formatted_grants_queries
+
+
 def run_deploy():
     args = Args()
 
@@ -159,37 +191,42 @@ def run_deploy():
         sys.exit(1)
 
     try:
-        for sql in CreateSchemaSQL.format(
-                mamonsu_version,
-                mamonsu_version.replace('.', '_'),
-                '[0-9A-F]{24}',
-                'wal' if Pooler.server_version_greater('10.0') else 'xlog',
-                'wal_lsn' if Pooler.server_version_greater('10.0') else 'xlog_location',
-                'waiting' if Pooler.server_version_less('9.6.0') else 'case when wait_event_type is null then false '
-                                                                      ' else true end  as waiting',
-                'flush_lag, replay_lag, write_lag,' if Pooler.server_version_greater('10.0') else '',
-                'wal_lsn' if Pooler.server_version_greater('10.0') else 'xlog_location',
-                'flush_lag INTERVAL, replay_lag INTERVAL, write_lag INTERVAL,' if Pooler.server_version_greater('10.0')
-                else '',
-                'lsn' if Pooler.server_version_greater('10.0') else 'location'
-        ).split(QuerySplit):
-            if args.args.verbose:
-                sys.stdout.write("\nExecuting query:\n{0}\n".format(sql))
-            Pooler.query(sql)
+        bootstrap_queries = fill_query_params(CreateSchemaDefaultSQL)
+        Pooler.query(bootstrap_queries)
+        #  if args.args.verbose:
+        #     sys.stdout.write("\nExecuting query:\n{0}\n".format(sql))
     except Exception as e:
-        sys.stderr.write("Query:\n{0}\nerror: {1}\n".format(sql, e))
+        sys.stderr.write("Bootstrap execution have exited with an error: {0}\n".format(e))
         sys.exit(2)
+
+    try:
+        bootstrap_extension_queries = fill_query_params(CreateSchemaExtensionSQL)
+        Pooler.query(bootstrap_extension_queries)
+        #  if args.args.verbose:
+        #     sys.stdout.write("\nExecuting query:\n{0}\n".format(sql))
+    except Exception as e:
+        sys.stderr.write(
+            "Bootstrap failed to create the function which required pg_buffercache extension.\n"
+            "Error: {0}\n".format(e))
+        sys.stderr.write("Please install pg_buffercache extension and rerun bootstrap "
+                         "if you want to get metrics from pg_buffercache view. \n")
+
     try:
-        for sql in GrantsOnSchemaSQL.format(
-                mamonsu_version.replace('.', '_'),
-                args.args.mamonsu_username,
-                'wal' if Pooler.server_version_greater('10.0') else 'xlog'
-        ).split(QuerySplit):
-            if args.args.verbose:
-                sys.stdout.write("\nExecuting query:\n{0}\n".format(sql))
-            Pooler.query(sql)
+        bootstrap_grant_queries = fill_grant_params(GrantsOnDefaultSchemaSQL, args)
+        Pooler.query(bootstrap_grant_queries)
+
     except Exception as e:
-        sys.stderr.write("Query:\n{0}\nerror: {1}\n".format(sql, e))
+        sys.stderr.write("Error: \n {0}\n".format(e))
+        sys.stderr.write("Please check user rights and rerun bootstrap: \n {0}\n".format(e))
         sys.exit(2)
 
+    try:
+        bootstrap_grant_extension_queries = fill_grant_params(GrantsOnExtensionSchemaSQL, args)
+        Pooler.query(bootstrap_grant_extension_queries)
+
+    except Exception as e:
+        sys.stderr.write("Bootstrap failed to grant execution command to "
+                         "the function which required pg_buffercache extension.\n")
+        sys.stderr.write("Error: \n {0}\n".format(e))
+
     sys.stdout.write("Bootstrap successfully completed\n")