38 lines
1 KiB
Python
38 lines
1 KiB
Python
"""JARVIS Core Database utilities."""
|
|
from datetime import datetime, timezone
|
|
from functools import partial
|
|
from typing import Any
|
|
|
|
import nanoid
|
|
from beanie import Document
|
|
from pydantic import Field, GetCoreSchemaHandler
|
|
from pydantic_core import CoreSchema, core_schema
|
|
|
|
NANOID_ALPHA = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
|
|
|
|
|
def get_now() -> datetime:
|
|
"""Get proper timestamp."""
|
|
return datetime.now(tz=timezone.utc)
|
|
|
|
|
|
def get_id() -> str:
|
|
"""Get nanoid."""
|
|
return nanoid.generate(NANOID_ALPHA, 12)
|
|
|
|
|
|
NowField = partial(Field, default_factory=get_now)
|
|
NanoField = partial(Field, default_factory=get_id)
|
|
|
|
|
|
class Snowflake(int):
|
|
@classmethod
|
|
def __get_pydantic_core_schema__(
|
|
cls, source_type: Any, handler: GetCoreSchemaHandler
|
|
) -> CoreSchema:
|
|
return core_schema.no_info_after_validator_function(cls, handler(int))
|
|
|
|
|
|
class SnowflakeDocument(Document):
|
|
class Settings:
|
|
bson_encoders = {Snowflake: str}
|