import typing from datetime import datetime, timedelta from typing import Dict, List, Tuple from motor.motor_asyncio import AsyncIOMotorClient from pymongo import ReplaceOne from pymongo.errors import BulkWriteError, DuplicateKeyError from sentry_sdk import capture_exception import config import context from server.core.cache.utils import is_image_caching_collection from server.core.constants import StorageField from server.core.utils import compress, count_time, decompress mongo_client: AsyncIOMotorClient = None def init() -> AsyncIOMotorClient: """ AWS DocDB reference Download CA wget https://s3.amazonaws.com/rds-downloads/rds-combined-ca-bundle.pem mongodb://{user}:{password}@{host}:{port}/{db}?ssl=true&ssl_ca_certs=rds-combined-ca-bundle.pem&replicaSet=rs0&readPreference=secondaryPreferred&retryWrites=false """ additional_kwargs = {} if config.MONGODB_ENABLE_SSL: additional_kwargs = { "ssl": True, "ssl_ca_certs": "./rds-combined-ca-bundle.pem", "replicaSet": "rs0", "readPreference": "secondaryPreferred", } if config.MONGODB_WAIT_WRITE_REPLICAS: if config.MONGODB_WAIT_JOURNAL: additional_kwargs["journal"] = config.MONGODB_WAIT_JOURNAL elif config.MONGODB_WAIT_FSYNC: additional_kwargs["fsync"] = config.MONGODB_WAIT_FSYNC global mongo_client # see all the arguments at https://pymongo.readthedocs.io/en/stable/api/pymongo/mongo_client.html mongo_client = AsyncIOMotorClient( host=config.MONGODB_HOST, port=config.MONGODB_PORT, authSource=config.MONGODB_DATABASE, username=config.MONGODB_USER, password=config.MONGODB_PASSWORD, authMechanism=config.MONGODB_AUTH_MECHANISM, maxPoolSize=config.MONGODB_MAX_POOL_SIZE, maxIdleTimeMS=config.MONGODB_CONNECTION_IDLE_TIME_MS, socketTimeoutMS=config.MONGODB_SOCKET_TIMEOUT_MS, connectTimeoutMS=config.MONGODB_CONNECTION_TIMEOUT_MS, serverSelectionTimeoutMS=config.MONGODB_SELECT_TIMEOUT_MS, waitQueueTimeoutMS=config.MONGODB_WAIT_QUEUE_TIMEOUT_MS, waitQueueMultiple=config.MONGODB_WAIT_QUEUE_MULTIPLE, retryReads=config.MONGODB_RETRY_READS, retryWrites=config.MONGODB_RETRY_WRITES, w=config.MONGODB_WAIT_WRITE_REPLICAS, wTimeoutMS=config.MONGODB_WRITE_TIMEOUT_MS, **additional_kwargs, ) return mongo_client async def check_health() -> typing.Dict[str, typing.Any]: """Invokes healthcheck.""" return await mongo_client.server_info() async def close() -> None: """Closes cache client session.""" return await mongo_client.close() @count_time(context.CACHE_TIME) async def get_documents( collection_name: str, keys: List[str], timeout: int or None, by_record_id: bool ) -> Tuple[Dict[str, dict], datetime or None, datetime or None]: """Get cached data from MongoDB. Args: collection_name: Collection name. keys: Collection items ID list. timeout: Cached data considered as actual if it was saved to mongo less then this interval (in seconds) ago, 0 or None - do not use cache, -1 - ignore timeout. by_record_id: Return results mapping by record ID or by object ID. Returns: ID to data mapping and datetime stats. """ if not timeout: return {}, None, None filter_condition = {StorageField.RECORD_ID: {"$in": keys}} if timeout > 0: conditions = [filter_condition] timeout_condition = {StorageField.CREATED_AT: {"$gt": datetime.utcnow() - timedelta(seconds=timeout)}} if is_image_caching_collection(collection_name) and timeout > config.TEMP_IMAGE_MAX_TTL: conditions.append( { "$or": [ {"$and": [{StorageField.IMAGES_SAVED: True}, timeout_condition]}, { StorageField.CREATED_AT: { "$gt": datetime.utcnow() - timedelta(seconds=config.TEMP_IMAGE_MAX_TTL) } }, ] } ) else: conditions.append(timeout_condition) filter_condition = {"$and": conditions} try: mongo_records = ( await mongo_client[config.MONGODB_DATABASE][collection_name] .find( filter_condition, {(StorageField.OBJECT_ID if by_record_id else StorageField.RECORD_ID): 0, StorageField.IMAGES_SAVED: 0}, ) .to_list(None) ) key_field = StorageField.RECORD_ID if by_record_id else StorageField.OBJECT_ID dates = [r[StorageField.CREATED_AT] for r in mongo_records] result = { r[key_field]: decompress(r.get(StorageField.COMPRESSION), r[StorageField.DATA]) for r in mongo_records } return result, min(dates) if dates else None, max(dates) if dates else None except Exception as e: capture_exception(e) return {}, None, None async def save_documents( collection_name: str, records: Dict[str, Tuple[str, dict] or Tuple[int, dict]], created_at: datetime ): """Save data to MongoDB. Args: collection_name: Collection name. records: Record ID to object ID and data mapping. created_at: Records created_at datetime. """ try: await mongo_client[config.MONGODB_DATABASE][collection_name].bulk_write( [ ReplaceOne( {StorageField.RECORD_ID: key}, { StorageField.RECORD_ID: key, StorageField.OBJECT_ID: value[0], StorageField.DATA: compress(config.MONGODB_COMPRESSION_LIB, value[1]), StorageField.CREATED_AT: created_at, StorageField.COMPRESSION: config.MONGODB_COMPRESSION_LIB, }, upsert=True, ) for key, value in records.items() ], ordered=False, ) except BulkWriteError as e: if "writeErrors" not in (e.details or {}) or any(i for i in e.details["writeErrors"] if i.get("code") != 11000): capture_exception(e) except Exception as e: capture_exception(e) async def save_document(collection_name: str, record_id: str, data: Dict, created_at: datetime): """Save single record to MongoDB. Args: collection_name: Collection name. record_id: Record ID. data: Object data. created_at: Record created_at datetime. """ try: await mongo_client[config.MONGODB_DATABASE][collection_name].replace_one( {StorageField.RECORD_ID: record_id}, { StorageField.RECORD_ID: record_id, StorageField.DATA: compress(config.MONGODB_COMPRESSION_LIB, data), StorageField.CREATED_AT: created_at, StorageField.COMPRESSION: config.MONGODB_COMPRESSION_LIB, }, upsert=True, ) except Exception as e: if not isinstance(e, DuplicateKeyError): capture_exception(e)