(feat): Insertion to mongoDB done

utilsv2/mongo.py

@@ -0,0 +1,71 @@
+import logging
+from typing import Any
+
+from pymongo import MongoClient
+from pymongo.collection import Collection
+from pymongo.errors import ConnectionFailure
+
+try:
+    from utilsv2.log import logger
+    from utilsv2.parser import massage_magnitudes
+except ModuleNotFoundError:
+    from log import logger
+    from parser import massage_magnitudes
+
+logger = logging.getLogger(__name__)
+
+
+def connect(uri) -> MongoClient:
+    try:
+        client = MongoClient(uri)
+        logger.info("Connected to the DB")
+    except ConnectionFailure as e:
+        logger.critical("Could not connect to the MongoDB")
+        raise e
+
+    return client
+
+
+def add_events(
+    client: MongoClient, collection: str, data: list[dict[str, Any]]
+) -> None:
+    db = client["main"]
+    coll = db[collection]
+
+    data = massage_magnitudes(data)
+
+    _res = coll.insert_many(data)
+
+    if _res.acknowledged:
+        logger.info(f"Added {len(_res.inserted_ids)} events.")
+    else:
+        logger.info("Could not add events to the database.")
+
+
+def add_stations(
+    client: MongoClient, collection: str, data: list[dict[str, Any]]
+) -> None:
+    db = client["main"]
+    coll = db[collection]
+
+    _res = coll.insert_many(data)
+
+    if _res.acknowledged:
+        logger.info(f"Added {len(_res.inserted_ids)} events.")
+    else:
+        logger.info("Could not add events to the database.")
+
+
+def get_ids(collection: Collection) -> set[Any]:
+    return set(collection.distinct("ID"))
+
+
+def close(client: MongoClient) -> None:
+    client.close()
+    logger.info("Closed the DB.")
+
+
+if __name__ == "__main__":
+    v = connect("mongodb://localhost:27017")
+
+    close(v)