Merge pull request #22 from CCRI-POPROX/karl/feature/opt-out

Add `opt_out` column to `expt_assignments` table

.pre-commit-config.yaml

@@ -28,9 +28,11 @@ repos:
       - id: ruff
         args: [--fix]
         name: auto-fix Python lint errors
+        exclude: ^poprox-db/migrations/versions
       # Run the formatter.
       - id: ruff-format
         name: format Python source
+        exclude: ^poprox-db/migrations/versions
   - repo: local
     hooks:
       - id: format-toml-files

poprox-db/migrations/versions/2024_07_26_0945-fc8d79481ee4_merge_head_revisions.py

@@ -0,0 +1,23 @@
+"""merge head revisions
+
+Revision ID: fc8d79481ee4
+Revises: 14be38931ae0, 5d566fff8aaa
+Create Date: 2024-07-26 09:45:33.246538
+
+"""
+
+from typing import Sequence, Union
+
+# revision identifiers, used by Alembic.
+revision: str = "fc8d79481ee4"
+down_revision: Union[str, None] = ("14be38931ae0", "5d566fff8aaa")
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    pass
+
+
+def downgrade() -> None:
+    pass

poprox-db/migrations/versions/2024_07_26_0947-53fc350d4cd3_rename_expt_allocations_to_expt_.py

@@ -0,0 +1,65 @@
+"""rename expt_allocations to expt_assignments
+
+Revision ID: 53fc350d4cd3
+Revises: fc8d79481ee4
+Create Date: 2024-07-26 09:47:21.351380
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "53fc350d4cd3"
+down_revision: Union[str, None] = "fc8d79481ee4"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.drop_constraint("fk_expt_allocations_group_id", "expt_allocations", type_="foreignkey")
+    op.drop_constraint("fk_expt_allocations_account_id", "expt_allocations", type_="foreignkey")
+
+    op.alter_column("expt_allocations", "allocation_id", new_column_name="assignment_id")
+    op.rename_table("expt_allocations", "expt_assignments")
+
+    op.create_foreign_key(
+        "fk_expt_assignments_account_id",
+        "expt_assignments",
+        "accounts",
+        ["account_id"],
+        ["account_id"],
+    )
+
+    op.create_foreign_key(
+        "fk_expt_assignments_group_id",
+        "expt_assignments",
+        "expt_groups",
+        ["group_id"],
+        ["group_id"],
+    )
+
+
+def downgrade() -> None:
+    op.drop_constraint("fk_expt_assignments_group_id", "expt_assignments", type_="foreignkey")
+    op.drop_constraint("fk_expt_assignments_account_id", "expt_assignments", type_="foreignkey")
+
+    op.alter_column("expt_assignments", "assignment_id", new_column_name="allocation_id")
+    op.rename_table("expt_assignments", "expt_allocations")
+
+    op.create_foreign_key(
+        "fk_expt_allocations_account_id",
+        "expt_allocations",
+        "accounts",
+        ["account_id"],
+        ["account_id"],
+    )
+
+    op.create_foreign_key(
+        "fk_expt_allocations_group_id",
+        "expt_allocations",
+        "expt_groups",
+        ["group_id"],
+        ["group_id"],
+    )

poprox-db/migrations/versions/2024_07_26_0953-211c6b3ac5de_add_opt_out_column_to_experiment_.py

@@ -0,0 +1,29 @@
+"""add opt-out column to experiment assignments
+
+Revision ID: 211c6b3ac5de
+Revises: 53fc350d4cd3
+Create Date: 2024-07-26 09:53:57.652145
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "211c6b3ac5de"
+down_revision: Union[str, None] = "53fc350d4cd3"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "expt_assignments",
+        sa.Column("opted_out", sa.Boolean, nullable=False, server_default=sa.false()),
+    )
+
+
+def downgrade() -> None:
+    op.drop_column("expt_assignments", "opted_out")

poprox-db/tests/test_stairstep.py

@@ -21,6 +21,9 @@ def test_migrations_stairway(alembic_config: Config, revision: Script):
     Does not require any maintenance - you just add it once to check 80% of typos
     and mistakes in migrations forever.
     """
+    if isinstance(revision.down_revision, tuple):
+        return  # don't test merge revisions
+
     upgrade(alembic_config, revision.revision)
 
     # We need -1 for downgrading first migration (its down_revision is None)

src/poprox_storage/concepts/experiment.py

@@ -83,7 +83,8 @@ def duration(self) -> timedelta:
         return self.end_date - self.start_date
 
 
-class Allocation(BaseModel):
-    allocation_id: UUID | None = None
+class Assignment(BaseModel):
+    assignment_id: UUID | None = None
     account_id: UUID
     group_id: UUID
+    opted_out: bool | None = None

src/poprox_storage/repositories/experiments.py

@@ -2,11 +2,11 @@
 from uuid import UUID
 
 import tomli
-from sqlalchemy import Connection, Table, and_, select
+from sqlalchemy import Connection, Table, and_, select, update
 
 from poprox_concepts import Account
 from poprox_storage.concepts.experiment import (
-    Allocation,
+    Assignment,
     Experiment,
     Group,
     Phase,
@@ -23,7 +23,7 @@ def __init__(self, connection: Connection):
         super().__init__(connection)
         self.tables: dict[str, Table] = self._load_tables(
             "experiments",
-            "expt_allocations",
+            "expt_assignments",
             "expt_groups",
             "expt_phases",
             "expt_recommenders",
@@ -39,8 +39,8 @@ def store_experiment(self, experiment: Experiment, assignments: dict[str, list[A
             for group in experiment.groups:
                 group.group_id = self._insert_expt_group(experiment_id, group)
                 for account in assignments.get(group.name, []):
-                    allocation = Allocation(account_id=account.account_id, group_id=group.group_id)
-                    self._insert_expt_allocation(allocation)
+                    allocation = Assignment(account_id=account.account_id, group_id=group.group_id)
+                    self._insert_expt_assignment(allocation)
 
             for recommender in experiment.recommenders:
                 recommender.recommender_id = self._insert_expt_recommender(
@@ -104,19 +104,18 @@ def fetch_active_expt_endpoint_urls(self, date: datetime.date | None = None) ->
 
         return recommender_lookup_by_group
 
-    def fetch_active_expt_allocations(self, date: datetime.date | None = None) -> dict[UUID, Allocation]:
-        allocations_tbl = self.tables["expt_allocations"]
+    def fetch_active_expt_assignments(self, date: datetime.date | None = None) -> dict[UUID, Assignment]:
+        assignments_tbl = self.tables["expt_assignments"]
 
         group_ids = self.get_active_expt_group_ids(date)
 
-        # Find accounts allocated to the groups that are assigned the active recommenders above
         group_query = select(
-            allocations_tbl.c.allocation_id, allocations_tbl.c.account_id, allocations_tbl.c.group_id
-        ).where(allocations_tbl.c.group_id.in_(group_ids))
+            assignments_tbl.c.assignment_id, assignments_tbl.c.account_id, assignments_tbl.c.group_id
+        ).where(and_(assignments_tbl.c.group_id.in_(group_ids), assignments_tbl.c.opted_out is False))
         result = self.conn.execute(group_query).fetchall()
         group_lookup_by_account = {
-            row.account_id: Allocation(
-                allocation_id=row.allocation_id, account_id=row.account_id, group_id=row.group_id
+            row.account_id: Assignment(
+                assignment_id=row.assignment_id, account_id=row.account_id, group_id=row.group_id
             )
             for row in result
         }
@@ -125,7 +124,25 @@ def fetch_active_expt_allocations(self, date: datetime.date | None = None) -> di
 
     get_active_expt_group_ids = fetch_active_expt_group_ids
     get_active_expt_endpoint_urls = fetch_active_expt_endpoint_urls
-    get_active_expt_allocations = fetch_active_expt_allocations
+    get_active_expt_allocations = fetch_active_expt_assignments
+
+    def update_expt_assignment_to_opt_out(self, account_id: UUID) -> UUID | None:
+        assignments_tbl = self.tables["expt_assignments"]
+
+        group_ids = self.get_active_expt_group_ids()
+
+        assignment_query = (
+            update(assignments_tbl)
+            .where(
+                and_(
+                    assignments_tbl.c.account_id == account_id,
+                    assignments_tbl.c.group_id.in_(group_ids),
+                    assignments_tbl.c.opted_out is False,
+                )
+            )
+            .values(opted_out=True)
+        )
+        self.conn.execute(assignment_query)
 
     def _insert_experiment(self, experiment: Experiment) -> UUID | None:
         return self._insert_model("experiments", experiment, exclude={"phases"}, commit=False)
@@ -185,13 +202,13 @@ def _insert_expt_treatment(
             commit=False,
         )
 
-    def _insert_expt_allocation(
+    def _insert_expt_assignment(
         self,
-        allocation: Allocation,
+        assignment: Assignment,
     ) -> UUID | None:
         return self._insert_model(
-            "expt_allocations",
-            allocation,
+            "expt_assignments",
+            assignment,
             commit=False,
         )