stress_at_work_analysis/participants/query_db.py

import datetime
from collections.abc import Collection

import pandas as pd
from sqlalchemy import and_

from config.models import Participant
from setup import db_engine, session


def get_usernames(
    tester: bool = False,
    active: bool = False,
    collection_start: datetime.date = datetime.date.fromisoformat("2020-01-01"),
    last_upload: datetime.date = datetime.date.today(),
) -> Collection:
    query_participant_usernames = session.query(Participant.username).filter(
        and_(
            Participant.tester == tester,
            Participant.active == active,
            Participant.collection_start_utc >= collection_start,
            Participant.last_upload_utc < last_upload,
        )
    )
    with db_engine.connect() as connection:
        df_participants = pd.read_sql(query_participant_usernames.statement, connection)
    return df_participants.values.flatten()