Add admin.maintain_tables() sql function.

- It increases the abstraction of the maintenance logic from the Go code
by handling internally the fetch of all metric tables and updating them
one by one and deleting entries for dropped tables.

- To achieve so it uses the same functions used before in the Go code

Author: 0xgouda

internal/sinks/postgres.go

@@ -507,55 +507,11 @@ func (pgw *PostgresWriter) DeleteOldPartitions() {
 // in each metric table in admin.all_distinct_dbname_metrics.
 // This is used to avoid listing the same source multiple times in Grafana dropdowns.
 func (pgw *PostgresWriter) MaintainUniqueSources() {
+	sql := "SELECT admin.maintain_tables() WHERE pg_try_advisory_lock(1571543679778230000)"
 	logger := log.GetLogger(pgw.ctx)
-
-	var lock bool
-	logger.Infof("Trying to get maintenance advisory lock...") // to only have one "maintainer" in case of a "push" setup, as can get costly
-	if err := pgw.sinkDb.QueryRow(pgw.ctx, "SELECT admin.try_get_maintenance_lock();").Scan(&lock); err != nil {
-		logger.Error("Getting maintenance advisory lock failed:", err)
-		return
-	}
-	if !lock {
-		logger.Info("Skipping maintenance as another instance has the advisory lock...")
-		return
-	}
-
-	logger.Info("Updating admin.all_distinct_dbname_metrics listing table...")
-	rows, _ := pgw.sinkDb.Query(pgw.ctx, "SELECT table_name FROM admin.get_top_level_metric_tables()")
-	allDistinctMetricTables, err := pgx.CollectRows(rows, pgx.RowTo[string])
-	if err != nil {
-		logger.Error(err)
-		return
-	}
-
-	// updates listing table entries for a single metric.
-	sqlUpdateListingTable := "SELECT * FROM admin.update_listing_table(metric_table_name => $1);"
-	for i, tableName := range allDistinctMetricTables {
-		metricName := strings.Replace(tableName, "public.", "", 1)
-		allDistinctMetricTables[i] = metricName // later usage in sqlDroppedTables requires no "public." prefix.
-		logger.Debugf("Updating admin.all_distinct_dbname_metrics listing for metric: %s", metricName)
-		var deletedRowsCnt, insertedRowsCnt int
-		err = pgw.sinkDb.QueryRow(pgw.ctx, sqlUpdateListingTable, tableName).Scan(&deletedRowsCnt, &insertedRowsCnt)
-		if err != nil {
-			logger.Errorf("Could not update admin.all_distinct_dbname_metrics listing for metric '%s': %s", metricName, err)
-		}
-		if deletedRowsCnt > 0 {
-			logger.Infof("Removed %d stale entries from admin.all_distinct_dbname_metrics listing table for metric: %s", deletedRowsCnt, metricName)
-		}
-		if insertedRowsCnt > 0 {
-			logger.Infof("Added %d entries to admin.all_distinct_dbname_metrics listing table for metric: %s", insertedRowsCnt, metricName)
-		}
-		time.Sleep(time.Minute)
-	}
-
-	// removes all entries for any non-existing table.
-	sqlRemoveDroppedTables := "SELECT admin.remove_dropped_tables_listing(existing_metrics => $1)"
-	var deletedRowsCnt int
-	err = pgw.sinkDb.QueryRow(pgw.ctx, sqlRemoveDroppedTables, allDistinctMetricTables).Scan(&deletedRowsCnt)
-	if err != nil {
-		logger.Errorf("Could not update admin.all_distinct_dbname_metrics listing table for dropped metric tables: %s", err)
-	} else if deletedRowsCnt > 0 {
-		logger.Infof("Removed %d stale entries for dropped metric tables from admin.all_distinct_dbname_metrics listing table", deletedRowsCnt)
+	logger.Info("Starting maintenance...")
+	if _, err := pgw.sinkDb.Exec(pgw.ctx, sql); err != nil {
+		logger.Error("Maintaining 'admin.all_distinct_dbname_metrics' listing table failed:", err)
 	}
 }
 

internal/sinks/postgres_test.go

@@ -709,16 +709,6 @@ func Test_MaintainUniqueSources_DeleteOldPartitions(t *testing.T) {
 		a.NoError(err)
 		a.Equal(1, numOfEntries)
 
-		message[0].DBName = "test_db_2"
-		pgw.flush(message)
-		// explicitly use `public.*` prefix.
-		_, err = pgw.sinkDb.Exec(pgw.ctx, "SELECT admin.update_listing_table(metric_table_name => 'public.test_metric_1');")
-		a.NoError(err)
-		// another entry should have been added.
-		err = conn.QueryRow(ctx, "SELECT count(*) FROM admin.all_distinct_dbname_metrics;").Scan(&numOfEntries)
-		a.NoError(err)
-		a.Equal(2, numOfEntries)
-
 		_, err = conn.Exec(ctx, "DROP TABLE test_metric_1;")
 		r.NoError(err)
 		// all entries should be deleted

internal/sinks/sql/admin_functions.sql

@@ -140,34 +140,40 @@ BEGIN
 END;
 $SQL$ LANGUAGE plpgsql;
 
-CREATE OR REPLACE FUNCTION admin.try_get_maintenance_lock(
-  OUT have_lock BOOLEAN
-)
+CREATE OR REPLACE FUNCTION admin.maintain_tables()
+RETURNS VOID
 AS 
 $SQL$
+DECLARE
+  rec record;
+  metric_name text;
+  existing_metrics text[];
 BEGIN
-  -- 1571543679778230000 is just a random bigint
-  SELECT pg_try_advisory_lock(1571543679778230000) INTO have_lock;
+  FOR rec IN SELECT * FROM admin.get_top_level_metric_tables()
+  LOOP
+    IF POSITION('public.' IN rec.table_name) = 1 THEN
+      metric_name = SUBSTRING(rec.table_name FROM POSITION('public.' IN rec.table_name) + LENGTH('public.'));
+    ELSE
+      metric_name = rec.table_name;
+    END IF;
+
+    SELECT array_append(existing_metrics, metric_name) INTO existing_metrics;
+
+    EXECUTE FORMAT($$SELECT admin.update_listing_table(metric_name => '%s')$$, metric_name);
+    SELECT PG_SLEEP(60);
+  END LOOP;
+
+  -- Delete entries for dropped tables
+  DELETE FROM admin.all_distinct_dbname_metrics WHERE metric != ALL(existing_metrics) OR existing_metrics IS NULL;
 END;
 $SQL$ LANGUAGE plpgsql;
 
-CREATE OR REPLACE FUNCTION admin.update_listing_table(
-  metric_table_name text,
-  OUT deleted_rows_cnt int,
-  OUT inserted_rows_cnt int
-)
+CREATE OR REPLACE FUNCTION admin.update_listing_table(metric_name text)
+RETURNS VOID
 AS
 $SQL$
-DECLARE
-  metric_name text;
 BEGIN
 
-  IF POSITION('public.' IN metric_table_name) > 0 THEN
-    metric_name = SUBSTRING(metric_table_name FROM POSITION('public.' IN metric_table_name) + LENGTH('public.'));
-  ELSE
-    metric_name = metric_table_name;
-  END IF;
-
   EXECUTE FORMAT(
   $$
     CREATE TEMP TABLE distinct_dbnames AS
@@ -186,30 +192,14 @@ BEGIN
       AND metric = '%s';
   $$, metric_name);
 
-  GET DIAGNOSTICS deleted_rows_cnt = ROW_COUNT;
-
   EXECUTE FORMAT(
   $$
     INSERT INTO admin.all_distinct_dbname_metrics 
       SELECT d.dbname, '%s' FROM distinct_dbnames AS d
       WHERE NOT EXISTS (SELECT * FROM admin.all_distinct_dbname_metrics WHERE dbname = d.dbname AND metric = '%s');
   $$, metric_name, metric_name);
 
-  GET DIAGNOSTICS inserted_rows_cnt = ROW_COUNT;
-
   DROP TABLE distinct_dbnames;
 
 END;
 $SQL$ LANGUAGE plpgsql;
-
-CREATE OR REPLACE FUNCTION admin.remove_dropped_tables_listing(
-  existing_metrics text[],
-  OUT deleted_rows_cnt int
-)
-AS 
-$SQL$
-BEGIN
-  DELETE FROM admin.all_distinct_dbname_metrics WHERE metric != ALL(existing_metrics);
-  GET DIAGNOSTICS deleted_rows_cnt = ROW_COUNT;
-END;
-$SQL$ LANGUAGE plpgsql;