Source code for util.mongodb

from config import Config
from typing import Any




class MongoDbUtil:
    """Utilities work working mongodb."""



    @staticmethod
    def get_discovery_db_connection():
        """Get connection to discovery database.
        :return: connection string"""

        from urllib.parse import quote_plus, urlunsplit
        from airflow.models.connection import Connection

        conn = Connection.get_connection_from_secrets(Config.RAD_DOC_DB_CONN)
        print("Retrieving connection", conn.conn_id)

        scheme = "mongodb"
        login = conn.login
        password = conn.password
        netloc = conn.host
        if login is not None and password is not None:
            netloc = f"{quote_plus(login)}:{quote_plus(password)}@{netloc}"
        if conn.port:
            netloc = f"{netloc}:{conn.port}"
        path = f"/{conn.schema}"
        connection_str = urlunsplit((scheme, netloc, path, "directConnection=true", ""))
        print(f"Connecting to {connection_str}")
        return connection_str




    @staticmethod
    def mongodb_get_collection(mongo_collection: str, mongo_db: str | None = None):
        """
        Get a collection from a mongo database.
        param mongo_collection: name of the collection
        param mongo_db: name of the database
        return: collection object

        """
        from pymongo import MongoClient

        conn = MongoDbUtil.get_discovery_db_connection()

        mongo_client: MongoClient = MongoClient(conn)
        print(mongo_client.server_info())
        if mongo_db:
            return mongo_client.get_database(mongo_db).get_collection(mongo_collection)

        return mongo_client.get_default_database().get_collection(mongo_collection)




    @staticmethod
    def mongodb_find(
        mongo_collection: str,
        query: dict,
        find_one: bool = False,
        mongo_db: str | None = None,
    ) -> Any | None:
        """
        Run a mongo find query and returns the results.
        param mongo_collection: name of the collection
        param query: query to run
        param find_one: if True, return only the first result
        param mongo_db: name of the database

        https://pymongo.readthedocs.io/en/stable/api/pymongo/collection.html#pymongo.collection.Collection.find
        """
        collection = MongoDbUtil.mongodb_get_collection(
            mongo_collection, mongo_db=mongo_db
        )

        if find_one:
            return collection.find_one(query)
        else:
            return collection.find(query)




    @staticmethod
    def mongodb_insert_one(
        mongo_collection: str, doc: dict, mongo_db: str | None = None, **kwargs
    ) -> Any:
        """
        Insert a single document into a mongo collection.

        https://pymongo.readthedocs.io/en/stable/api/pymongo/collection.html#pymongo.collection.Collection.insert_one
        """
        collection = MongoDbUtil.mongodb_get_collection(
            mongo_collection, mongo_db=mongo_db
        )
        result = collection.insert_one(doc, **kwargs)

        return result.inserted_id