feat(data): add user enrollment data table

birdxplorer/models.py

@@ -466,8 +466,10 @@ class ModelingPopulation(str, Enum):
     treatment = "EXPANSION"
 
 
-UserEnrollmentLastStateChangeTimeStamp = Union[TwitterTimestamp, Literal["0"], Literal["103308100"]]
-UserEnrollmentLastEarnOutTimestamp = Union[TwitterTimestamp, Literal["1"]]
+UserEnrollmentLastStateChangeTimeStamp = Union[
+    TwitterTimestamp, Literal["0"], Literal["103308100"], Literal[0], Literal[103308100]
+]
+UserEnrollmentLastEarnOutTimestamp = Union[TwitterTimestamp, Literal["1"], Literal[1]]
 
 
 class UserEnrollment(BaseModel):

birdxplorer/storage.py

@@ -1,11 +1,54 @@
+from sqlalchemy import BigInteger, Column, Float, Integer, String
+from sqlalchemy.engine import Engine, create_engine
+from sqlalchemy.orm import DeclarativeBase, Session
+
+from .exceptions import UserEnrollmentNotFoundError
 from .models import ParticipantId, UserEnrollment
 from .settings import GlobalSettings
 
 
+class Base(DeclarativeBase):
+    ...
+
+
+class UserEnrollmentT(Base):
+    __tablename__ = "user_enrollment"
+    participant_id = Column(String, primary_key=True)
+    enrollment_state = Column(String)
+    successful_rating_needed_to_earn_in = Column(Integer)
+    timestamp_of_last_state_change = Column(BigInteger)
+    timestamp_of_last_earn_out = Column(BigInteger)
+    modeling_population = Column(String)
+    modeling_group = Column(Float)
+
+
 class Storage:
+    def __init__(self, engine: Engine) -> None:
+        self._engine = engine
+
+    def __delete__(self) -> None:
+        self._engine.dispose()
+
+    @property
+    def engine(self) -> Engine:
+        return self._engine
+
     def get_user_enrollment_by_participant_id(self, participant_id: ParticipantId) -> UserEnrollment:
-        raise NotImplementedError
+        with Session(self.engine) as sess:
+            row = sess.get(UserEnrollmentT, str(participant_id))
+            if row is None:
+                raise UserEnrollmentNotFoundError(participant_id=participant_id)
+            return UserEnrollment(
+                participant_id=ParticipantId.from_str(str(row.participant_id)),
+                enrollment_state=row.enrollment_state,
+                successful_rating_needed_to_earn_in=row.successful_rating_needed_to_earn_in,
+                timestamp_of_last_state_change=row.timestamp_of_last_state_change,
+                timestamp_of_last_earn_out=row.timestamp_of_last_earn_out,
+                modeling_population=row.modeling_population,
+                modeling_group=row.modeling_group,
+            )
 
 
 def gen_storage(settings: GlobalSettings) -> Storage:
-    return Storage()
+    engine = create_engine(settings.storage_settings.sqlalchemy_database_url.unicode_string())
+    return Storage(engine=engine)

tests/conftest.py

@@ -1,10 +1,17 @@
 from collections.abc import Generator
-from typing import Type
+from typing import Dict, List, Type, Union
 
 from polyfactory.factories.pydantic_factory import ModelFactory
 from pytest import fixture
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.orm import Session
 
-from birdxplorer.settings import GlobalSettings
+from birdxplorer.settings import GlobalSettings, PostgresStorageSettings
+from birdxplorer.storage import Base, UserEnrollmentT
+
+TEST_DATABASE_NAME = "test"
 
 
 @fixture
@@ -13,3 +20,153 @@ class GlobalSettingsFactory(ModelFactory[GlobalSettings]):
         __model__ = GlobalSettings
 
     yield GlobalSettingsFactory
+
+
+def create_engine_from_settings(settings: PostgresStorageSettings) -> Engine:
+    if not settings.sqlalchemy_database_url:
+        raise ValueError("SQLAlchemy database URL is not set")
+    return create_engine(settings.sqlalchemy_database_url.unicode_string())
+
+
+@fixture(scope="session")
+def settings_for_default() -> Generator[PostgresStorageSettings, None, None]:
+    settings = PostgresStorageSettings(
+        host="127.0.0.1", port=5432, username="postgres", password="postgres", database="postgres"
+    )
+    yield settings
+
+
+@fixture(scope="session")
+def settings_for_test() -> Generator[PostgresStorageSettings, None, None]:
+    settings = PostgresStorageSettings(
+        host="127.0.0.1", port=5432, username="postgres", password="postgres", database=TEST_DATABASE_NAME
+    )
+    yield settings
+
+
+@fixture(scope="session")
+def db_for_test(settings_for_default: PostgresStorageSettings) -> Generator[int, None, None]:
+    engine = create_engine_from_settings(settings_for_default)
+    conn = engine.connect()
+    conn.execute(text("commit"))
+    try:
+        conn.execute(text(f"drop database {TEST_DATABASE_NAME}"))
+    except SQLAlchemyError:
+        pass
+    finally:
+        conn.close()
+
+    conn = engine.connect()
+    conn.execute(text("commit"))
+    conn.execute(text(f"create database {TEST_DATABASE_NAME}"))
+    conn.close()
+    yield 1
+
+    conn = engine.connect()
+    conn.execute(text("commit"))
+    conn.execute(text(f"drop database {TEST_DATABASE_NAME}"))
+    conn.close()
+
+    engine.dispose()
+
+
+@fixture(scope="session")
+def engine_for_test(settings_for_test: PostgresStorageSettings, db_for_test: int) -> Generator[Engine, None, None]:
+    engine = create_engine_from_settings(settings_for_test)
+    yield engine
+    engine.dispose()
+
+
+@fixture(scope="session")
+def schema_for_test(settings_for_test: PostgresStorageSettings, db_for_test: int) -> Generator[int, None, None]:
+    engine = create_engine_from_settings(settings_for_test)
+    Base.metadata.create_all(engine)
+
+    yield 1
+    engine.dispose()
+
+
+@fixture(scope="session")
+def user_enrollment_data_list() -> Generator[List[Dict[str, Union[int, str, float]]], None, None]:
+    yield [
+        {
+            "participantId": "D1B8692FB8A8E940F231F606D316881740AF41D73909FE28D5555DA36760A215",
+            "enrollmentState": "newUser",
+            "successfulRatingNeededToEarnIn": 5,
+            "timestampOfLastStateChange": 1679950000000,
+            "timestampOfLastEarnOut": 1,
+            "modelingPopulation": "CORE",
+            "modelingGroup": 13.0,
+        },
+        {
+            "participantId": "49A4DF10193C2FD2D189515D21CEB0BE30F1D2F4297B7D45021A74EF830241A1",
+            "enrollmentState": "earnedIn",
+            "successfulRatingNeededToEarnIn": 5,
+            "timestampOfLastStateChange": 1679960000000,
+            "timestampOfLastEarnOut": 1,
+            "modelingPopulation": "CORE",
+            "modelingGroup": 13.0,
+        },
+        {
+            "participantId": "34166A0D0B81720C6797EF9CD7896EAAD777C140FE718BB9A3CF3236234469A7",
+            "enrollmentState": "newUser",
+            "successfulRatingNeededToEarnIn": 5,
+            "timestampOfLastStateChange": 1679970000000,
+            "timestampOfLastEarnOut": 1,
+            "modelingPopulation": "CORE",
+            "modelingGroup": 13.0,
+        },
+        {
+            "participantId": "C571B31494AC67D57F269B340D689A23B62E5BD73DC51538A6EE917A1EA1D633",
+            "enrollmentState": "earnedIn",
+            "successfulRatingNeededToEarnIn": 5,
+            "timestampOfLastStateChange": 1679980000000,
+            "timestampOfLastEarnOut": 1,
+            "modelingPopulation": "CORE",
+            "modelingGroup": 13.0,
+        },
+        {
+            "participantId": "0EB851180C7BC42F126573C062EC735A3FF073089499A856BE6FC3E2BE823AAD",
+            "enrollmentState": "newUser",
+            "successfulRatingNeededToEarnIn": 5,
+            "timestampOfLastStateChange": 1679990000000,
+            "timestampOfLastEarnOut": 1,
+            "modelingPopulation": "CORE",
+            "modelingGroup": 13.0,
+        },
+    ]
+
+
+@fixture(scope="session")
+def user_enrollment_records(
+    user_enrollment_data_list: List[Dict[str, Union[int, str, float]]], engine_for_test: Engine, schema_for_test: int
+) -> Generator[int, None, None]:
+    with Session(engine_for_test) as session:
+        session.query(UserEnrollmentT).delete()
+        for data in user_enrollment_data_list:
+            session.execute(
+                text(
+                    """
+                    insert into user_enrollment (
+                        participant_id,
+                        enrollment_state,
+                        successful_rating_needed_to_earn_in,
+                        timestamp_of_last_state_change,
+                        timestamp_of_last_earn_out,
+                        modeling_population,
+                        modeling_group
+                    ) values (
+                        :participantId,
+                        :enrollmentState,
+                        :successfulRatingNeededToEarnIn,
+                        :timestampOfLastStateChange,
+                        :timestampOfLastEarnOut,
+                        :modelingPopulation,
+                        :modelingGroup
+                    )
+                    """,
+                ),
+                data,
+            )
+        session.commit()
+    yield 1

tests/docker-compose.yml

@@ -0,0 +1,14 @@
+version: '3.1'
+services:
+  db:
+    image: postgres
+    volumes:
+      - bx-db:/var/lib/postgresql/data
+    environment:
+      - POSTGRES_PASSWORD=postgres
+    restart: unless-stopped
+    ports:
+      - 5432:5432
+
+volumes:
+  bx-db:

tests/test_storage.py

@@ -0,0 +1,26 @@
+from typing import Dict, List, Union
+
+from birdxplorer.models import ParticipantId, UserEnrollment
+from birdxplorer.settings import GlobalSettings, PostgresStorageSettings
+from birdxplorer.storage import gen_storage
+
+
+def test_get_user_enrollment_by_participant_id_returns_data(
+    user_enrollment_data_list: List[Dict[str, Union[int, str, float]]],
+    user_enrollment_records: int,
+    settings_for_test: PostgresStorageSettings,
+) -> None:
+    s = GlobalSettings()
+    s.storage_settings = settings_for_test
+    storage = gen_storage(settings=s)
+    for data in user_enrollment_data_list:
+        pid = ParticipantId.from_str(str(data["participantId"]))
+        result = storage.get_user_enrollment_by_participant_id(pid)
+        assert isinstance(result, UserEnrollment)
+        assert result.participant_id == data["participantId"]
+        assert result.enrollment_state == data["enrollmentState"]
+        assert result.successful_rating_needed_to_earn_in == data["successfulRatingNeededToEarnIn"]
+        assert result.timestamp_of_last_state_change == data["timestampOfLastStateChange"]
+        assert result.timestamp_of_last_earn_out == data["timestampOfLastEarnOut"]
+        assert result.modeling_population == data["modelingPopulation"]
+        assert result.modeling_group == data["modelingGroup"]