blob: 3b16ef53111bee9830cf996fed01948d0e09b8cc (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
|
"""Our database module"""
from time import sleep
from sys import exit as app_exit
from dataclasses import dataclass
from motor.motor_asyncio import (
AsyncIOMotorClient,
AsyncIOMotorCollection,
)
from bson import ObjectId
@dataclass()
class DBClient:
"""Class to hold database connections for us."""
client: AsyncIOMotorClient
collection: AsyncIOMotorCollection
def __init__(self, username: str, password: str, collection: str) -> None:
self.client = AsyncIOMotorClient(f"mongodb://{username}:{password}@mongodb:27017/production", timeoutMS=2000)
self.collection = self.client["production"][collection]
async def check_server(self) -> None:
"""Try query the DB and exit the program if we fail after 5 times.
:return: None
"""
for i in range(5):
try:
await self.collection.find_one({"_id": ObjectId("507f1f77bcf86cd799439011")})
print("Connection to DB - OK")
break
except: # pylint: disable=bare-except
print(f"WARNING failed to connect to DB - {i} / 4", flush=True)
sleep(1)
else:
print("Could not connect to DB - mongodb://REDACTED_USERNAME:REDACTED_PASSWORD@mongodb:27017/production")
app_exit(1)
|
