diff --git a/api/main.py b/api/main.py index 5c0d448..317cbf9 100644 --- a/api/main.py +++ b/api/main.py @@ -60,7 +60,7 @@ from utils_v2.api.async_quart import ( ) # GMail-related utils: -from utils_v2.goog.controllers.gmail.gmail_client import AsyncGMailClient +from utils_v2.goog.controllers.gmail.gmail_client import AsyncGmailClient # Core Controller Models: from controllers.core.message import CoreMessageController @@ -534,13 +534,13 @@ async def app_startup(**kwargs): # ┗┛┗┛┛┗┛┗┗ ┗┗┗┛┛ ┛ ┗┻┛┗┗┻ ┗┛┗┗┗ ┛┗┗┛ # Create an instance to handle GMail-related activities: - current_app.gmail_client = AsyncGMailClient( + current_app.gmail_client = AsyncGmailClient( service_name = "gmail", oauth_json = script_cred["google"]["oauth"]["tcaoff"], http_client = current_app.http_client, redirect_url = r"https://api.thecaoffice.com/converse/mail/callback/gmail", debug = enable_debugging, - debug_prefix = "GMail (M) | ", + debug_prefix = "Gmail (M) | ", debug_only_errors = False ) diff --git a/controllers_v2/core/auth_token.py b/controllers_v2/core/auth_token.py index eea1659..01bcf36 100644 --- a/controllers_v2/core/auth_token.py +++ b/controllers_v2/core/auth_token.py @@ -21,7 +21,7 @@ N/A """ - +import datetime # ***************************************************************************************************************** # ***** **** # *** IMPORT *** @@ -55,6 +55,7 @@ import httpx # To work with MongoDB: from bson import ObjectId +from pymongo import UpdateOne # ***************************************************************************************************************** @@ -571,6 +572,117 @@ class CoreAuthTokenController(CoreBaseModel): # Done here: return [CoreAuthTokenModel(**token) for token in tokens] + async def get_batches_to_sync( + self, + mongo_data_conn: AsyncMongo, + limit: int = 25, + batch_timeout_seconds: int | float = 120, + additional_filter: dict = None + ) -> List[CoreAuthTokenModel]: + + """ + Use this when some kind of account that has been integrated needs to perform data synchronization at periodic + intervals. One simple example is mails. The aim is to pick batches of auth-tokens + :param mongo_data_conn: The database connection (MongoDB) to use to perform the action. + :param limit: The max. no. of records to pick. + :param batch_timeout_seconds: The no. of seconds to not pick an element that has been marked as a member of + another batch. + :param additional_filter: Any addition filters to use. + :return: The retrieved record that has the token, and information about the service and client if found, else + None when there is no matching record. + """ + + # Prepare the filter: + now = date_time.get_current_utc_date_time(as_string = False) + batch_timeout_ts = now - datetime.timedelta(seconds = int(batch_timeout_seconds)) + filter_json = { + "$and": [ + { + "$or": [ + {"syncAfterTs": None}, # ............ The time after which sync'ing is allowed is null. + {"syncAfterTs": {"$lte": now}}, # ... The time after which sync'ing is allowed has passed. + ] + }, + { + "$or": [ + {"batchId": None}, # ....................... Not currently a member of a batch. + {"batchTs": {"$lt": batch_timeout_ts}} # ... Is a member of a batch, but it has timed-out. + ] + } + ] + } + if self._base_filter: + for k, v in self._base_filter.items(): filter_json[k] = v + if additional_filter: + for k, v in additional_filter.items(): filter_json[k] = v + + # Generate a batch-id: + batch_id = ObjectId() + + # fetch the credentials that match the conditions: + tokens = await mongo_data_conn.find_many( + collection = self.AUTH_COLLECTION, + filter = filter_json, + sort = {"lastSyncTs": 1}, + limit = limit, + projection = {"_id": True} + ) + token_ids = [t["_id"] for t in tokens] + + # Mark 'n' no. of auth-tokens as a member of your batch: + await mongo_data_conn.update_many( + collection = self.AUTH_COLLECTION, + filter = {"_id": {"$in": token_ids}}, + update = {"$set": {"batchId": batch_id, "batchTs": now}}, + upsert = False + ) + + # Fetch the credentials to use for this batch: + tokens = await mongo_data_conn.find_many( + collection = self.AUTH_COLLECTION, + filter = {"batchId": batch_id}, + limit = len(token_ids) + ) + + # Done here: + return [CoreAuthTokenModel(**token) for token in tokens] + + async def release_token_from_batch_by_id( + self, + mongo_data_conn: AsyncMongo, + token_id: ObjectId | str, + sync_after_ts: datetime.datetime | None, + last_sync_ts: datetime.datetime | None = None + ) -> bool: + + """ + To release a particular auth-tokn from a sync'ing batch. This helps in gracefully completing the cycle. + :param mongo_data_conn: The database connection (MongoDB) to use to perform the action. + :param token_id: The ObjectId of the document in MongoDb that holds the auth-token. + :param sync_after_ts: The time (UTC) after which this same token can be picked again for sync'ing. + :param last_sync_ts: The time at which this latest sync'ing was done. Leave it to null for the current time, + else pass a custom timestamp. + :return: True if successful, else False. + """ + + # Process the inputs: + if isinstance(last_sync_ts, datetime.datetime): date_time.to_timezone(last_sync_ts, date_time.TIMEZONE_UTC) + else: last_sync_ts = date_time.get_current_utc_date_time(as_string = False) + + # release the auth-token from the batch: + return await mongo_data_conn.update_one( + collection = self.AUTH_COLLECTION, + filter = {"_id": ObjectId(token_id)}, + update = { + "$set": { + "batchId": None, + "syncAfterTs": sync_after_ts, + "lastSyncTs": last_sync_ts + } + }, + upsert = False + ) + # ┏┓┳┓┳┳┳┓ ┳┓ ┓ # ┃ ┣┫┃┃┃┃ ━━ ┃┃┏┓┃┏┓╋┏┓ # ┗┛┛┗┗┛┻┛ ┻┛┗ ┗┗ ┗┗ diff --git a/controllers_v2/message/mail/all_mail.py b/controllers_v2/message/mail/all_mail.py index 874d8e0..463e185 100644 --- a/controllers_v2/message/mail/all_mail.py +++ b/controllers_v2/message/mail/all_mail.py @@ -58,7 +58,7 @@ from models.message.mail.sync import MailSyncOneResult, MailSyncManyResults from models.message.mail.send import MailSendOneResult # Mail Client(s): -from utils_v2.goog.controllers.gmail.gmail_client import AsyncGMailClient +from utils_v2.goog.controllers.gmail.gmail_client import AsyncGmailClient # To work with datatypes: from typing import List, Any @@ -138,9 +138,9 @@ class AllMailController(MailController): """ # Prepare the combined base filter: - sms_filter = {} - for k, v in (base_filter or {}).items(): sms_filter[k] = v - sms_filter["serviceType"] = "sms" + mail_filter = {} + for k, v in (base_filter or {}).items(): mail_filter[k] = v + mail_filter["serviceType"] = "email" # Invoke the parent's constructor: MailController.__init__( @@ -148,7 +148,7 @@ class AllMailController(MailController): cache = cache, alert_url = alert_url, http_client = http_client, - base_filter = sms_filter, + base_filter = mail_filter, debug = debug, debug_prefix = debug_prefix, debug_only_errors = debug_only_errors @@ -169,7 +169,7 @@ class AllMailController(MailController): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, user_info: CoreUserInfoModel, inbound_data: OAuthMailAuthorizationRequestData, session_token: str @@ -193,7 +193,7 @@ class AllMailController(MailController): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, request_url: str, inbound_data: dict, session_token: str = None @@ -216,7 +216,7 @@ class AllMailController(MailController): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, http_client: httpx.AsyncClient, auth_token: CoreAuthTokenModel, force_refresh: bool = False, @@ -257,7 +257,7 @@ class AllMailController(MailController): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, auth_token: CoreAuthTokenModel, user_info: CoreUserInfoModel | None, llm: CoreLLMController = None, diff --git a/controllers_v2/message/mail/base.py b/controllers_v2/message/mail/base.py index 6483d0c..2c483c9 100644 --- a/controllers_v2/message/mail/base.py +++ b/controllers_v2/message/mail/base.py @@ -58,7 +58,7 @@ from models.message.mail.sync import MailSyncOneResult, MailSyncManyResults from models.message.mail.send import MailSendOneResult # Mail Client(s): -from utils_v2.goog.controllers.gmail.gmail_client import AsyncGMailClient +from utils_v2.goog.controllers.gmail.gmail_client import AsyncGmailClient # To work with datatypes: from typing import List, Any @@ -190,7 +190,7 @@ class MailController(CoreMessageController, ABC): for k, v in (base_filter or {}).items(): sms_filter[k] = v sms_filter["serviceType"] = "email" - # Invoke the parent's constructor: + # Invoke the parents' constructor: CoreMessageController.__init__( self, cache = cache, @@ -300,7 +300,7 @@ class MailController(CoreMessageController, ABC): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, user_info: CoreUserInfoModel, inbound_data: OAuthMailAuthorizationRequestData, session_token: str @@ -325,7 +325,7 @@ class MailController(CoreMessageController, ABC): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, request_url: str, inbound_data: dict, session_token: str = None @@ -349,7 +349,7 @@ class MailController(CoreMessageController, ABC): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, http_client: httpx.AsyncClient, auth_token: CoreAuthTokenModel, force_refresh: bool = False, @@ -426,7 +426,7 @@ class MailController(CoreMessageController, ABC): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, auth_token: CoreAuthTokenModel, user_info: CoreUserInfoModel | None, llm: CoreLLMController = None, diff --git a/controllers_v2/message/mail/gmail.py b/controllers_v2/message/mail/gmail.py index dd85fec..0a153aa 100644 --- a/controllers_v2/message/mail/gmail.py +++ b/controllers_v2/message/mail/gmail.py @@ -57,7 +57,7 @@ from models.message.mail.sync import MailSyncOneResult, MailSyncManyResults from models.message.mail.send import MailSendOneResult # Mail Client(s): -from utils_v2.goog.controllers.gmail.gmail_client import AsyncGMailClient, SCOPES_GMAIL_MAIL_MANAGEMENT +from utils_v2.goog.controllers.gmail.gmail_client import AsyncGmailClient, SCOPES_GMAIL_MAIL_MANAGEMENT from utils_v2.goog.models.auth_tokens import GoogleAuthTokens # To work with datatypes: @@ -163,7 +163,7 @@ class GmailController(MailController): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, user_info: CoreUserInfoModel, inbound_data: OAuthMailAuthorizationRequestData, session_token: str @@ -230,7 +230,7 @@ class GmailController(MailController): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, request_url: str, inbound_data: dict, session_token: str = None @@ -368,7 +368,7 @@ class GmailController(MailController): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, http_client: httpx.AsyncClient, auth_token: CoreAuthTokenModel, force_refresh: bool = False, @@ -451,7 +451,7 @@ class GmailController(MailController): mongo_data_conn: AsyncMongo, user_info: CoreUserInfoModel, auth_token: CoreAuthTokenModel, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, google_tokens: GoogleAuthTokens, message_id: str, llm: CoreLLMController = None, @@ -495,6 +495,7 @@ class GmailController(MailController): f"Gmail message '{message_id}' already " f"sync'd on {mail_records[0].syncTs} (UTC)." ) + print(sync_result) return sync_result # Now that we know that we have to fetch the mail from GMail: @@ -561,7 +562,7 @@ class GmailController(MailController): self, sql_conn: AsyncMySQL, mongo_data_conn: AsyncMongo, - mail_client: AsyncGMailClient, + mail_client: AsyncGmailClient, auth_token: CoreAuthTokenModel, user_info: CoreUserInfoModel | None, llm: CoreLLMController = None, diff --git a/cron/message/mail/auto_sync.py b/cron/message/mail/auto_sync.py index d7f51f9..a97b164 100644 --- a/cron/message/mail/auto_sync.py +++ b/cron/message/mail/auto_sync.py @@ -6,12 +6,12 @@ DATE: - Thursday, 2nd Jan., 2025 + Friday, 17th Jan., 2025. OBJECTIVE: - To broadcast live tick updates to connected clients. It doesn't matter which stockbroker we are getting the - ticks from as long as we are reading standardized ticks from the Kafka queue. + To automatically synchronize the mails by fetching data from the user's third-party mail client's server to your + database. This is done in batches of auth-tokens. REFERENCES: @@ -39,6 +39,7 @@ sys.path.append("..") # System-level activities: import io import os +import socket # My utils: from utils_v2.string import json @@ -46,10 +47,19 @@ from utils_v2.string import regex from utils_v2.system import files from utils_v2.date_time import date_time from utils_v2.database.async_mongo_v2 import AsyncMongo +from utils_v2.database.async_mysql_v2 import AsyncMySQL from utils_v2.queue.kafka.controllers.async_kafka import ConsumerKafka, get_ssl_context from utils_v2.cache.async_redis_cache_v2 import AsyncRedisCache from utils_v2.serialization.json_serializer import JSONSerializer +# Controllers: +from controllers_v2.message.mail.all_mail import AllMailController +from controllers_v2.message.mail.gmail import GmailController +from controllers_v2.core.ai.llm import CoreLLMController + +# Mail clients: +from utils_v2.goog.controllers.gmail.gmail_client import AsyncGmailClient + # To make HTTP calls: import httpx @@ -58,11 +68,9 @@ import datetime import time # Models: +from models.core.auth_token import CoreAuthTokenModel from models.core.user import CoreUserInfoModel - -# To work with SocketIO: -import socket -import socketio +from models.message.mail.sync import MailSyncOneResult, MailSyncManyResults # For asynchronous activities: import asyncio @@ -82,8 +90,8 @@ from icecream import IceCreamDebugger # Debugging: -printer = IceCreamDebugger(prefix = "Tick-Out | ", includeContext = True) -no_context_printer = IceCreamDebugger(prefix = "Tick-Out | ", includeContext = False) +printer = IceCreamDebugger(prefix = "Mail Sync. | ", includeContext = True) +no_context_printer = IceCreamDebugger(prefix = "Mail Sync. | ", includeContext = False) # To make API calls: http_client = httpx.AsyncClient( @@ -102,16 +110,6 @@ http_client = httpx.AsyncClient( # General: SERVER_HOSTNAME = str(socket.gethostname()) -# For SocketIO: -# Namespaces: -NAMESPACE_MODULE = None -NAMESPACE_PASSTHROUGH = "/passthrough" -# Events: -EVENT_CONNECT = "connect" -EVENT_DISCONNECT = "disconnect" -EVENT_ECHO = "echo" -EVENT_TICKS = "ticks" - # ***************************************************************************************************************** # ***** **** @@ -120,27 +118,24 @@ EVENT_TICKS = "ticks" # ***************************************************************************************************************** -# For SocketIO: -ALLOWED_ORIGINS = [] -sio = socketio.AsyncServer( - cors_allowed_origins = "*", - async_mode = "asgi" -) -app = socketio.ASGIApp(sio) - -# Redis: -redis_cache: AsyncRedisCache | None = None - -# For kafka: -kafka_consumer: ConsumerKafka | None = None - # Session-awareness and maintenance of this script's state: SCRIPT_DATA = {} exclusive_lock = asyncio.Semaphore(1) -CONNECTED_CLIENTS = {} -FLAGS = { - "initDone": False -} + +# For databases: +data_mongo: AsyncMongo | None = None +sql_writer: AsyncMySQL | None = None + +# For caching: +redis_cache: AsyncRedisCache | None = None + +# Controllers: +mail_controller: AllMailController | None = None +gmail_controller: GmailController | None = None +llm_controller: CoreLLMController | None = None + +# Mail clients: +gmail_client: AsyncGmailClient | None = None # ***************************************************************************************************************** @@ -150,113 +145,6 @@ FLAGS = { # ***************************************************************************************************************** -def origin_is_allowed(origin: str) -> bool: - - """ - To check if a given origin is in the allowed list. - :param origin: The origin of your request. - :return: True if allowed, else False. - """ - - # Start by assuming failure: - is_allowed = False - - # Check through all the allowed origins: - for allowed in ALLOWED_ORIGINS: - try: - if regex.match(origin, allowed): - is_allowed = True - break - except Exception as exception: - printer(exception) - - # Done here: - return is_allowed - - -# --------------------------------------------------------------------------------------------------------------------- - - -async def send_ticks(ticks: List[dict]) -> None: - - """ - Here's where we decide which client gets which tick and send it out. - WARNING: WE ARE ASSUMING THAT NO FURTHER FORMATING/COMPUTATION IS REQUIRED OTHER THAN SELECTING WHICH SUBSETS OF - TICKS TO SEND TO WHICH CLIENTS. FOR US THE TICKS ALREADY HAVE ALL THE DATA NEEDED TO BE SEND TO - RESPECTIVE CLIENTS. - :param ticks: The list of individual tick updates to send out to the clients. - :return: None - """ - - # Currently we're just broadcasting - # all the data to all the clients: - await sio.emit( - event = EVENT_TICKS, - data = ticks, - namespace = NAMESPACE_MODULE - ) - - -# --------------------------------------------------------------------------------------------------------------------- - - -async def ticks_from_kafka( - consumer: ConsumerKafka, - fetch_count: int = 100, - fetch_timeout: float = 1.0 -) -> None: - - """ - This function must run in the background forever and just keep listening for ticks on Kafka and keep relaying them - to all the connected clients as per their watchlists. - :param consumer: The preconfigured Kafka consumer that can listen for ticks in asynchronous mode. - :param fetch_count: How many messages to consume in one go. - :param fetch_timeout: How long to wait (in seconds) while consuming messages from Kafka. - :return: None - """ - - printer("Starting Kafka consumer (ticks).") - - # Do the next part infinitely: - while True: - - # Note the time: - now_utc = date_time.get_current_utc_date_time().timestamp() - - # Get messages form Kafka: - ticks = await consumer.consume( - count = fetch_count, - timeout = fetch_timeout - ) - - # If there are no updates to give: - if not ticks: continue - - # Each message must be treated as an array of tick updates (list of dicts). - # In case the producer is sending each individual tick as a separate message, - # we normalize it to be a list: - tasks = [send_ticks(t.value if isinstance(t.value, list) else [t.value]) for t in ticks] - results = await asyncio.gather(*tasks) - - # Analyze the ticks: - # latency = [abs(now_utc - t.value.get("rcvdTs", t.value["tradeTs"])) for t in ticks] - latency = [abs(now_utc - t.ts.timestamp()) for t in ticks] - avg_latency = sum(latency) / len(latency) - total_ticks = len(ticks) - # late_cutoff_seconds = 3.0 - # - # late_ticks = 0 - # for tick in ticks: - # if now_utc - tick.value["tradeTs"] > late_cutoff_seconds: - # late_ticks += 1 - # ticks_str = f"COUNT: {total_ticks: >5,} | LATE: {late_ticks: >5,} ({(late_ticks/total_ticks)*100.0:.2f}%)" - ticks_str = f"COUNT: {total_ticks: >5,} | AVG. LATENCY: {avg_latency:.5f}" - no_context_printer(ticks_str) - - -# --------------------------------------------------------------------------------------------------------------------- - - async def init( script_id: str, debug: bool @@ -271,36 +159,17 @@ async def init( # Declare the required global variables: global SCRIPT_DATA - global ALLOWED_ORIGINS + global data_mongo + global sql_writer global redis_cache - global kafka_consumer + global mail_controller + global gmail_controller + global llm_controller + global gmail_client # Basic stuff: if debug: printer.enable() - printer("Initializing.") - - # ┏┓ • • - # ┃┃┏┓┓┏┓┓┏┓┏ - # ┗┛┛ ┗┗┫┗┛┗┛ - # ┛ - - response = await http_client.post( - url = r"https://api.thecaoffice.com/ca/get/title", - headers = {"Origin": "https://thecaoffice.com/"}, - data = { - "domainName": "127.0.0.1:1234", - "screenWidth": 1920, - "screenHeight": 1080 - } - ) - if response.status_code not in [200]: - print("FATAL: ALLOWED ORIGINS NOT FETCHED!") - return False - ALLOWED_ORIGINS = [origin["domain"] for origin in response.json().get("data", {}).get("rs2", [])] - printer(ALLOWED_ORIGINS) - if len(ALLOWED_ORIGINS) < 1: - print("FATAL: ALLOWED ORIGINS IS EMPTY!") - return False + no_context_printer("Initializing.") # ┏┓ ┓ ┓ ┳┓ # ┃ ┏┓┏┓┏┫ ┏┓┏┓┏┫ ┃┃┏┓╋┏┓ @@ -327,58 +196,105 @@ async def init( SCRIPT_DATA = response.json().get("data") # Done with this step: - printer("Cred and Data loaded.") + no_context_printer("Cred and Data loaded.") - # ┓┏┓ ┏┓ ┏┓┓• - # ┃┫ ┏┓╋┃┏┏┓ ┃ ┃┓┏┓┏┓╋┏ - # ┛┗┛┗┻┛┛┗┗┻ ┗┛┗┗┗ ┛┗┗┛ + # ┳┳┓ • ┳┓┳┓ + # ┃┃┃┏┓┏┓┓┏┓┃┃┣┫ + # ┛ ┗┗┻┛ ┗┗┻┻┛┻┛ - # Create the consumer that will listen to changes in watchlist: - consumer_creds = script_cred["kafka"]["consumer"] - kafka_consumer = ConsumerKafka( - topic = consumer_creds["topic"], - bootstrap_servers = consumer_creds["config"]["bootstrapServers"], - security_protocol = consumer_creds["config"].get("securityProtocol", "PLAINTEXT"), - ssl_context = get_ssl_context( - ca_file = consumer_creds["config"].get("caFile"), - cert_file = consumer_creds["config"].get("certFile"), - key_file = consumer_creds["config"].get("keyFile"), - ), - serializer = JSONSerializer(), + sql_writer = AsyncMySQL( + pool_size = script_cred["mariaDb"]["write"]["poolSize"], + host = script_cred["mariaDb"]["write"]["host"], + user = script_cred["mariaDb"]["write"]["user"], + password = script_cred["mariaDb"]["write"]["password"], + database = script_cred["mariaDb"]["write"]["database"] + ) + if not await sql_writer.connect(): + print("FATAL: MARIA-DB NOT CONNECTED!") + return False + no_context_printer("MariaDB ready.") + + # ┳┳┓ + # ┃┃┃┏┓┏┓┏┓┏┓ + # ┛ ┗┗┛┛┗┗┫┗┛ + # ┛ + + data_mongo = AsyncMongo( + connection_string = script_cred["mongoDb"]["data"]["connectionString"], + database_name = script_cred["mongoDb"]["data"]["dbName"], + max_connections = script_cred["mongoDb"]["data"]["poolSize"], debug = debug ) - if not await kafka_consumer.connect(): - print("FATAL: KAFKA CONSUMER NOT CREATED!") + if not await data_mongo.connect(): + print("FATAL: MONGO-DB NOT CONNECTED!") return False - printer("Kafka consumer ready.") + no_context_printer("MongoDB ready.") # ┳┓ ┓• ┏┓ ┓ # ┣┫┏┓┏┫┓┏ ━━ ┃ ┏┓┏┣┓┏┓ # ┛┗┗ ┗┻┗┛ ┗┛┗┻┗┛┗┗ redis_cache = AsyncRedisCache( - connection_string = script_cred["redisCache"]["general"]["connectionString"], + connection_string = script_cred["redisCache"]["funcReturn"]["connectionString"], serializer = JSONSerializer(), debug = debug, debug_prefix = "General Cache | " ) if not await redis_cache.connect(): - print("FATAL: REDIS CACHE NOT CREATED!") + print("FATAL: REDIS CACHE NOT CONNECTED!") return False - printer("Redis cache ready.") + no_context_printer("Redis cache ready.") - # ┳┓ ┓ ┓ ┏┳┓ ┓ - # ┣┫┏┓┏┃┏┏┓┏┓┏┓┓┏┏┓┏┫ ┃ ┏┓┏┃┏┏ - # ┻┛┗┻┗┛┗┗┫┛ ┗┛┗┻┛┗┗┻ ┻ ┗┻┛┛┗┛ - # ┛ + # ┏┓ ┓┓ ┓┏┏┓ + # ┃ ┏┓┏┓╋┏┓┏┓┃┃┏┓┏┓┏ ┃┃┏┛ + # ┗┛┗┛┛┗┗┛ ┗┛┗┗┗ ┛ ┛ ┗┛┗━ - # Start the background task that will receive ticks from the Kafka queue and broadcast them to the respective - # connected clients: - sio.start_background_task( - ticks_from_kafka, - consumer = kafka_consumer, - fetch_count = 1_000, - fetch_timeout = 1.0 + # Messages / Mail Controllers: + mail_controller = AllMailController( + cache = redis_cache, + http_client = http_client, + alert_url = SCRIPT_DATA["alerts"]["url"], + debug = False + ) + gmail_controller = GmailController( + cache = redis_cache, + http_client = http_client, + alert_url = SCRIPT_DATA["alerts"]["url"], + debug = False + ) + + # ┏┓ ┓ ┏┓┓• + # ┃ ┏┓┏┓┏┓┏┓┏╋┏┓┏┓┏ ┏┓┏┓┏┫ ┃ ┃┓┏┓┏┓╋┏ + # ┗┛┗┛┛┗┛┗┗ ┗┗┗┛┛ ┛ ┗┻┛┗┗┻ ┗┛┗┗┗ ┛┗┗┛ + + # Create an instance to handle GMail-related activities: + gmail_client = AsyncGmailClient( + service_name = "gmail", + oauth_json = script_cred["google"]["oauth"]["tcaoff"], + http_client = http_client, + redirect_url = r"https://api.thecaoffice.com/converse/mail/callback/gmail", + debug = False, + debug_prefix = "Gmail (M) | ", + debug_only_errors = True + ) + + # ┏┓┳ ┳┳┓ • + # ┣┫┃ ┃┃┃┏┓┏┓┓┏ + # ┛┗┻ ┛ ┗┗┻┗┫┗┗ + # ┛ + + # For LLMs: + llm_controller = CoreLLMController( + llm_creds = { + "model": script_cred["openAi"]["model"], + "openai_api_key": script_cred["openAi"]["openai_api_key"] + }, + cache = redis_cache, + alert_url = SCRIPT_DATA["alerts"]["url"], + http_client = http_client, + debug = debug, + debug_prefix = "AI (LLM) | ", + debug_only_errors = True ) # ┳┓ @@ -386,103 +302,101 @@ async def init( # ┻┛┗┛┛┗┗ # If everything went well, we return with success: - printer("Initialization done.") + no_context_printer("Initialization done.") return True # --------------------------------------------------------------------------------------------------------------------- -@sio.on(event = EVENT_CONNECT, namespace = NAMESPACE_MODULE) -async def on_connect(sid, environ, *args) -> bool: +async def sync_one_account(auth_token: CoreAuthTokenModel): """ - The event handler for when a new connection request comes in. - :param sid: The session id of the incoming request (generated by SocketIO). - :param environ: The set of headers and other connection-specific values. - :param args: Any extra input coming from the connection request. - :return: True to accept a connection request, False to reject it. + To sync one third-party mail account. + :param auth_token: The auth-token that gives access to the mail account. + :return: ?? """ - # declare the required global variables: - global CONNECTED_CLIENTS + # printer("Sync'ing account.", auth_token.clientUserId["email"], auth_token.authTokenId) - # Initialize the script if needed: - async with exclusive_lock: - if not FLAGS.get("initDone"): - FLAGS["initDone"] = await init( - script_id = os.environ["SCRIPT_ID"], - debug = True if os.environ["DEBUG"].lower() == "true" else False - ) + # Note down the time at which the attempt to sync the account is being made: + now = date_time.get_current_utc_date_time(as_string = False) - # If the initialization failed, we cannot accept the incoming request: - if not FLAGS.get("initDone"): - printer("SOCKET REJECTED: Init. pending.", sid) - return False + # Start by assuming failure: + client_controller = None + client_connector = None + sync_results = MailSyncManyResults() - # Check the origin of the incoming request: - printer("Checking origin.") - origin = environ.get("HTTP_ORIGIN", "???") - if not origin_is_allowed(origin): - printer("SOCKET REJECTED: Bad origin.", sid, origin) - return False + # Figure out which mail client has to be used: + match auth_token.client: + case "gmail": client_controller, client_connector = gmail_controller, gmail_client + case _: client_controller, client_connector = None, None - # Get the session token from the incoming request: - printer("Checking session token.") - session_token = environ.get("HTTP_X_SESSION_TOKEN") - if not session_token and len(args) > 0: session_token = args[0].get("X-Session-Token") - if not session_token: - printer("SOCKET REJECTED: No session token.", sid) - return False + # Try to sync mails: + if client_controller is not None and client_connector is not None: + try: sync_results = await client_controller.sync_mails( + sql_conn = sql_writer, + mongo_data_conn = data_mongo, + mail_client = client_connector, + auth_token = auth_token, + user_info = None, + llm = llm_controller, + force_sync = False, + start_date = now - datetime.timedelta(days = 1), + end_date = now, + max_count = 100, + session_token = None + ) + except Exception as exception: + pass + # printer(exception) + # else: printer("Invalid/unimplemented client.", auth_token.client) - # Get the user's details from the session token: - printer("Fetching user info.") - user_info = await redis_cache.get(key = session_token) - if not user_info: - printer("SOCKET REJECTED: Invalid session token.", sid) - return False - user_info = CoreUserInfoModel(**user_info) - - # Get the user's watchlist and note down the details. - # Consider the following structure for a user's info: - redis_key = f"io_{session_token}" - async with exclusive_lock: - CONNECTED_CLIENTS[sid] = { - "user": user_info, - "redisKey": redis_key, - "rooms": [] - } - sid_cached = await redis_cache.set( - key = redis_key, - value = {"server": SERVER_HOSTNAME, "socket_id": sid} + # Release the auth-token from the batch: + await mail_controller.release_token_from_batch_by_id( + mongo_data_conn = data_mongo, + token_id = auth_token.authTokenId, + # sync_after_ts = now + datetime.timedelta(seconds = auth_token.syncFreq or 300), + sync_after_ts = now + datetime.timedelta(seconds = 1), + last_sync_ts = now ) # Done here: - printer("SOCKET ACCEPTED.", sid, sid_cached) - return True + if sync_results.totalCount: printer(auth_token.clientUserId["email"], sync_results.successCount, sync_results.totalCount) # --------------------------------------------------------------------------------------------------------------------- -@sio.on(event = EVENT_DISCONNECT, namespace = NAMESPACE_MODULE) -async def handle_disconnect(sid, reason) -> None: +async def sync_accounts(batch_size: int) -> None: """ - To handle a disconnect event. Automatically triggered when a client disconnects from the server. - :param sid: The session id of the client (generated by SocketIO on connecting). - :param reason: The hint about why the disconnection happened. + To go into an indefinite loop and keep sync'ing mails for many accounts. + :param batch_size: The no. of accounts to pick in every batch. :return: None. """ - # declare the required global variables: - global CONNECTED_CLIENTS + while True: - # register the disconnect in the global variable, and on the cache server: - client_info = {} - async with exclusive_lock: client_info = CONNECTED_CLIENTS.pop(sid, None) - sid_uncached = await redis_cache.delete(key = client_info["redisKey"]) if client_info else False - printer("SOCKET DISCONNECTED", sid, reason, sid_uncached) + # Get batches of auth-tokens to work with: + # no_context_printer("Getting batch.") + # no_context_printer(batch_size) + auth_tokens_batch = await mail_controller.get_batches_to_sync( + mongo_data_conn = data_mongo, + limit = batch_size, + batch_timeout_seconds = 10 + ) + # no_context_printer(len(auth_tokens_batch)) + + # Process each batch: + now = date_time.get_current_utc_date_time(as_string = False) + tasks = [sync_one_account(auth_token = auth_token) for auth_token in auth_tokens_batch] + results = await asyncio.gather(*tasks) + + # Small delay to not overload the database: + # if len(auth_tokens_batch) == 0: await asyncio.sleep(5.0) + await asyncio.sleep(1.0) + # no_context_printer() # ***************************************************************************************************************** @@ -494,34 +408,16 @@ async def handle_disconnect(sid, reason) -> None: if __name__ == "__main__": - printer("Main.") - # To get args from the terminal: import argparse - # To run the ASGI: - import uvicorn - from multiprocessing import freeze_support - # Get the config from the command-line: - parser = argparse.ArgumentParser(description = f"SocketIO to serve live market data (and a general passthrough).") + parser = argparse.ArgumentParser(description = f"To automatically sync. the mails for all users.") parser.add_argument( - "-w", "--workers", + "-b", "--batch-size", type = int, - help = "The no. of threads to spin up for this instance!", - default = 2 - ) - parser.add_argument( - "-a", "--host", - type = str, - help = "The host for the app. e.g.: '0.0.0.0' or '127.0.0.1'.", - default = "127.0.0.1" - ) - parser.add_argument( - "-p", "--port", - type = int, - help = "The port no. to bind the app to.", - default = 8080 + help = "How many auth-tokens to load at once to sync.", + default = 25 ) parser.add_argument( "-s", "--script-id", @@ -536,21 +432,18 @@ if __name__ == "__main__": ) args = parser.parse_args() - # Note down the config; - os.environ["SCRIPT_ID"] = args.script_id - os.environ["DEBUG"] = str(args.debug) - # Startup message: printer.enable() - printer(str(args.debug)) + debugging_enabled = args.debug + printer(debugging_enabled) printer.disable() - # Run the gateway: - freeze_support() - uvicorn.run( - app = "tick_out:app", - workers = args.workers, - host = args.host, - port = args.port - ) + async def main(): + if await init( + script_id = args.script_id, + debug = args.debug + ): await sync_accounts( + batch_size = args.batch_size + ) + asyncio.run(main()) diff --git a/models/core/auth_token.py b/models/core/auth_token.py index 164cade..25ae7ec 100644 --- a/models/core/auth_token.py +++ b/models/core/auth_token.py @@ -179,12 +179,30 @@ class CoreAuthTokenModel(BaseModel): default = None ) + syncAfterTs: AwareDatetime | None = Field( + description = "the time (utc) after which this account becomes eligible for re-sync'ing", + frozen = False, + default = True + ) + lastSyncTs: AwareDatetime | None = Field( description = "the time at which this client's updates were last polled", frozen = False, default = None ) + batchId: ObjectId | None = Field( + description = "a batch id assigned when creds are picked for account data sync'ing", + frozen = False, + default = None + ) + + batchTs: AwareDatetime | None = Field( + description = "the time (utc) at which this record was picked for data sync'ing", + frozen = False, + default = None + ) + # ┏┓ ┏• # ┃ ┏┓┏┓╋┓┏┓ # ┗┛┗┛┛┗┛┗┗┫ @@ -204,6 +222,7 @@ class CoreAuthTokenModel(BaseModel): @field_validator( "firstRequestTs", "lastRequestTs", "firstRefreshTs", "lastRefreshTs", + "batchTs", "lastSyncTs", "syncAfterTs", mode = "before" ) def parse_date_time(cls, value):