summaryrefslogtreecommitdiff
path: root/components/database.py
blob: a65a4c56f68f40354c937de7f22c5c472edcfb85 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import atexit
from os import getenv
from typing import Any, Dict
from dotenv import load_dotenv
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

load_dotenv('.env')

client: MongoClient[Any] = MongoClient(
    "mongodb://%s:%s@localhost/admin" % (getenv('MONGO_USER'), getenv('MONGO_PASS')),
    tz_aware=True,
)
database: Database[Any] = client.get_database(getenv('YT_DB') or "youtube")
subscriptions: Collection[SubsDict] = database.get_collection("subscriptions")
users: Collection[UserDict] = database.get_collection("users")

@atexit.register
def _cleanup() -> None:
    client.close()