From a01298873a09f3ae2f055cbf86fc7ded006cd715 Mon Sep 17 00:00:00 2001
From: Matt Hammerly <matt.hammerly@sentry.io>
Date: Thu, 6 Jun 2024 13:53:39 -0700
Subject: [PATCH] add partitioning to pg_telemetry SimpleMetric model

---
 shared/django_apps/db_settings.py                        | 2 +-
 .../django_apps/{user_measurements => }/partitioning.py  | 9 +++++++++
 shared/django_apps/pg_telemetry/models.py                | 8 +++++++-
 3 files changed, 17 insertions(+), 2 deletions(-)
 rename shared/django_apps/{user_measurements => }/partitioning.py (76%)

diff --git a/shared/django_apps/db_settings.py b/shared/django_apps/db_settings.py
index e56af718..a1c7c8ad 100644
--- a/shared/django_apps/db_settings.py
+++ b/shared/django_apps/db_settings.py
@@ -162,7 +162,7 @@
 
 # Allows to use the pgpartition command
 PSQLEXTRA_PARTITIONING_MANAGER = (
-    "shared.django_apps.user_measurements.partitioning.manager"
+    "shared.django_apps.partitioning.manager"
 )
 
 DATABASE_ROUTERS = [
diff --git a/shared/django_apps/user_measurements/partitioning.py b/shared/django_apps/partitioning.py
similarity index 76%
rename from shared/django_apps/user_measurements/partitioning.py
rename to shared/django_apps/partitioning.py
index f1fb4136..3f61ee48 100644
--- a/shared/django_apps/user_measurements/partitioning.py
+++ b/shared/django_apps/partitioning.py
@@ -6,6 +6,7 @@
 )
 from psqlextra.partitioning.config import PostgresPartitioningConfig
 
+from shared.django_apps.pg_telemetry.models import SimpleMetric
 from shared.django_apps.user_measurements.models import UserMeasurement
 
 # Overlapping partitions will cause errors - https://www.postgresql.org/docs/current/ddl-partitioning.html#DDL-PARTITIONING-DECLARATIVE -> "create partitions"
@@ -23,5 +24,13 @@
                 max_age=relativedelta(months=12),
             ),
         ),
+        PostgresPartitioningConfig(
+            model=SimpleMetric,
+            strategy=PostgresCurrentTimePartitioningStrategy(
+                size=PostgresTimePartitionSize(months=1),
+                count=3,
+                max_age=relativedelta(months=6),
+            ),
+        ),
     ]
 )
diff --git a/shared/django_apps/pg_telemetry/models.py b/shared/django_apps/pg_telemetry/models.py
index fb045322..05207448 100644
--- a/shared/django_apps/pg_telemetry/models.py
+++ b/shared/django_apps/pg_telemetry/models.py
@@ -1,8 +1,10 @@
 from django.conf import settings
 from django.db import models
+from psqlextra.models import PostgresPartitionedModel
+from psqlextra.types import PostgresPartitioningMethod
 
 
-class BaseModel(models.Model):
+class BaseModel(PostgresPartitionedModel):
     """
     Base model for timeseries metrics. It provides a timestamp field which
     represents the time that the data sample was captured at and a few metadata
@@ -13,6 +15,10 @@ class BaseModel(models.Model):
     Timescale for a time, we'll pick one.
     """
 
+    class PartitioningMeta:
+        method = PostgresPartitioningMethod.RANGE
+        key = ["timestamp"]
+
     class Meta:
         abstract = True