import os import time from datetime import datetime, timezone import boto3 from botocore.exceptions import ClientError _table = None def _get_table(): global _table if _table is None: _table = boto3.resource("dynamodb").Table(os.environ["TABLE_NAME"]) return _table def _ttl_90_days(): return int(time.time()) + (90 * 24 * 60 * 60) def message_exists(message_id): resp = _get_table().get_item( Key={"pk": f"MSG#{message_id}", "sk": "MSG"}, ProjectionExpression="pk", ) return "Item" in resp def save_message(msg_data): item = { "pk": f"MSG#{msg_data['message_id']}", "sk": "MSG", "thread_id": msg_data["thread_id"], "from_email": msg_data["from_email"], "from_name": msg_data["from_name"], "to_emails": msg_data["to_emails"], "cc_emails": msg_data["cc_emails"], "subject": msg_data["subject"], "snippet": msg_data["snippet"], "internal_date": msg_data["internal_date"], "classification": msg_data["classification"], "classification_reason": msg_data["classification_reason"], "bypassed_wo": msg_data["bypassed_wo"], "notified": msg_data.get("notified", False), "classified_date": msg_data["classified_date"], "slack_ts": msg_data.get("slack_ts", ""), "ttl": _ttl_90_days(), } try: _get_table().put_item( Item=item, ConditionExpression="attribute_not_exists(pk)", ) except ClientError as e: if e.response["Error"]["Code"] == "ConditionalCheckFailedException": return raise def save_thread_state(thread_id, message, adam_email): now_iso = datetime.now(timezone.utc).isoformat() is_adam = message["from_email"] == adam_email if is_adam: _get_table().update_item( Key={"pk": f"THD#{thread_id}", "sk": "THD"}, UpdateExpression=( "SET last_message_from = :from_email, " "last_message_date = :now, " "subject = :subject, " "adam_last_reply_date = :now, " "ttl = :ttl " "REMOVE unanswered_since" ), ExpressionAttributeValues={ ":from_email": message["from_email"], ":now": now_iso, ":subject": message["subject"], ":ttl": _ttl_90_days(), }, ) else: _get_table().update_item( Key={"pk": f"THD#{thread_id}", "sk": "THD"}, UpdateExpression=( "SET last_message_from = :from_email, " "last_message_date = :now, " "subject = :subject, " "unanswered_since = if_not_exists(unanswered_since, :now), " "ttl = :ttl" ), ExpressionAttributeValues={ ":from_email": message["from_email"], ":now": now_iso, ":subject": message["subject"], ":ttl": _ttl_90_days(), }, ) def get_sync_metadata(): resp = _get_table().get_item( Key={"pk": "META#sync", "sk": "META"} ) return resp.get("Item") def save_sync_metadata(history_id): _get_table().put_item( Item={ "pk": "META#sync", "sk": "META", "history_id": str(history_id), "last_sync_at": datetime.now(timezone.utc).isoformat(), "initial_sync_done": True, } ) def get_todays_messages(date_str): items = [] resp = _get_table().query( IndexName="by-date", KeyConditionExpression="classified_date = :d AND sk = :sk", ExpressionAttributeValues={":d": date_str, ":sk": "MSG"}, ) items.extend(resp.get("Items", [])) while "LastEvaluatedKey" in resp: resp = _get_table().query( IndexName="by-date", KeyConditionExpression="classified_date = :d AND sk = :sk", ExpressionAttributeValues={":d": date_str, ":sk": "MSG"}, ExclusiveStartKey=resp["LastEvaluatedKey"], ) items.extend(resp.get("Items", [])) return items def get_conversation(thread_ts): resp = _get_table().get_item( Key={"pk": f"CONV#{thread_ts}", "sk": "CONV"} ) item = resp.get("Item") return item["messages"] if item else [] def save_conversation(thread_ts, messages): _get_table().put_item(Item={ "pk": f"CONV#{thread_ts}", "sk": "CONV", "messages": messages, "updated_at": datetime.now(timezone.utc).isoformat(), "ttl": _ttl_90_days(), }) def get_unanswered_threads(threshold_hours): cutoff = datetime.now(timezone.utc).timestamp() - (threshold_hours * 3600) items = [] scan_kwargs = { "FilterExpression": "begins_with(pk, :prefix) AND attribute_exists(unanswered_since)", "ExpressionAttributeValues": {":prefix": "THD#"}, } resp = _get_table().scan(**scan_kwargs) items.extend(resp.get("Items", [])) while "LastEvaluatedKey" in resp: scan_kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"] resp = _get_table().scan(**scan_kwargs) items.extend(resp.get("Items", [])) return [ item for item in items if datetime.fromisoformat(item["unanswered_since"]).timestamp() < cutoff ]