From c4d4b31daeef32cefe3453b78c1a671b31cc4cca Mon Sep 17 00:00:00 2001 From: Raphael Paul Laude Date: Sun, 21 Jul 2024 16:54:33 -0400 Subject: [PATCH] test create doc fixing some issues with the testing set up along the way --- backend/app/alembic/env.py | 5 -- .../6f0559e497ba_assignments_table.py | 71 +++++++++++++++++++ backend/app/main.py | 27 +++++-- backend/app/models.py | 21 +++--- backend/app/test_main.py | 29 +++++--- 5 files changed, 124 insertions(+), 29 deletions(-) create mode 100644 backend/app/alembic/versions/6f0559e497ba_assignments_table.py diff --git a/backend/app/alembic/env.py b/backend/app/alembic/env.py index 65ac97cc..c65c4665 100644 --- a/backend/app/alembic/env.py +++ b/backend/app/alembic/env.py @@ -34,11 +34,6 @@ def get_url(): - database_url = os.getenv("DATABASE_URL", None) - - if database_url: - return database_url - user = os.getenv("POSTGRES_USER", "postgres") password = os.getenv("POSTGRES_PASSWORD", "") server = os.getenv("POSTGRES_SERVER", "db") diff --git a/backend/app/alembic/versions/6f0559e497ba_assignments_table.py b/backend/app/alembic/versions/6f0559e497ba_assignments_table.py new file mode 100644 index 00000000..1fdf7e9e --- /dev/null +++ b/backend/app/alembic/versions/6f0559e497ba_assignments_table.py @@ -0,0 +1,71 @@ +"""assignments table + +Revision ID: 6f0559e497ba +Revises: 966d8d72887e +Create Date: 2024-07-21 16:08:29.504177 + +""" + +from typing import Sequence, Union + +from alembic import op +from sqlmodel.sql import sqltypes +import sqlalchemy as sa + +from app.models import UUIDType + + +# revision identifiers, used by Alembic. +revision: str = "6f0559e497ba" +down_revision: Union[str, None] = "966d8d72887e" +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.create_table( + "document", + sa.Column( + "created_at", + sa.TIMESTAMP(timezone=True), + server_default=sa.text("CURRENT_TIMESTAMP"), + nullable=False, + ), + sa.Column( + "updated_at", + sa.TIMESTAMP(timezone=True), + server_default=sa.text("CURRENT_TIMESTAMP"), + nullable=False, + ), + sa.Column("document_id", UUIDType(), nullable=False), + sa.PrimaryKeyConstraint("document_id"), + sa.UniqueConstraint("document_id"), + ) + op.create_table( + "assignments", + sa.Column("document_id", sa.UUID(), nullable=False), + sa.Column("geo_id", sqltypes.AutoString(), nullable=False), + sa.Column("zone", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint( + ["document_id"], + ["document.document_id"], + ), + sa.PrimaryKeyConstraint("document_id", "geo_id"), + postgresql_partition_by="LIST (document_id)", + ) + op.alter_column("gerrydbtable", "uuid", existing_type=sa.UUID(), nullable=False) + op.drop_column("gerrydbtable", "id") + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "gerrydbtable", + sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False), + ) + op.alter_column("gerrydbtable", "uuid", existing_type=sa.UUID(), nullable=True) + op.drop_table("assignments") + op.drop_table("document") + # ### end Alembic commands ### diff --git a/backend/app/main.py b/backend/app/main.py index b93ac970..89bc8d39 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -5,11 +5,12 @@ from sqlalchemy.dialects.postgresql import insert from typing import List import logging +from uuid import uuid4 import sentry_sdk from app.core.db import engine from app.core.config import settings -from app.models import Assignments, Document +from app.models import Assignments, Document, DocumentPublic if settings.ENVIRONMENT == "production": sentry_sdk.init( @@ -58,26 +59,40 @@ async def db_is_alive(session: Session = Depends(get_session)): ) -@app.post("/create_document") +@app.post( + "/create_document", + response_model=DocumentPublic, + status_code=status.HTTP_201_CREATED, +) async def create_document(session: Session = Depends(get_session)): - doc = Document() + # To be created in the database + document_id = str(uuid4().hex).replace("-", "") + print(document_id) + doc = Document.model_validate({"document_id": document_id}) session.add(doc) session.commit() session.refresh(doc) - document_id = doc.document_id + + if not doc.document_id: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Document creation failed", + ) + + document_id: str = doc.document_id.replace("-", "") # Also create the partition in one go. session.execute( text( f""" CREATE TABLE assignments_{document_id} PARTITION OF assignments - VALUES IN ('{document_id}') + FOR VALUES IN ('{document_id}') """ ) ) return doc -@app.post("/update_assignments") +@app.patch("/update_assignments") async def update_assignments( assignments: List[Assignments], session: Session = Depends(get_session) ): diff --git a/backend/app/models.py b/backend/app/models.py index b5a2d895..dd675227 100644 --- a/backend/app/models.py +++ b/backend/app/models.py @@ -1,5 +1,6 @@ from datetime import datetime from typing import Optional +from pydantic import BaseModel from sqlmodel import Field, SQLModel, UUID, TIMESTAMP, text, Column @@ -29,17 +30,21 @@ class TimeStampMixin(SQLModel): ) -class GerryDBTableBase(TimeStampMixin, SQLModel): - id: int = Field(default=None, primary_key=True) +class GerryDBTable(TimeStampMixin, SQLModel, table=True): + uuid: str = Field(sa_column=Column(UUIDType, unique=True, primary_key=True)) + name: str = Field(nullable=False, unique=True) -class GerryDBTable(GerryDBTableBase, table=True): - uuid: str = Field(sa_column=Column(UUIDType, unique=True)) - name: str = Field(nullable=False, unique=True) +class Document(TimeStampMixin, SQLModel, table=True): + document_id: str | None = Field( + sa_column=Column(UUIDType, unique=True, primary_key=True) + ) -class Document(TimeStampMixin, SQLModel): - document_id: str | None = Field(sa_column=Column(UUIDType, unique=True)) +class DocumentPublic(BaseModel): + document_id: str + created_at: datetime + updated_at: datetime class Assignments(SQLModel, table=True): @@ -47,4 +52,4 @@ class Assignments(SQLModel, table=True): document_id: str = Field(foreign_key="document.document_id", primary_key=True) geo_id: str = Field(primary_key=True) zone: int - __table_args__ = {"postgres_partition_by": "document_id"} + __table_args__ = {"postgresql_partition_by": "document_id"} diff --git a/backend/app/test_main.py b/backend/app/test_main.py index ce7111c1..fed7a0bb 100644 --- a/backend/app/test_main.py +++ b/backend/app/test_main.py @@ -8,6 +8,7 @@ from sqlalchemy import text from sqlalchemy.exc import OperationalError, ProgrammingError import subprocess +import uuid client = TestClient(app) @@ -65,18 +66,17 @@ def engine_fixture(request): except (OperationalError, ProgrammingError): pass - if ENVIRONMENT != "test": - subprocess.run(["alembic", "upgrade", "head"], check=True, env=my_env) + subprocess.run(["alembic", "upgrade", "head"], check=True, env=my_env) def teardown(): - close_connections_query = f""" - SELECT pg_terminate_backend(pg_stat_activity.pid) - FROM pg_stat_activity - WHERE pg_stat_activity.datname = '{POSTGRES_TEST_DB}' - AND pid <> pg_backend_pid(); - """ - conn.execute(text(close_connections_query)) - conn.execute(text(f"DROP DATABASE {POSTGRES_TEST_DB}")) + # close_connections_query = f""" + # SELECT pg_terminate_backend(pg_stat_activity.pid) + # FROM pg_stat_activity + # WHERE pg_stat_activity.datname = '{POSTGRES_TEST_DB}' + # AND pid <> pg_backend_pid(); + # """ + # conn.execute(text(close_connections_query)) + # conn.execute(text(f"DROP DATABASE {POSTGRES_TEST_DB}")) conn.close() request.addfinalizer(teardown) @@ -109,3 +109,12 @@ def test_db_is_alive(client): response = client.get("/db_is_alive") assert response.status_code == 200 assert response.json() == {"message": "DB is alive"} + + +def test_new_document(client): + print(TEST_SQLALCHEMY_DATABASE_URI) + response = client.post("/create_document") + assert response.status_code == 201 + document_id = response.json().get("document_id", None) + assert document_id is not None + assert isinstance(uuid.UUID(document_id), uuid.UUID)