browsertrix/backend/db.py
Ilya Kreymer bcbc40059e
Refactor backend data model to support UUID (fixes #118) (#119)
* uuid fix: (fixes #118)
- update all mongo models to use UUID type as main '_id' (users continue to use 'id' as defined by fastapi-users)
- update all foreign doc references to use UUID instead of string
- api handlers convert str->uuid as needed
api fix:
- fix single crawl api, add CrawlOut response model
- fix collections api
- fix standalone-docker apis
- for manual job, set user to current user, overriding the setting from crawlconfig

* additional fixes:
- rename username -> userName to indicate not the login 'username'
- rename user -> userid, archive -> aid for crawlconfig + crawls
- ensure invites correctly convert str -> uuid as needed
- filter out unset values from browsertrix-crawler config

* convert remaining user -> userid variables
ensure archive id is passed to crawl_manager as str (via archive.id_str)

* remove bulk crawlconfig delete
* add support for `stopping` state when gracefully stopping crawl
* for get crawl endpoint, check stopped crawls first, then running
2022-01-29 19:00:11 -08:00

60 lines
1.4 KiB
Python

"""
Browsertrix API Mongo DB initialization
"""
import os
from typing import Optional
import motor.motor_asyncio
from pydantic import BaseModel, UUID4
DATABASE_URL = (
f"mongodb://root:example@{os.environ.get('MONGO_HOST', 'localhost')}:27017"
)
# ============================================================================
def init_db():
"""initializde the mongodb connector"""
client = motor.motor_asyncio.AsyncIOMotorClient(
DATABASE_URL, uuidRepresentation="standard"
)
mdb = client["browsertrixcloud"]
return mdb
# ============================================================================
class BaseMongoModel(BaseModel):
"""Base pydantic model that is also a mongo doc"""
id: Optional[UUID4]
@property
def id_str(self):
""" Return id as str """
return str(self.id)
@classmethod
def from_dict(cls, data):
"""convert dict from mongo to an Archive"""
if not data:
return None
data["id"] = data.pop("_id")
return cls(**data)
def serialize(self, **opts):
"""convert Archive to dict"""
return self.dict(
exclude_unset=True, exclude_defaults=True, exclude_none=True, **opts
)
def to_dict(self, **opts):
"""convert to dict for mongo"""
res = self.dict(**opts)
res["_id"] = res.pop("id", "")
return res