feat: migrate backend persistence from TinyDB to MongoDB Atlas
- Implement lazy MongoDB client initialization in backend/db/mongo_client.py. - Create Gunicorn configuration to ensure MongoDB client is initialized per worker. - Refactor database access layer to support MongoDB with a new MongoLockedTable adapter. - Add migration script to transfer existing TinyDB data to MongoDB, preserving idempotency. - Update tracking event handling to ensure deterministic ordering with a monotonic sequence. - Modify tests to use mongomock for MongoDB integration and ensure existing tests pass. - Add integration test script to run tests against a local MongoDB Docker container. - Document environment variables and migration process in specs/feat-database-migration.md.
This commit is contained in:
+493
-40
@@ -1,13 +1,42 @@
|
||||
# python
|
||||
import os
|
||||
from config.paths import get_database_dir
|
||||
import threading
|
||||
from config.paths import get_database_dir
|
||||
from tinydb import TinyDB
|
||||
from tinydb.queries import QueryInstance
|
||||
|
||||
from db.mongo_client import get_mongo_client, get_mongo_db_name
|
||||
|
||||
try:
|
||||
from tinydb.table import Document
|
||||
except ImportError: # pragma: no cover - tinydb version compatibility
|
||||
from tinydb.database import Document
|
||||
|
||||
|
||||
def _stable_clause_key(clause: dict) -> str:
|
||||
"""Return a stable string key for sorting MongoDB filter clauses."""
|
||||
import json
|
||||
return json.dumps(clause, sort_keys=True, default=str)
|
||||
|
||||
|
||||
try:
|
||||
from pymongo import ASCENDING
|
||||
except ImportError: # pragma: no cover - pymongo is a required dependency
|
||||
ASCENDING = 1
|
||||
|
||||
|
||||
USE_MONGODB = os.environ.get('USE_MONGODB', 'false').lower() == 'true'
|
||||
# Resolve the MongoDB database name once at module load so runtime changes to
|
||||
# DB_ENV/DATA_ENV in tests do not switch databases mid-process.
|
||||
_mongo_db_name = get_mongo_db_name() if USE_MONGODB else None
|
||||
base_dir = get_database_dir()
|
||||
os.makedirs(base_dir, exist_ok=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TinyDB-backed table wrapper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class LockedTable:
|
||||
"""
|
||||
Thread-safe wrapper around a TinyDB table. All callable attribute access
|
||||
@@ -65,6 +94,361 @@ class LockedTable:
|
||||
with self._lock:
|
||||
return self._table.truncate()
|
||||
|
||||
def close(self):
|
||||
with self._lock:
|
||||
return self._table.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TinyDB Query -> MongoDB filter translator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_MONGO_OP_MAP = {
|
||||
'==': '$eq',
|
||||
'!=': '$ne',
|
||||
'<': '$lt',
|
||||
'<=': '$lte',
|
||||
'>': '$gt',
|
||||
'>=': '$gte',
|
||||
}
|
||||
|
||||
_NEGATED_OPS = {
|
||||
'$eq': '$ne',
|
||||
'$ne': '$eq',
|
||||
'$lt': '$gte',
|
||||
'$lte': '$gt',
|
||||
'$gt': '$lte',
|
||||
'$gte': '$lt',
|
||||
}
|
||||
|
||||
|
||||
def _field_to_mongo(field_path: tuple) -> str:
|
||||
"""Map a TinyDB field path to a MongoDB field name.
|
||||
|
||||
The model ``id`` field is stored as the MongoDB ``_id`` field, so queries
|
||||
on ``id`` are translated to queries on ``_id``.
|
||||
"""
|
||||
if len(field_path) == 1:
|
||||
return '_id' if field_path[0] == 'id' else field_path[0]
|
||||
return '.'.join('_id' if p == 'id' else p for p in field_path)
|
||||
|
||||
|
||||
def _negate_condition(cond: dict) -> dict | None:
|
||||
"""Return a MongoDB condition that negates a single-field condition."""
|
||||
if len(cond) != 1:
|
||||
return None
|
||||
field, inner = next(iter(cond.items()))
|
||||
if not isinstance(inner, dict) or len(inner) != 1:
|
||||
return None
|
||||
op, value = next(iter(inner.items()))
|
||||
if op in _NEGATED_OPS:
|
||||
return {field: {_NEGATED_OPS[op]: value}}
|
||||
return None
|
||||
|
||||
|
||||
def _hash_to_mongo_filter(query_hash) -> dict | None:
|
||||
"""Translate a TinyDB query hash tuple to a MongoDB filter document.
|
||||
|
||||
Returns ``None`` when the query construct cannot be expressed as a native
|
||||
MongoDB filter, signalling that the caller should fall back to in-memory
|
||||
TinyDB evaluation.
|
||||
"""
|
||||
if not isinstance(query_hash, tuple) or len(query_hash) == 0:
|
||||
return None
|
||||
|
||||
op = query_hash[0]
|
||||
|
||||
if op in _MONGO_OP_MAP:
|
||||
field = _field_to_mongo(query_hash[1])
|
||||
value = query_hash[2]
|
||||
return {field: {_MONGO_OP_MAP[op]: value}}
|
||||
|
||||
if op == 'exists':
|
||||
field = _field_to_mongo(query_hash[1])
|
||||
return {field: {'$exists': True}}
|
||||
|
||||
if op == 'one_of':
|
||||
field = _field_to_mongo(query_hash[1])
|
||||
return {field: {'$in': list(query_hash[2])}}
|
||||
|
||||
if op == 'any':
|
||||
field = _field_to_mongo(query_hash[1])
|
||||
return {field: {'$in': list(query_hash[2])}}
|
||||
|
||||
if op == 'all':
|
||||
field = _field_to_mongo(query_hash[1])
|
||||
return {field: {'$all': list(query_hash[2])}}
|
||||
|
||||
if op == 'matches':
|
||||
field = _field_to_mongo(query_hash[1])
|
||||
# TinyDB matches() anchors the regex at the start of the string.
|
||||
return {field: {'$regex': f'^{query_hash[2]}'}}
|
||||
|
||||
if op == 'search':
|
||||
field = _field_to_mongo(query_hash[1])
|
||||
return {field: {'$regex': query_hash[2]}}
|
||||
|
||||
if op == 'and':
|
||||
merged: dict = {}
|
||||
for sub_hash in query_hash[1]:
|
||||
sub = _hash_to_mongo_filter(sub_hash)
|
||||
if sub is None:
|
||||
return None
|
||||
for field, inner in sub.items():
|
||||
if field in merged:
|
||||
if isinstance(merged[field], dict) and isinstance(inner, dict):
|
||||
merged[field].update(inner)
|
||||
elif isinstance(merged[field], list) and isinstance(inner, list):
|
||||
merged[field].extend(inner)
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
merged[field] = (
|
||||
dict(inner) if isinstance(inner, dict)
|
||||
else list(inner) if isinstance(inner, list)
|
||||
else inner
|
||||
)
|
||||
return merged
|
||||
|
||||
if op == 'or':
|
||||
clauses = [_hash_to_mongo_filter(sub_hash) for sub_hash in query_hash[1]]
|
||||
if any(c is None for c in clauses):
|
||||
return None
|
||||
return {'$or': sorted(clauses, key=_stable_clause_key)}
|
||||
|
||||
if op == 'not':
|
||||
inner = _hash_to_mongo_filter(query_hash[1])
|
||||
if inner is None:
|
||||
return None
|
||||
negated = _negate_condition(inner)
|
||||
if negated is not None:
|
||||
return negated
|
||||
return None
|
||||
|
||||
# Unsupported operation (test, fragment, noop, etc.) -> fall back.
|
||||
return None
|
||||
|
||||
|
||||
def _query_to_mongo_filter(query) -> dict | None:
|
||||
"""Translate a TinyDB QueryInstance to a MongoDB filter, if possible."""
|
||||
if isinstance(query, dict):
|
||||
return query
|
||||
if isinstance(query, QueryInstance):
|
||||
return _hash_to_mongo_filter(query._hash)
|
||||
return None
|
||||
|
||||
|
||||
def _evaluate_in_memory(docs, query) -> list:
|
||||
"""Evaluate a TinyDB query against an in-memory list of documents."""
|
||||
if callable(query):
|
||||
return [doc for doc in docs if query(doc)]
|
||||
return docs
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MongoDB-backed table wrapper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class MongoLockedTable:
|
||||
"""Drop-in replacement for ``LockedTable`` that delegates to MongoDB.
|
||||
|
||||
The adapter preserves the existing table API while mapping the model
|
||||
``id`` field to MongoDB's ``_id`` field on reads and writes.
|
||||
"""
|
||||
|
||||
def __init__(self, collection_name: str):
|
||||
self.collection_name = collection_name
|
||||
|
||||
def _collection(self):
|
||||
client = get_mongo_client()
|
||||
return client[_mongo_db_name][self.collection_name]
|
||||
|
||||
@staticmethod
|
||||
def _doc_to_mongo(doc: dict) -> dict:
|
||||
"""Store a copy of ``doc`` with ``id`` promoted to MongoDB ``_id``.
|
||||
|
||||
This avoids storing both ``_id`` and ``id`` with identical values.
|
||||
The original ``id`` field is removed from the stored document.
|
||||
"""
|
||||
if doc is None:
|
||||
return None
|
||||
d = dict(doc)
|
||||
if 'id' in d:
|
||||
d['_id'] = d.pop('id')
|
||||
return d
|
||||
|
||||
@staticmethod
|
||||
def _doc_from_mongo(doc: dict):
|
||||
"""Return a TinyDB-compatible Document with ``doc_id`` set to ``_id``.
|
||||
|
||||
Restores the model ``id`` field from MongoDB's ``_id`` and exposes
|
||||
TinyDB's ``doc_id`` attribute so callers that rely on it continue to
|
||||
work.
|
||||
"""
|
||||
if doc is None:
|
||||
return None
|
||||
d = dict(doc)
|
||||
doc_id = d.pop('_id', None)
|
||||
if doc_id is not None:
|
||||
d['id'] = doc_id
|
||||
return Document(d, doc_id=doc_id)
|
||||
|
||||
def _mongo_filter(self, cond):
|
||||
"""Translate a TinyDB query or dict to a MongoDB filter."""
|
||||
return _query_to_mongo_filter(cond)
|
||||
|
||||
def all(self):
|
||||
return [self._doc_from_mongo(doc) for doc in self._collection().find({})]
|
||||
|
||||
def search(self, cond):
|
||||
mongo_filter = self._mongo_filter(cond)
|
||||
if mongo_filter is not None:
|
||||
cursor = self._collection().find(mongo_filter)
|
||||
return [self._doc_from_mongo(doc) for doc in cursor]
|
||||
|
||||
# Fallback: fetch all and evaluate the TinyDB query in Python.
|
||||
docs = list(self._collection().find({}))
|
||||
matched = _evaluate_in_memory(
|
||||
[self._doc_from_mongo(doc) for doc in docs], cond
|
||||
)
|
||||
return matched
|
||||
|
||||
def get(self, cond):
|
||||
mongo_filter = self._mongo_filter(cond)
|
||||
if mongo_filter is not None:
|
||||
doc = self._collection().find_one(mongo_filter)
|
||||
return self._doc_from_mongo(doc)
|
||||
|
||||
docs = list(self._collection().find({}))
|
||||
for doc in docs:
|
||||
d = self._doc_from_mongo(doc)
|
||||
if callable(cond) and cond(d):
|
||||
return d
|
||||
return None
|
||||
|
||||
def insert(self, document: dict):
|
||||
doc = self._doc_to_mongo(document)
|
||||
result = self._collection().insert_one(doc)
|
||||
return str(result.inserted_id)
|
||||
|
||||
def insert_multiple(self, documents: list):
|
||||
if not documents:
|
||||
return []
|
||||
docs = [self._doc_to_mongo(d) for d in documents]
|
||||
result = self._collection().insert_many(docs)
|
||||
return [str(iid) for iid in result.inserted_ids]
|
||||
|
||||
def update(self, fields, cond=None, doc_ids=None):
|
||||
is_callable = callable(fields)
|
||||
|
||||
if doc_ids is not None:
|
||||
mongo_filter = {'_id': {'$in': list(doc_ids)}}
|
||||
target_ids = [str(did) for did in doc_ids]
|
||||
if not target_ids:
|
||||
return []
|
||||
|
||||
if is_callable:
|
||||
# Fetch, apply callable in-memory, and replace each document.
|
||||
updated_ids = []
|
||||
for doc in self._collection().find(mongo_filter):
|
||||
d = self._doc_from_mongo(doc)
|
||||
fields(d)
|
||||
new_doc = self._doc_to_mongo(d)
|
||||
new_doc.pop('_id', None)
|
||||
self._collection().update_one(
|
||||
{'_id': doc['_id']}, {'$set': new_doc}
|
||||
)
|
||||
updated_ids.append(str(doc['_id']))
|
||||
return updated_ids
|
||||
|
||||
update_doc = self._doc_to_mongo(fields) or {}
|
||||
update_doc.pop('_id', None)
|
||||
update_doc.pop('id', None)
|
||||
if update_doc:
|
||||
self._collection().update_many(mongo_filter, {'$set': update_doc})
|
||||
return target_ids
|
||||
|
||||
mongo_filter = self._mongo_filter(cond)
|
||||
if mongo_filter is not None and not is_callable:
|
||||
update_doc = self._doc_to_mongo(fields) or {}
|
||||
update_doc.pop('_id', None)
|
||||
update_doc.pop('id', None)
|
||||
target_ids = [
|
||||
str(doc['_id'])
|
||||
for doc in self._collection().find(mongo_filter, {'_id': 1})
|
||||
]
|
||||
if target_ids and update_doc:
|
||||
self._collection().update_many(
|
||||
mongo_filter, {'$set': update_doc}
|
||||
)
|
||||
return target_ids
|
||||
|
||||
# Fallback: evaluate the query in-memory and update one at a time.
|
||||
docs = list(self._collection().find({}))
|
||||
updated_ids = []
|
||||
for doc in docs:
|
||||
d = self._doc_from_mongo(doc)
|
||||
match = cond(d) if callable(cond) else (mongo_filter is not None)
|
||||
if not match:
|
||||
continue
|
||||
if is_callable:
|
||||
fields(d)
|
||||
new_doc = self._doc_to_mongo(d)
|
||||
new_doc.pop('_id', None)
|
||||
self._collection().update_one(
|
||||
{'_id': doc['_id']}, {'$set': new_doc}
|
||||
)
|
||||
else:
|
||||
update_doc = self._doc_to_mongo(fields) or {}
|
||||
update_doc.pop('_id', None)
|
||||
update_doc.pop('id', None)
|
||||
if update_doc:
|
||||
self._collection().update_one(
|
||||
{'_id': doc['_id']}, {'$set': update_doc}
|
||||
)
|
||||
updated_ids.append(str(doc['_id']))
|
||||
return updated_ids
|
||||
|
||||
def remove(self, cond):
|
||||
mongo_filter = self._mongo_filter(cond)
|
||||
if mongo_filter is not None:
|
||||
target_ids = [
|
||||
str(doc['_id'])
|
||||
for doc in self._collection().find(mongo_filter, {'_id': 1})
|
||||
]
|
||||
if target_ids:
|
||||
self._collection().delete_many(mongo_filter)
|
||||
return target_ids
|
||||
|
||||
# Fallback: evaluate the query in-memory and delete one at a time.
|
||||
docs = list(self._collection().find({}))
|
||||
removed_ids = []
|
||||
for doc in docs:
|
||||
d = self._doc_from_mongo(doc)
|
||||
if callable(cond) and cond(d):
|
||||
self._collection().delete_one({'_id': doc['_id']})
|
||||
removed_ids.append(str(doc['_id']))
|
||||
return removed_ids
|
||||
|
||||
def truncate(self):
|
||||
self._collection().delete_many({})
|
||||
|
||||
def close(self):
|
||||
# MongoDB clients are shared and long-lived; nothing to close here.
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Collection factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_table(json_path: str, collection_name: str):
|
||||
if USE_MONGODB:
|
||||
return MongoLockedTable(collection_name)
|
||||
db = TinyDB(json_path, indent=2)
|
||||
return LockedTable(db)
|
||||
|
||||
|
||||
# Setup DB files next to this module
|
||||
|
||||
child_path = os.path.join(base_dir, 'children.json')
|
||||
@@ -86,46 +470,116 @@ refresh_tokens_path = os.path.join(base_dir, 'refresh_tokens.json')
|
||||
push_subscriptions_path = os.path.join(base_dir, 'push_subscriptions.json')
|
||||
digest_action_tokens_path = os.path.join(base_dir, 'digest_action_tokens.json')
|
||||
|
||||
# Use separate TinyDB instances/files for each collection
|
||||
_child_db = TinyDB(child_path, indent=2)
|
||||
_task_db = TinyDB(task_path, indent=2)
|
||||
_routine_db = TinyDB(routine_path, indent=2)
|
||||
_routine_items_db = TinyDB(routine_items_path, indent=2)
|
||||
_routine_schedules_db = TinyDB(routine_schedules_path, indent=2)
|
||||
_routine_extensions_db = TinyDB(routine_extensions_path, indent=2)
|
||||
_reward_db = TinyDB(reward_path, indent=2)
|
||||
_image_db = TinyDB(image_path, indent=2)
|
||||
_pending_rewards_db = TinyDB(pending_reward_path, indent=2)
|
||||
_pending_confirmations_db = TinyDB(pending_confirmations_path, indent=2)
|
||||
_users_db = TinyDB(users_path, indent=2)
|
||||
_tracking_events_db = TinyDB(tracking_events_path, indent=2)
|
||||
_child_overrides_db = TinyDB(child_overrides_path, indent=2)
|
||||
_chore_schedules_db = TinyDB(chore_schedules_path, indent=2)
|
||||
_task_extensions_db = TinyDB(task_extensions_path, indent=2)
|
||||
_refresh_tokens_db = TinyDB(refresh_tokens_path, indent=2)
|
||||
_push_subscriptions_db = TinyDB(push_subscriptions_path, indent=2)
|
||||
_digest_action_tokens_db = TinyDB(digest_action_tokens_path, indent=2)
|
||||
# Expose table objects backed by TinyDB or MongoDB based on USE_MONGODB
|
||||
child_db = _make_table(child_path, 'children')
|
||||
task_db = _make_table(task_path, 'tasks')
|
||||
routine_db = _make_table(routine_path, 'routines')
|
||||
routine_items_db = _make_table(routine_items_path, 'routine_items')
|
||||
routine_schedules_db = _make_table(routine_schedules_path, 'routine_schedules')
|
||||
routine_extensions_db = _make_table(routine_extensions_path, 'routine_extensions')
|
||||
reward_db = _make_table(reward_path, 'rewards')
|
||||
image_db = _make_table(image_path, 'images')
|
||||
pending_reward_db = _make_table(pending_reward_path, 'pending_rewards')
|
||||
pending_confirmations_db = _make_table(pending_confirmations_path, 'pending_confirmations')
|
||||
users_db = _make_table(users_path, 'users')
|
||||
tracking_events_db = _make_table(tracking_events_path, 'tracking_events')
|
||||
child_overrides_db = _make_table(child_overrides_path, 'child_overrides')
|
||||
chore_schedules_db = _make_table(chore_schedules_path, 'chore_schedules')
|
||||
task_extensions_db = _make_table(task_extensions_path, 'task_extensions')
|
||||
refresh_tokens_db = _make_table(refresh_tokens_path, 'refresh_tokens')
|
||||
push_subscriptions_db = _make_table(push_subscriptions_path, 'push_subscriptions')
|
||||
digest_action_tokens_db = _make_table(digest_action_tokens_path, 'digest_action_tokens')
|
||||
|
||||
# Expose table objects wrapped with locking
|
||||
child_db = LockedTable(_child_db)
|
||||
task_db = LockedTable(_task_db)
|
||||
routine_db = LockedTable(_routine_db)
|
||||
routine_items_db = LockedTable(_routine_items_db)
|
||||
routine_schedules_db = LockedTable(_routine_schedules_db)
|
||||
routine_extensions_db = LockedTable(_routine_extensions_db)
|
||||
reward_db = LockedTable(_reward_db)
|
||||
image_db = LockedTable(_image_db)
|
||||
pending_reward_db = LockedTable(_pending_rewards_db)
|
||||
pending_confirmations_db = LockedTable(_pending_confirmations_db)
|
||||
users_db = LockedTable(_users_db)
|
||||
tracking_events_db = LockedTable(_tracking_events_db)
|
||||
child_overrides_db = LockedTable(_child_overrides_db)
|
||||
chore_schedules_db = LockedTable(_chore_schedules_db)
|
||||
task_extensions_db = LockedTable(_task_extensions_db)
|
||||
refresh_tokens_db = LockedTable(_refresh_tokens_db)
|
||||
push_subscriptions_db = LockedTable(_push_subscriptions_db)
|
||||
digest_action_tokens_db = LockedTable(_digest_action_tokens_db)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Index management
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
COLLECTION_INDEXES = {
|
||||
# NOTE: The model ``id`` field is stored as MongoDB's primary key ``_id``,
|
||||
# so no separate unique index on ``id`` is needed. Only secondary indexes
|
||||
# for frequently queried fields are defined here.
|
||||
'children': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
],
|
||||
'tasks': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
],
|
||||
'routines': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
],
|
||||
'routine_items': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
],
|
||||
'routine_schedules': [],
|
||||
'routine_extensions': [],
|
||||
'rewards': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
],
|
||||
'images': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
],
|
||||
'pending_rewards': [
|
||||
{'keys': [('child_id', ASCENDING)]},
|
||||
],
|
||||
'pending_confirmations': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
{'keys': [('child_id', ASCENDING)]},
|
||||
{'keys': [('entity_id', ASCENDING), ('entity_type', ASCENDING)]},
|
||||
],
|
||||
'users': [],
|
||||
'tracking_events': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
{'keys': [('child_id', ASCENDING)]},
|
||||
{'keys': [('entity_id', ASCENDING), ('entity_type', ASCENDING)]},
|
||||
],
|
||||
'child_overrides': [
|
||||
{'keys': [('child_id', ASCENDING)]},
|
||||
{'keys': [('entity_id', ASCENDING), ('entity_type', ASCENDING)]},
|
||||
],
|
||||
'chore_schedules': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
{'keys': [('child_id', ASCENDING)]},
|
||||
],
|
||||
'task_extensions': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
{'keys': [('child_id', ASCENDING)]},
|
||||
],
|
||||
'refresh_tokens': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
{'keys': [('token', ASCENDING)], 'unique': True, 'sparse': True},
|
||||
],
|
||||
'push_subscriptions': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
],
|
||||
'digest_action_tokens': [
|
||||
{'keys': [('user_id', ASCENDING)]},
|
||||
{'keys': [('token', ASCENDING)], 'unique': True, 'sparse': True},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def ensure_mongodb_indexes(client=None, db_name=None):
|
||||
"""Create required indexes on all MongoDB collections.
|
||||
|
||||
Safe to call repeatedly: MongoDB treats index creation as idempotent.
|
||||
"""
|
||||
if not USE_MONGODB:
|
||||
return
|
||||
|
||||
client = client or get_mongo_client()
|
||||
db_name = db_name or _mongo_db_name
|
||||
db = client[db_name]
|
||||
|
||||
for collection_name, indexes in COLLECTION_INDEXES.items():
|
||||
coll = db[collection_name]
|
||||
for spec in indexes:
|
||||
keys = spec['keys']
|
||||
kwargs = {k: v for k, v in spec.items() if k != 'keys'}
|
||||
coll.create_index(keys, **kwargs)
|
||||
|
||||
|
||||
# Clear test collections at import time so tests start with a clean slate.
|
||||
if os.environ.get('DB_ENV', 'prod') == 'test':
|
||||
child_db.truncate()
|
||||
task_db.truncate()
|
||||
@@ -145,4 +599,3 @@ if os.environ.get('DB_ENV', 'prod') == 'test':
|
||||
refresh_tokens_db.truncate()
|
||||
push_subscriptions_db.truncate()
|
||||
digest_action_tokens_db.truncate()
|
||||
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
# python
|
||||
"""Lazy MongoDB client factory and database-name helpers.
|
||||
|
||||
The client is intentionally **not** created at module import. Use
|
||||
``get_mongo_client()`` to obtain a cached singleton. For Gunicorn multi-worker
|
||||
deployments call ``init_mongo_client()`` from a ``post_fork`` hook so each
|
||||
worker process owns its own connection pool rather than inheriting the parent
|
||||
process's client.
|
||||
"""
|
||||
import os
|
||||
import threading
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from pymongo import MongoClient
|
||||
from pymongo.uri_parser import parse_uri
|
||||
|
||||
|
||||
_mongo_client_lock = threading.Lock()
|
||||
_mongo_client = None
|
||||
|
||||
|
||||
def _create_mongo_client():
|
||||
"""Build a fail-fast MongoClient from environment variables."""
|
||||
uri = os.environ.get('MONGO_URI')
|
||||
if not uri:
|
||||
raise RuntimeError(
|
||||
'MONGO_URI environment variable is required when USE_MONGODB=true.'
|
||||
)
|
||||
|
||||
# mongomock is used for unit/integration tests without a real server.
|
||||
if uri.lower().startswith('mongomock') or uri.lower() == 'mongomock':
|
||||
try:
|
||||
import mongomock
|
||||
except ImportError as exc: # pragma: no cover - test dependency
|
||||
raise RuntimeError(
|
||||
'mongomock is required for test MongoDB mode. '
|
||||
'Install it with: pip install mongomock'
|
||||
) from exc
|
||||
return mongomock.MongoClient()
|
||||
|
||||
return MongoClient(
|
||||
uri,
|
||||
serverSelectionTimeoutMS=5000,
|
||||
connectTimeoutMS=5000,
|
||||
maxPoolSize=20,
|
||||
)
|
||||
|
||||
|
||||
def init_mongo_client():
|
||||
"""Create a fresh MongoClient and store it as the process singleton.
|
||||
|
||||
Call this from a Gunicorn ``post_fork`` hook so each worker process gets
|
||||
its own client after forking. It can also be called in tests to reset the
|
||||
shared client to a known state.
|
||||
"""
|
||||
global _mongo_client
|
||||
with _mongo_client_lock:
|
||||
_mongo_client = _create_mongo_client()
|
||||
return _mongo_client
|
||||
|
||||
|
||||
def get_mongo_client():
|
||||
"""Return the cached process-level MongoClient, creating it lazily once."""
|
||||
global _mongo_client
|
||||
if _mongo_client is None:
|
||||
with _mongo_client_lock:
|
||||
if _mongo_client is None:
|
||||
_mongo_client = _create_mongo_client()
|
||||
return _mongo_client
|
||||
|
||||
|
||||
def _db_name_from_uri(uri: str) -> str | None:
|
||||
"""Extract the database name from a MongoDB connection URI, if present."""
|
||||
if not uri or uri.lower().startswith('mongomock'):
|
||||
return None
|
||||
try:
|
||||
parsed = parse_uri(uri)
|
||||
return parsed.get('database') or None
|
||||
except Exception:
|
||||
# Fallback to a simple path-based parse for non-standard URIs.
|
||||
try:
|
||||
path = urlparse(uri).path
|
||||
return path.lstrip('/') or None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def get_mongo_db_name() -> str:
|
||||
"""Resolve the MongoDB database name from env vars or the connection URI.
|
||||
|
||||
Precedence:
|
||||
1. ``MONGO_DB_NAME`` environment variable.
|
||||
2. Database name parsed from ``MONGO_URI``.
|
||||
3. Default based on ``DATA_ENV`` / ``DB_ENV``:
|
||||
* prod -> ``chore_db``
|
||||
* e2e -> ``chore_db_e2e``
|
||||
* test -> ``chore_db_test``
|
||||
"""
|
||||
env_name = os.environ.get('MONGO_DB_NAME')
|
||||
if env_name:
|
||||
return env_name
|
||||
|
||||
uri = os.environ.get('MONGO_URI', '')
|
||||
db_name = _db_name_from_uri(uri)
|
||||
if db_name:
|
||||
return db_name
|
||||
|
||||
env = (os.environ.get('DATA_ENV') or os.environ.get('DB_ENV', 'prod')).lower()
|
||||
if env == 'prod':
|
||||
return 'chore_db'
|
||||
if env == 'e2e':
|
||||
return 'chore_db_e2e'
|
||||
return 'chore_db_test'
|
||||
+22
-8
@@ -1,4 +1,5 @@
|
||||
"""Helper functions for tracking events database operations."""
|
||||
import itertools
|
||||
import logging
|
||||
from typing import Optional, List
|
||||
from tinydb import Query
|
||||
@@ -8,6 +9,10 @@ from models.tracking_event import TrackingEvent, EntityType, ActionType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Monotonic sequence used as a deterministic tiebreaker when tracking events
|
||||
# share the same ``occurred_at``/``created_at`` timestamps (common in tests).
|
||||
_tracking_event_seq = itertools.count()
|
||||
|
||||
|
||||
def insert_tracking_event(event: TrackingEvent) -> str:
|
||||
"""
|
||||
@@ -20,7 +25,9 @@ def insert_tracking_event(event: TrackingEvent) -> str:
|
||||
The event ID
|
||||
"""
|
||||
try:
|
||||
tracking_events_db.insert(event.to_dict())
|
||||
event_dict = event.to_dict()
|
||||
event_dict['_seq'] = next(_tracking_event_seq)
|
||||
tracking_events_db.insert(event_dict)
|
||||
logger.info(f"Tracking event created: {event.action} {event.entity_type} {event.entity_id} for child {event.child_id}")
|
||||
return event.id
|
||||
except Exception as e:
|
||||
@@ -61,12 +68,16 @@ def get_tracking_events_by_child(
|
||||
all_results = tracking_events_db.search(query_condition)
|
||||
total = len(all_results)
|
||||
|
||||
# Sort by occurred_at desc, then created_at desc
|
||||
all_results.sort(key=lambda x: (x.get('occurred_at', ''), x.get('created_at', 0)), reverse=True)
|
||||
|
||||
# Sort by occurred_at desc, then created_at desc, then _seq desc for
|
||||
# deterministic ordering when timestamps collide (common in fast tests).
|
||||
all_results.sort(
|
||||
key=lambda x: (x.get('occurred_at', ''), x.get('created_at', 0), x.get('_seq', 0)),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
paginated = all_results[offset:offset + limit]
|
||||
events = [TrackingEvent.from_dict(r) for r in paginated]
|
||||
|
||||
|
||||
return events, total
|
||||
|
||||
|
||||
@@ -99,11 +110,14 @@ def get_tracking_events_by_user(
|
||||
all_results = tracking_events_db.search(query_condition)
|
||||
total = len(all_results)
|
||||
|
||||
all_results.sort(key=lambda x: (x.get('occurred_at', ''), x.get('created_at', 0)), reverse=True)
|
||||
|
||||
all_results.sort(
|
||||
key=lambda x: (x.get('occurred_at', ''), x.get('created_at', 0), x.get('_seq', 0)),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
paginated = all_results[offset:offset + limit]
|
||||
events = [TrackingEvent.from_dict(r) for r in paginated]
|
||||
|
||||
|
||||
return events, total
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user