Add plan geometry

README.md

@@ -56,6 +56,7 @@ docker network ls --format {{.Name}} |grep pytest | awk '{print $1}' | xargs -I
 3. If you want to change *all* tables in a schema (i.e. edit *all* the code tables, or add a field to *all* the data tables), the abstract base classes are in [base.py](./database/base.py).
 4. If you only want to change/add *one* code table or one data table, please edit/add the right table in [codes.py](./database/codes.py) or [models.py](./database/models.py).
 5. To get the changes tested and usable in your functions, create a new database revision with `make revision name="describe_your_changes"`, e.g. `make revision name="add_plan_object_table"`. This creates a new random id (`uuid`) for your migration, and a revision file `YYYY-MM-DD-HHMM-uuid-add_plan_object_table` in the [alembic versions dir](./database/migrations/versions). Please check that the autogenerated revision file seems to do approximately sensible things.
+    - Specifically, when adding geometry fields, please note [GeoAlchemy2 bug with Alembic](https://geoalchemy-2.readthedocs.io/en/latest/alembic.html#interactions-between-alembic-and-geoalchemy-2), which means you will have to *manually remove* `op.create_index` and `op.drop_index` in the revision file. This is because GeoAlchemy2 already automatically creates geometry index whenever adding a geometry column.
 6. Run tests with `make pytest` to check that the revision file runs correctly. At minimum, you may have to change the tested table counts (codes_count and hame_count) in [migration tests](./database/test/test_db.py) to reflect the correct number of tables in the database.
 7. Run `make rebuild` and `make test-create-db` to start development instance with the new model.
 <!-- 8. To update the [database documentation](./backend/databasemodel/dbdoc/README.md) to reflect the changes, install [tbls](https://github.com/k1LoW/tbls) and run `tbls doc --force`. -->

database/base.py

@@ -2,11 +2,15 @@
 from datetime import datetime
 from typing import Optional
 
+from geoalchemy2 import Geometry
+from shapely.geometry import Polygon
 from sqlalchemy.dialects.postgresql import JSONB, UUID
 from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 from sqlalchemy.sql import func
 from typing_extensions import Annotated
 
+PROJECT_SRID = 3067
+
 
 class Base(DeclarativeBase):
     """
@@ -16,6 +20,7 @@ class Base(DeclarativeBase):
     type_annotation_map = {
         uuid.UUID: UUID(as_uuid=True),
         dict[str, str]: JSONB,
+        Polygon: Geometry(geometry_type="POLYGON", srid=PROJECT_SRID),
     }
 
 

database/migrations/script.py.mako

@@ -9,6 +9,7 @@ from typing import Sequence, Union
 
 from alembic import op
 import sqlalchemy as sa
+import geoalchemy2
 ${imports if imports else ""}
 
 # revision identifiers, used by Alembic.

database/migrations/versions/2024_01_29_1341-9f82c38f45a9_add_plan_geometry.py

@@ -0,0 +1,44 @@
+"""add plan geometry
+
+Revision ID: 9f82c38f45a9
+Revises: 6592d88a81df
+Create Date: 2024-01-29 13:41:44.839255
+
+"""
+from typing import Sequence, Union
+
+import geoalchemy2
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "9f82c38f45a9"
+down_revision: Union[str, None] = "6592d88a81df"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "plan",
+        sa.Column(
+            "geom",
+            geoalchemy2.types.Geometry(
+                geometry_type="POLYGON",
+                srid=3067,
+                from_text="ST_GeomFromEWKT",
+                name="geometry",
+                nullable=False,
+            ),
+            nullable=False,
+        ),
+        schema="hame",
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("plan", "geom", schema="hame")
+    # ### end Alembic commands ###

database/models.py

@@ -3,6 +3,7 @@
 
 from base import VersionedBase, language_str
 from codes import LifeCycleStatus
+from shapely.geometry import Polygon
 from sqlalchemy import ForeignKey
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
@@ -20,3 +21,4 @@ class Plan(VersionedBase):
         ForeignKey("codes.lifecycle_status.id")
     )
     lifecycle_status: Mapped[LifeCycleStatus] = relationship(back_populates="plans")
+    geom: Mapped[Polygon]

database/test/test_db.py

@@ -55,15 +55,31 @@ def assert_database_is_alright(
         assert (os.environ.get("ADMIN_USER"), "UPDATE") in grants
         assert (os.environ.get("ADMIN_USER"), "DELETE") in grants
 
-        # TODO: Do we need to check constraint naming?
-        # cur.execute(
-        #     "SELECT con.conname FROM pg_catalog.pg_constraint con INNER JOIN pg_catalog.pg_class rel ON rel.oid = con.conrelid "
-        #     "INNER JOIN pg_catalog.pg_namespace nsp ON nsp.oid = connamespace WHERE "
-        #     f"nsp.nspname = 'kooste' AND rel.relname = '{table_name}';"
-        # )
-        # constraints = cur.fetchall()
-        # if constraints:
-        #     assert (f"{table_name}_pk",) in constraints
+        # Check indexes
+        cur.execute(
+            f"SELECT * FROM pg_indexes WHERE schemaname = 'hame' AND tablename = '{table_name}';"
+        )
+        indexes = cur.fetchall()
+        cur.execute(
+            f"SELECT column_name FROM information_schema.columns WHERE table_schema = 'hame' AND table_name = '{table_name}';"
+        )
+        columns = cur.fetchall()
+        if ("id",) in columns:
+            assert (
+                "hame",
+                table_name,
+                f"{table_name}_pkey",
+                None,
+                f"CREATE UNIQUE INDEX {table_name}_pkey ON hame.{table_name} USING btree (id)",
+            ) in indexes
+        if ("geom",) in columns:
+            assert (
+                "hame",
+                table_name,
+                f"idx_{table_name}_geom",
+                None,
+                f"CREATE INDEX idx_{table_name}_geom ON hame.{table_name} USING gist (geom)",
+            ) in indexes
 
     # Check code tables
     cur.execute("SELECT tablename, tableowner FROM pg_tables WHERE schemaname='codes';")

requirements.in

@@ -1,4 +1,6 @@
 alembic
 boto3
+geoalchemy2
 psycopg2-binary
+shapely
 sqlalchemy

requirements.txt

@@ -12,6 +12,8 @@ botocore==1.34.15
     # via
     #   boto3
     #   s3transfer
+geoalchemy2==0.14.3
+    # via -r requirements.in
 jmespath==1.0.1
     # via
     #   boto3
@@ -20,18 +22,25 @@ mako==1.3.0
     # via alembic
 markupsafe==2.1.3
     # via mako
+numpy==1.26.3
+    # via shapely
+packaging==23.2
+    # via geoalchemy2
 psycopg2-binary==2.9.9
     # via -r requirements.in
 python-dateutil==2.8.2
     # via botocore
 s3transfer==0.10.0
     # via boto3
+shapely==2.0.2
+    # via -r requirements.in
 six==1.16.0
     # via python-dateutil
 sqlalchemy==2.0.25
     # via
     #   -r requirements.in
     #   alembic
+    #   geoalchemy2
 typing-extensions==4.9.0
     # via
     #   alembic