(20250117) Day-end push.

This commit is contained in:
2025-01-17 18:57:46 +05:30
parent 8ab63d036d
commit 4b2cdbee76
7 changed files with 356 additions and 331 deletions
+3 -3
View File
@@ -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
)
+113 -1
View File
@@ -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
)
# ┏┓┳┓┳┳┳┓ ┳┓ ┓
# ┃ ┣┫┃┃┃┃ ━━ ┃┃┏┓┃┏┓╋┏┓
# ┗┛┛┗┗┛┻┛ ┻┛┗ ┗┗ ┗┗
+9 -9
View File
@@ -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,
+6 -6
View File
@@ -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,
+7 -6
View File
@@ -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,
+199 -306
View File
@@ -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.")
# ┳┓ ┓ ┏┳┓ ┓
# ┣┫┏┓┏┃┏┏┓┏┓┏┓┓┏┓┏ ┃ ┏┓┏┃┏
# ┻┛┗┻┗┛┗┗┛ ┗┛┗┻┛┗┗┻ ┻ ┗┻┛┛┗
# ┏┓ ┓┓ ┓┏┏
# ┃ ┏┓┏┓┏┓┏┓┃┃┏┓┏┓┏ ┃┏
# ┗┛┗┛┗┗┛ ┗┛┗┗┗ ┛ ┛ ┗┛┗
# 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
)
# ┏┓┳ ┳┳┓ •
# ┣┫┃ ┃┃┃┏┓┏┓┓┏
# ┛┗┻ ┛ ┗┗┻┗┫┗┗
# ┛
# 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
# 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)
# Start by assuming failure:
client_controller = None
client_connector = None
sync_results = MailSyncManyResults()
# 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
# 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)
# If the initialization failed, we cannot accept the incoming request:
if not FLAGS.get("initDone"):
printer("SOCKET REJECTED: Init. pending.", sid)
return False
# 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
# 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
# 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())
+19
View File
@@ -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):