Add database setup code

1 files changed, 16 insertions, 0 deletions

diff --git a/data-collection/components/database.py b/data-collection/components/database.py
new file mode 100644
index 0000000..72f63e6
--- /dev/null
+++ b/data-collection/components/database.py
@@ -0,0 +1,16 @@
+import atexit
+from typing import Any, Dict
+from pymongo import MongoClient
+from pymongo.database import Database
+from pymongo.collection import Collection
+from components.subscriptions.typing import SubsDict
+from components.users.typing import UserDict
+
+client: MongoClient[Any] = MongoClient("mongodb://localhost", tz_aware=True)
+database: Database[Any] = client.get_database("youtube")
+subscriptions: Collection[SubsDict] = database.get_collection("subscriptions")
+users: Collection[UserDict] = database.get_collection("users")
+
+@atexit.register
+def _cleanup() -> None:
+    client.close()