diff --git a/airflow/dags/metrcis/load_firestore_to_snowflake.py b/airflow/dags/metrcis/load_firestore_to_snowflake.py index af03d8cc..ef9c89d2 100644 --- a/airflow/dags/metrcis/load_firestore_to_snowflake.py +++ b/airflow/dags/metrcis/load_firestore_to_snowflake.py @@ -54,14 +54,15 @@ def load_request_data_from_firestore() -> list[tuple[str, int, bool, int]]: .stream() ) - rows: list[tuple[str, int, bool, int]] = [] + rows: list[tuple[str, int | None, bool, datetime]] = [] for doc in docs: doc_dict = doc.to_dict() uuid = doc_dict["uuid"] - score = doc_dict.get("score") or 0 + score = doc_dict.get("score") status = doc_dict.get("status") == "complete" response_received_at = datetime.fromtimestamp(doc_dict.get("response_received_at")) - rows.append((uuid, score, status, response_received_at)) + client = doc_dict.get("client") + rows.append((uuid, score, status, response_received_at, client)) return rows @@ -77,9 +78,9 @@ def write_request_data_to_snowflake(rows: list[tuple[str, int, bool, int]]) -> N insert_sql = f""" INSERT INTO - {METRICS_SNOWFLAKE_DB_DATABASE}.{METRICS_SNOWFLAKE_DB_SCHEMA}.request(uuid, score, success, created_at) + {METRICS_SNOWFLAKE_DB_DATABASE}.{METRICS_SNOWFLAKE_DB_SCHEMA}.request(uuid, score, success, created_at, client) VALUES - (?, ?, ?, ?) + (?, ?, ?, ?, ?) """ conn.cursor().executemany(insert_sql, rows) diff --git a/api/ask_astro/models/request.py b/api/ask_astro/models/request.py index 97f79b7c..06be7f83 100644 --- a/api/ask_astro/models/request.py +++ b/api/ask_astro/models/request.py @@ -57,7 +57,7 @@ class AskAstroRequest(BaseModel): False, description="Whether the request is an example", ) - client: str = Field(None, description="The client type used to send the request -- webapp/slack") + client: str | None = Field(None, description="The client type used to send the request -- webapp/slack") def to_firestore(self) -> dict[str, Any]: """ @@ -107,7 +107,7 @@ def from_dict(cls, response_dict: dict[str, Any]) -> AskAstroRequest: response_received_at=response_dict.get("response_received_at"), is_processed=response_dict.get("is_processed", False), is_example=response_dict.get("is_example", False), - client=response_dict["client"], + client=response_dict.get("client", None), )