Partition Column_Value by geoid

gerrydb_meta/crud/column.py

@@ -5,13 +5,14 @@
 from datetime import datetime, timezone
 from typing import Any, Collection, Tuple
 
-from sqlalchemy import exc, insert, update
+from sqlalchemy import exc, insert, update, text
 from sqlalchemy.orm import Session
 
 from gerrydb_meta import models, schemas
 from gerrydb_meta.crud.base import NamespacedCRBase, normalize_path
 from gerrydb_meta.enums import ColumnType
 from gerrydb_meta.exceptions import ColumnValueTypeError, CreateValueError
+from gerrydb_meta.utils import create_column_value_partition_text
 
 log = logging.getLogger()
 
@@ -80,6 +81,9 @@ def create(
             canonical_ref.col_id = col.col_id
             db.flush()
 
+            # create partition
+            db.execute(create_column_value_partition_text(column_id=col.col_id))
+
             # Create additional aliases (non-canonical references) to the column.
             if obj_in.aliases:
                 self._add_aliases(
@@ -202,6 +206,10 @@ def set_values(
 
         # Add the new column values and invalidate the old ones where present.
         geo_ids = [geo.geo_id for geo, _ in values]
+
+        # make sure partition exists for column
+        db.execute(create_column_value_partition_text(column_id=col.col_id))
+
         with_tuples = (
             db.query(
                 models.ColumnValue.col_id,

gerrydb_meta/models.py

@@ -5,7 +5,7 @@
 from uuid import UUID, uuid4
 
 from geoalchemy2 import Geography as SqlGeography
-from sqlalchemy import JSON, BigInteger, Boolean, CheckConstraint, DateTime
+from sqlalchemy import JSON, BigInteger, Boolean, CheckConstraint, DateTime, text, event
 from sqlalchemy import Enum as SqlEnum
 from sqlalchemy import (
     ForeignKey,
@@ -17,7 +17,7 @@
     UniqueConstraint,
 )
 from sqlalchemy.dialects import postgresql
-from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship, Session
 from sqlalchemy.sql import func
 
 from gerrydb_meta.enums import (
@@ -27,8 +27,10 @@
     ScopeType,
     ViewRenderStatus,
 )
+from gerrydb_meta.utils import create_column_value_partition_text
 
-metadata_obj = MetaData(schema="gerrydb")
+SCHEMA = "gerrydb"
+metadata_obj = MetaData(schema=SCHEMA)
 
 
 class Base(DeclarativeBase):
@@ -547,7 +549,10 @@ class ColumnSetMember(Base):
 
 class ColumnValue(Base):
     __tablename__ = "column_value"
-    __table_args__ = (UniqueConstraint("col_id", "geo_id", "valid_from"),)
+    __table_args__ = (
+        UniqueConstraint("col_id", "geo_id", "valid_from"),
+        {"postgresql_partition_by": "LIST (col_id)"},
+    )
 
     col_id: Mapped[int] = mapped_column(
         Integer,

gerrydb_meta/utils.py

@@ -0,0 +1,8 @@
+from sqlalchemy import text
+from gerrydb_meta import models
+
+
+def create_column_value_partition_text(column_id: int):
+    table_name = models.ColumnValue.__table__.name
+    sql = f"CREATE TABLE IF NOT EXISTS {models.SCHEMA}.{table_name}_{column_id} PARTITION OF {models.SCHEMA}.{table_name} FOR VALUES IN ({column_id})"
+    return text(sql)

tests/test_utils.py

@@ -0,0 +1,12 @@
+from gerrydb_meta.utils import create_column_value_partition_text
+from sqlalchemy import text
+
+
+def test_create_column_value_partition_text():
+    column_id = 42
+    got = create_column_value_partition_text(column_id=column_id)
+    wanted = text(
+        "CREATE TABLE IF NOT EXISTS gerrydb.column_value_42 PARTITION OF gerrydb.column_value FOR VALUES IN (42)"
+    )
+    # different object instances, so compare string form
+    assert str(got) == str(wanted)