"""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}