Spaces:
Sleeping
Sleeping
File size: 6,867 Bytes
b60402f |
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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 |
"""
Database utilities for ClipboardHealthAI application.
This module provides database-related utilities, including:
- Custom Pydantic types for MongoDB ObjectID handling
- Base model classes for MongoDB document models with serialization support
- Background task for periodically refreshing Snowflake database connections
"""
from datetime import datetime
from enum import Enum
from typing import Any, Dict, Type
from bson import ObjectId
from pydantic import AnyUrl, BaseModel, Field, GetCoreSchemaHandler
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema
class PyObjectId:
"""
Custom type for handling MongoDB ObjectId in Pydantic models.
This class provides validation and serialization for MongoDB's ObjectId,
making it compatible with Pydantic model validation and JSON serialization.
"""
@classmethod
def __get_pydantic_core_schema__(
cls, _source: type, _handler: GetCoreSchemaHandler
) -> core_schema.CoreSchema:
"""
Define the core schema for Pydantic validation.
Args:
source: The source type
handler: Schema handler instance
Returns:
CoreSchema: The schema for validation
"""
return core_schema.with_info_after_validator_function(
cls.validate, core_schema.str_schema() # type: ignore
)
@classmethod
def __get_pydantic_json_schema__(
cls, _schema: core_schema.CoreSchema, _handler: GetCoreSchemaHandler
) -> JsonSchemaValue:
"""
Define the JSON schema representation.
Args:
schema: The core schema
handler: Schema handler instance
Returns:
JsonSchemaValue: The JSON schema representation
"""
return {"type": "string"}
@classmethod
def validate(cls, value: str) -> ObjectId:
"""
Validate and convert a string to MongoDB ObjectId.
Args:
value: String representation of ObjectId
Returns:
ObjectId: MongoDB ObjectId instance
Raises:
ValueError: If the value is not a valid ObjectId
"""
if not ObjectId.is_valid(value):
raise ValueError(f"Invalid ObjectId: {value}")
return ObjectId(value)
def __getattr__(self, item):
"""
Delegate attribute access to the wrapped ObjectId.
Args:
item: The attribute name
Returns:
The attribute value from the wrapped ObjectId
"""
return getattr(self.__dict__["value"], item)
def __init__(self, value: str | None = None):
"""
Initialize with a string value or create a new ObjectId.
Args:
value: Optional string representation of ObjectId
"""
if value is None:
self.value = ObjectId()
else:
self.value = self.validate(value)
def __str__(self):
"""
Convert to string representation.
Returns:
str: String representation of the ObjectId
"""
return str(self.value)
class MongoBaseModel(BaseModel):
"""
Base model for MongoDB documents with serialization support.
This class extends Pydantic's BaseModel to provide MongoDB-specific
serialization/deserialization and handling of special types.
"""
id: str = Field(default_factory=lambda: str(PyObjectId()))
class Config: # pylint: disable=R0903
"""
Configuration for the model.
"""
arbitrary_types_allowed = True
def to_mongo(self) -> Dict[str, Any]:
"""
Convert the model instance to a MongoDB-compatible dictionary.
Handles special types like nested models, enums, datetimes, and URLs.
Returns:
Dict[str, Any]: A dictionary suitable for MongoDB storage
"""
def model_to_dict(model: BaseModel) -> Dict[str, Any]:
doc = {}
for name, value in model._iter(): # pylint: disable=W0212
key = model.__fields__[name].alias or name
if isinstance(value, BaseModel):
doc[key] = model_to_dict(value)
elif isinstance(value, list) and all(isinstance(i, BaseModel) for i in value):
doc[key] = [model_to_dict(item) for item in value] # type: ignore
elif value and isinstance(value, Enum):
doc[key] = value.value
elif isinstance(value, datetime):
doc[key] = value.isoformat() # type: ignore
elif value and isinstance(value, AnyUrl):
doc[key] = str(value) # type: ignore
else:
doc[key] = value
return doc
result = model_to_dict(self)
return result
@classmethod
def from_mongo(cls, data: Dict[str, Any]):
"""
Create a model instance from MongoDB document data.
Handles special types conversion, particularly for enum values.
Args:
data: Dictionary containing MongoDB document data
Returns:
MongoBaseModel: An instance of the model class
"""
def restore_enums(inst: Any, model_cls: Type[BaseModel]) -> None:
for name, field in model_cls.__fields__.items(): # type: ignore
value = getattr(inst, name)
if (
field
and isinstance(field.annotation, type)
and issubclass(field.annotation, Enum)
):
setattr(inst, name, field.annotation(value))
elif isinstance(value, BaseModel):
restore_enums(value, value.__class__)
elif isinstance(value, list):
for i, item in enumerate(value):
if isinstance(item, BaseModel):
restore_enums(item, item.__class__)
elif isinstance(field.annotation, type) and issubclass(
field.annotation, Enum
):
value[i] = field.annotation(item)
elif isinstance(value, dict):
for k, v in value.items():
if isinstance(v, BaseModel):
restore_enums(v, v.__class__)
elif isinstance(field.annotation, type) and issubclass(
field.annotation, Enum
):
value[k] = field.annotation(v)
if data is None:
return None
instance = cls(**data)
restore_enums(instance, instance.__class__)
return instance
|