Mais coisas, estatisticas

utilsv2/mongo.py

@@ -14,7 +14,7 @@ except ModuleNotFoundError:
 logger = logging.getLogger(__name__)
 
 
-def connect(uri) -> MongoClient:
+def connect(uri: str) -> MongoClient:
     try:
         client = MongoClient(uri)
         logger.info("Connected to the DB")
@@ -26,10 +26,9 @@ def connect(uri) -> MongoClient:
 
 
 def add_events(
-    client: MongoClient, collection: str, data: list[dict[str, Any]]
+    client: MongoClient, collection: str, data: list[dict[str, Any]], db: str = "main"
 ) -> None:
-    db = client["main"]
-    coll = db[collection]
+    coll: Collection = client[db][collection]
 
     _res = coll.insert_many(data)
 
@@ -40,10 +39,9 @@ def add_events(
 
 
 def add_stations(
-    client: MongoClient, collection: str, data: list[dict[str, Any]]
+    client: MongoClient, collection: str, data: list[dict[str, Any]], db: str = "main"
 ) -> None:
-    db = client["main"]
-    coll = db[collection]
+    coll: Collection = client[db][collection]
 
     _res = coll.insert_many(data)
 
@@ -62,19 +60,28 @@ def close(client: MongoClient) -> None:
     logger.info("Closed the DB.")
 
 
-def query_all(cli: MongoClient, collection: str) -> Any:
-    coll: Collection = cli.main[collection]
+def query_all(client: MongoClient, collection: str, db: str = "main") -> Any:
+    coll: Collection = client[db][collection]
 
     result = coll.find({})
 
-    for doc in result:
-        print(doc)
+    return list(result)
 
 
-def filter_query(cli: MongoClient, collection: str, filter_by):
-    coll: Collection = cli.main[collection]
+def filter_query(
+    client: MongoClient, collection: str, filter_by: dict[str, Any], db: str = "main"
+):
+    coll: Collection = client[db][collection]
 
-    res = coll.find({""})
+    res = coll.find(
+        filter_by, {"DateTime": 1, "Magnitudes": 1, "Depth": 1, "GAP": 1}
+    ).sort({"DateTime": 1})
+
+    if not res._empty:
+        res = list(res)
+        logger.info(f"Retrieved {len(res)} elements.")
+
+    return res
 
 
 if __name__ == "__main__":