feat: add custom function to get info from pg_stat_replication for bootsrap

Author: casperr0

mamonsu/plugins/pgsql/driver/pool.py

@@ -36,6 +36,19 @@ class Pool(object):
             " from public.pg_buffercache",
             'select size, twice_used, dirty from public.mamonsu_buffer_cache()'
         ),
+        'wal_lag_lsn': (
+            "SELECT application_name, " \
+            " flush_lag, replay_lag, write_lag, " \
+            " pg_wal_lsn_diff(pg_current_wal_lsn(), replay_lsn) AS total_lag " \
+            " FROM pg_stat_replication;",
+            'select public.mamonsu_count_wal_lag_lsn()'
+        ),
+        'xlog_lag_lsn': (
+            "SELECT application_name, " \
+            "pg_xlog_location_diff(pg_current_xlog_location(), replay_location) AS total_lag  " \
+            "FROM pg_stat_replication;",
+            'select public.mamonsu_count_xlog_lag_lsn()'
+        ),
     }
 
     def __init__(self, params=None):
@@ -66,7 +79,7 @@ def server_version(self, db=None):
         if db in self._cache['server_version']['storage']:
             return self._cache['server_version']['storage'][db]
         result = bytes(
-                self.query('show server_version', db)[0][0], 'utf-8')
+            self.query('show server_version', db)[0][0], 'utf-8')
         self._cache['server_version']['storage'][db] = '{0}'.format(
             result.decode('ascii'))
         return self._cache['server_version']['storage'][db]

mamonsu/plugins/pgsql/xlog.py

@@ -54,7 +54,7 @@ def run(self, zbx):
 
             if Pooler.server_version_greater('10.0'):
                 result = Pooler.query(self.query_wal_lsn_diff)
-                result_lags = Pooler.query(self.query_wal_lag_lsn)
+                result_lags = Pooler.run_sql_type('wal_lag_lsn')
                 if result_lags:
                     lags = []
                     for info in result_lags:
@@ -71,7 +71,7 @@ def run(self, zbx):
                     del lags
             else:
                 result = Pooler.query(self.query_xlog_lsn_diff)
-                result_lags = Pooler.query(self.query_xlog_lag_lsn)
+                result_lags = Pooler.run_sql_type('xlog_lag_lsn')
                 if result_lags:
                     lags = []
                     for info in result_lags:

mamonsu/tools/bootstrap/sql.py

@@ -123,6 +123,15 @@
 SELECT COUNT(*) AS count_prepared,
 coalesce (ROUND(MAX(EXTRACT (EPOCH FROM (now() - prepared)))),0)::bigint AS oldest_prepared  
 FROM pg_catalog.pg_prepared_xacts$$ LANGUAGE SQL SECURITY DEFINER;
+
+CREATE OR REPLACE FUNCTION public.mamonsu_count_{3}_lag_lsn()
+RETURNS TABLE(application_name TEXT, {8} total_lag NUMERIC ) AS $$
+SELECT 
+application_name,
+{6}
+pg_{7}_diff(pg_current_{7}(), replay_{9}) AS total_lag 
+FROM pg_stat_replication
+$$ LANGUAGE SQL SECURITY DEFINER;
 """
 
 GrantsOnSchemaSQL = """
@@ -153,4 +162,6 @@
 GRANT EXECUTE ON FUNCTION public.mamonsu_get_connections_states() TO {1};
 
 GRANT EXECUTE ON FUNCTION public.mamonsu_prepared_transaction() TO {1};
+
+GRANT EXECUTE ON FUNCTION public.mamonsu_count_{2}_lag_lsn() TO {1};
 """

mamonsu/tools/bootstrap/start.py

@@ -161,7 +161,12 @@ def run_deploy():
                 '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'
+                                                                      ' 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))