feat(sql): Speed up LatestForecast, enable pg_cron (#113)

Author: devsjc

Makefile

@@ -197,7 +197,7 @@ gen.proto.python: ${PROTOC}
 .PHONY: run.db # Run an instance of Postgres with the required extensions
 run.db:
 	docker build -f internal/server/postgres/infra/Containerfile internal/server/postgres/infra -t data-platform-pgdb:local
-	docker run --rm -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres -e POSTGRES_DB=postgres -p "5400:5432" data-platform-pgdb:local
+	docker run --rm -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres -e POSTGRES_DB=postgres -p "5400:5432" data-platform-pgdb:local postgres -c 'shared_preload_libraries=pg_cron' -c 'cron.database_name=postgres'
 
 .PHONY: run.notebook # Run a python notebook to inspect the API
 run.notebook: gen.proto.python

internal/server/postgres/infra/Containerfile

@@ -6,7 +6,6 @@ RUN apt-get update \
     && apt-get install -y --no-install-recommends \
         postgresql-16-postgis-3 \
         postgresql-16-postgis-3-scripts \
+    && apt-get install -y --no-install-recommends postgresql-16-cron \
     && rm -rf /var/lib/apt/lists/*
 
-COPY init_postgis.sh /docker-entrypoint-initdb.d/10_init_postgis.sh
-COPY init_partman_bgw.sh /docker-entrypoint-initdb.d/99_init_partman_bgw.sh

internal/server/postgres/infra/init_partman_bgw.sh

@@ -58,7 +58,13 @@ func setupTestMain(ctx context.Context, m *testing.M) (code int, err error) {
 			"POSTGRES_PASSWORD": "postgres",
 			"POSTGRES_DB":       "postgres",
 		},
-		Cmd:          []string{"postgres", "-c", "fsync=off"},
+		Cmd: []string{
+			"postgres",
+			"-c",
+			"fsync=off",
+			"-c",
+			"shared_preload_libraries=pg_cron",
+		},
 		ExposedPorts: []string{"5432/tcp"},
 		WaitingFor: wait.ForAll(
 			wait.ForLog(

internal/server/postgres/infra/init_postgis.sh

@@ -16,9 +16,12 @@
  * One geometry can have multiple sources, e.g. the UK nation geometry can have solar, wind, etc.
  */
 
-CREATE SCHEMA loc;
 CREATE EXTENSION IF NOT EXISTS btree_gist;
-CREATE EXTENSION IF NOT EXISTS postgis;
+CREATE EXTENSION IF NOT EXISTS postgis WITH SCHEMA public;
+CREATE SCHEMA IF NOT EXISTS topology;
+CREATE EXTENSION IF NOT EXISTS postgis_topology WITH SCHEMA topology;
+
+CREATE SCHEMA loc;
 
 /*- Lookups -----------------------------------------------------------------------------------*/
 
@@ -164,3 +167,6 @@ CREATE INDEX ON loc.sources_mv USING gist (sys_period);
 
 -- +goose Down
 DROP SCHEMA loc CASCADE;
+DROP EXTENSION IF EXISTS postgis_topology;
+DROP EXTENSION IF EXISTS postgis;
+DROP EXTENSION IF EXISTS btree_gist;

internal/server/postgres/package_test.go

@@ -8,6 +8,10 @@
  * with these providers provide access to the data in order to test the accuracy of predictions.
 */
 
+CREATE SCHEMA IF NOT EXISTS partman;
+CREATE EXTENSION IF NOT EXISTS pg_partman WITH SCHEMA partman;
+CREATE EXTENSION IF NOT EXISTS pg_cron;
+
 CREATE SCHEMA obs;
 
 /*- Tables ----------------------------------------------------------------------------------*/
@@ -83,6 +87,15 @@ SET
     infinite_time_partitions = TRUE
 WHERE parent_table = 'obs.observed_generation_values';
 SELECT partman.run_maintenance('obs.observed_generation_values');
+-- Schedule regular maintenance for the partitioned observed generation values table.
+SELECT cron.schedule('partman-maintenance', '@hourly', $$CALL partman.run_maintenance_proc()$$);
+SELECT cron.schedule('cron-details-cleanup', '0 12 * * *', $$DELETE FROM cron.job_run_details WHERE end_time < now() - interval '7 days'$$);
+
 
 -- +goose Down
+SELECT cron.unschedule('partman-maintenance');
 DROP SCHEMA obs CASCADE;
+
+DROP EXTENSION IF EXISTS pg_cron CASCADE;
+DROP EXTENSION IF EXISTS pg_partman CASCADE;
+DROP SCHEMA IF EXISTS partman CASCADE;

internal/server/postgres/sql/migrations/00002_locations.sql

@@ -125,7 +125,7 @@ CREATE TABLE pred.forecasts (
     PRIMARY KEY (forecast_uuid),
     UNIQUE (geometry_uuid, source_type_id, forecaster_id, init_time_utc)
 );
-CREATE INDEX ON pred.forecasts USING GIST (target_period);
+CREATE INDEX idx_forecasts_filter ON pred.forecasts USING GIST (geometry_uuid, source_type_id, target_period);
 
 /*
  * Table to store predicted generation values.
@@ -200,3 +200,4 @@ SELECT partman.run_maintenance('pred.predicted_generation_values');
 
 -- +goose Down
 DROP SCHEMA pred CASCADE;
+

internal/server/postgres/sql/migrations/00003_observations.sql

@@ -106,23 +106,33 @@ INSERT INTO pred.predicted_generation_values (
  * and source type made by all forecasters. Only forecasts that are older than the pivot time
  * minus the specified horizon are considered.
  */
+WITH latest_forecasts AS (
+    SELECT DISTINCT ON (f.forecaster_id)
+        f.forecast_uuid,
+        f.init_time_utc,
+        f.source_type_id,
+        f.geometry_uuid,
+        f.forecaster_id
+    FROM pred.forecasts AS f
+    WHERE f.geometry_uuid = $1
+        AND f.source_type_id = $2
+        AND f.init_time_utc
+        <= sqlc.arg(pivot_timestamp)::TIMESTAMP - MAKE_INTERVAL(mins => sqlc.arg(horizon_mins)::INTEGER)
+        AND f.target_period @> sqlc.arg(pivot_timestamp)::TIMESTAMP
+    ORDER BY
+        f.forecaster_id ASC,
+        f.forecast_uuid DESC
+)
+-- Only join to forecaster table to sort by name once forecasts have been filtered
 SELECT DISTINCT ON (fr.forecaster_name)
-    f.forecast_uuid,
-    f.init_time_utc,
-    f.source_type_id,
-    f.geometry_uuid,
+    lf.*,
     fr.forecaster_name,
     fr.forecaster_version,
-    UUIDV7_EXTRACT_TIMESTAMP(f.forecast_uuid) AS created_at_utc
-FROM pred.forecasts AS f
+    UUIDV7_EXTRACT_TIMESTAMP(lf.forecast_uuid) AS created_at_utc
+FROM latest_forecasts AS lf
     INNER JOIN pred.forecasters AS fr USING (forecaster_id)
-WHERE f.geometry_uuid = $1
-    AND f.source_type_id = $2
-    AND f.init_time_utc <= sqlc.arg(pivot_timestamp)::TIMESTAMP - MAKE_INTERVAL(mins => sqlc.arg(horizon_mins)::INTEGER)
-    AND f.target_period @> sqlc.arg(pivot_timestamp)::TIMESTAMP
 ORDER BY
-    fr.forecaster_name ASC,
-    f.init_time_utc DESC;
+    fr.forecaster_name ASC;
 
 -- name: ListForecasts :many
 /* ListForecasts retrieves all the forecasts for a given location, source type, and forecaster