From d2d92af72903b94a07bd0f7a41d1514b4cc276b1 Mon Sep 17 00:00:00 2001 From: Riku Oja Date: Mon, 29 Jan 2024 13:46:01 +0200 Subject: [PATCH 1/2] Add geoalchemy2 and plan geometry column --- README.md | 1 + database/base.py | 5 +++ database/migrations/script.py.mako | 1 + ..._29_1341-9f82c38f45a9_add_plan_geometry.py | 44 +++++++++++++++++++ database/models.py | 2 + requirements.in | 2 + requirements.txt | 9 ++++ 7 files changed, 64 insertions(+) create mode 100644 database/migrations/versions/2024_01_29_1341-9f82c38f45a9_add_plan_geometry.py diff --git a/README.md b/README.md index dd44360..48a7fd2 100644 --- a/README.md +++ b/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. diff --git a/database/base.py b/database/base.py index cf252e0..33eeb1c 100644 --- a/database/base.py +++ b/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), } diff --git a/database/migrations/script.py.mako b/database/migrations/script.py.mako index fbc4b07..a3c29b0 100644 --- a/database/migrations/script.py.mako +++ b/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. diff --git a/database/migrations/versions/2024_01_29_1341-9f82c38f45a9_add_plan_geometry.py b/database/migrations/versions/2024_01_29_1341-9f82c38f45a9_add_plan_geometry.py new file mode 100644 index 0000000..e754632 --- /dev/null +++ b/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 ### diff --git a/database/models.py b/database/models.py index 41d034c..9746501 100644 --- a/database/models.py +++ b/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] diff --git a/requirements.in b/requirements.in index 0b9c702..235ba5e 100644 --- a/requirements.in +++ b/requirements.in @@ -1,4 +1,6 @@ alembic boto3 +geoalchemy2 psycopg2-binary +shapely sqlalchemy diff --git a/requirements.txt b/requirements.txt index 1836025..c0efedc 100644 --- a/requirements.txt +++ b/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 From 125e8d182778e8b8c08746f34834172cb236ef1f Mon Sep 17 00:00:00 2001 From: Riku Oja Date: Mon, 29 Jan 2024 14:32:49 +0200 Subject: [PATCH 2/2] Add test to check database indexes --- database/test/test_db.py | 34 +++++++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/database/test/test_db.py b/database/test/test_db.py index a0e7512..7bbb446 100644 --- a/database/test/test_db.py +++ b/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';")