|
|
@@ -12,8 +12,7 @@ object (preferred), one client per process:
|
|
|
|
async with MongoDB(conn_string, database) as db:
|
|
|
|
async with MongoDB(conn_string, database) as db:
|
|
|
|
await db.get_documents("users", {})
|
|
|
|
await db.get_documents("users", {})
|
|
|
|
|
|
|
|
|
|
|
|
`Mongo` remains a back-compat alias of `MongoDB` — existing `Mongo(...)` call sites keep
|
|
|
|
`Mongo` is a back-compat alias of `MongoDB`; old `Mongo(...)` call sites still work.
|
|
|
|
working unchanged.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
module proxy (back-compat), arm once then call bare:
|
|
|
|
module proxy (back-compat), arm once then call bare:
|
|
|
|
import mongo # not `from mongo import ...`
|
|
|
|
import mongo # not `from mongo import ...`
|
|
|
@@ -21,12 +20,11 @@ module proxy (back-compat), arm once then call bare:
|
|
|
|
await mongo.get_documents("users", {"active": True})
|
|
|
|
await mongo.get_documents("users", {"active": True})
|
|
|
|
|
|
|
|
|
|
|
|
errors:
|
|
|
|
errors:
|
|
|
|
wrapped methods log and swallow, returning a safe default
|
|
|
|
wrapped methods log and swallow, returning a safe default (False / [] / {} / 0 /
|
|
|
|
(False / [] / {} / 0 / None). .collection(name) / db[name] return
|
|
|
|
None). .collection(name) / db[name] return the raw motor collection: full driver
|
|
|
|
the raw motor collection: full driver surface, raises, nothing swallowed.
|
|
|
|
surface, raises, nothing swallowed. mongo SWALLOWS by design — the one deliberate
|
|
|
|
NOTE: mongo SWALLOWS by design — this is the one deliberate difference from the
|
|
|
|
difference from the fail-loud redis/psql/mysql trio, kept so existing consumers'
|
|
|
|
newer datastore trio (redis/psql/mysql), which is fail-loud. mongo's contract is
|
|
|
|
branch-on-result control flow doesn't break.
|
|
|
|
kept as-is so existing consumers' branch-on-result control flow doesn't break.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
naming consistency with the trio (all additive — old names still work):
|
|
|
|
naming consistency with the trio (all additive — old names still work):
|
|
|
|
- class is `MongoDB` (was `Mongo`, kept as alias)
|
|
|
|
- class is `MongoDB` (was `Mongo`, kept as alias)
|
|
|
@@ -44,6 +42,7 @@ import logging
|
|
|
|
from typing import Any, List, Optional
|
|
|
|
from typing import Any, List, Optional
|
|
|
|
|
|
|
|
|
|
|
|
from pymongo import ReturnDocument
|
|
|
|
from pymongo import ReturnDocument
|
|
|
|
|
|
|
|
from pymongo.errors import BulkWriteError
|
|
|
|
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection, AsyncIOMotorDatabase
|
|
|
|
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection, AsyncIOMotorDatabase
|
|
|
|
|
|
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
@@ -53,21 +52,39 @@ class MongoDB:
|
|
|
|
"""async mongodb wrapper; one client per process, attach to bot as bot.db"""
|
|
|
|
"""async mongodb wrapper; one client per process, attach to bot as bot.db"""
|
|
|
|
|
|
|
|
|
|
|
|
def __init__(self, connection_string: str, database: str):
|
|
|
|
def __init__(self, connection_string: str, database: str):
|
|
|
|
|
|
|
|
# motor builds the client eagerly (no I/O yet); a bad URI raises InvalidURI
|
|
|
|
|
|
|
|
# here with nothing built. an invalid db name raises below, after the client
|
|
|
|
|
|
|
|
# exists — close it before propagating so it doesn't outlive the raise.
|
|
|
|
self._client = AsyncIOMotorClient(connection_string)
|
|
|
|
self._client = AsyncIOMotorClient(connection_string)
|
|
|
|
|
|
|
|
try:
|
|
|
|
self._db = self._client[database]
|
|
|
|
self._db = self._client[database]
|
|
|
|
|
|
|
|
except BaseException:
|
|
|
|
|
|
|
|
self._client.close()
|
|
|
|
|
|
|
|
raise
|
|
|
|
|
|
|
|
|
|
|
|
async def connect(self) -> "MongoDB":
|
|
|
|
async def connect(self) -> "MongoDB":
|
|
|
|
"""validate the connection with a ping and return self
|
|
|
|
"""validate the connection with a ping and return self
|
|
|
|
|
|
|
|
|
|
|
|
motor connects lazily, so this is optional — call it to fail early on a bad
|
|
|
|
optional — motor connects lazily, so this just fails early on a bad
|
|
|
|
URI/credentials rather than on the first real op (parallel to the trio's
|
|
|
|
URI/credentials (parallel to the trio's connect()). raises, unlike the
|
|
|
|
connect()). raises on a bad connection, unlike the swallowing wrapped methods.
|
|
|
|
swallowing wrapped methods.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
an instance is single-use: pymongo 4.x `close()` is irreversible and
|
|
|
|
|
|
|
|
connect() doesn't rebuild the client. construct a fresh MongoDB instead of
|
|
|
|
|
|
|
|
reusing one after close()/`async with` exit.
|
|
|
|
"""
|
|
|
|
"""
|
|
|
|
await self._client.admin.command("ping")
|
|
|
|
await self._client.admin.command("ping")
|
|
|
|
return self
|
|
|
|
return self
|
|
|
|
|
|
|
|
|
|
|
|
async def __aenter__(self) -> "MongoDB":
|
|
|
|
async def __aenter__(self) -> "MongoDB":
|
|
|
|
|
|
|
|
try:
|
|
|
|
return await self.connect()
|
|
|
|
return await self.connect()
|
|
|
|
|
|
|
|
except BaseException:
|
|
|
|
|
|
|
|
# __aexit__ is not called when __aenter__ raises, so a failed ping here
|
|
|
|
|
|
|
|
# would otherwise leak the motor client built in __init__. close it
|
|
|
|
|
|
|
|
# before propagating.
|
|
|
|
|
|
|
|
self.close()
|
|
|
|
|
|
|
|
raise
|
|
|
|
|
|
|
|
|
|
|
|
async def __aexit__(self, exc_type, exc, tb) -> None:
|
|
|
|
async def __aexit__(self, exc_type, exc, tb) -> None:
|
|
|
|
self.close()
|
|
|
|
self.close()
|
|
|
@@ -101,13 +118,17 @@ class MongoDB:
|
|
|
|
# -------------------------------------------------------------------------
|
|
|
|
# -------------------------------------------------------------------------
|
|
|
|
# collection / index management
|
|
|
|
# collection / index management
|
|
|
|
|
|
|
|
|
|
|
|
async def create_collection(self, collection: str, indexes=None) -> bool:
|
|
|
|
async def create_collection(self, collection: str, index=None) -> bool:
|
|
|
|
"""create a collection, optionally seeding an index"""
|
|
|
|
"""create a collection, optionally seeding ONE index
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
`index` is a single index key-spec (e.g. `[("a", 1), ("b", 1)]` builds one
|
|
|
|
|
|
|
|
compound index, NOT two); call `create_index` per index for multiples.
|
|
|
|
|
|
|
|
"""
|
|
|
|
try:
|
|
|
|
try:
|
|
|
|
await self._db.create_collection(collection)
|
|
|
|
await self._db.create_collection(collection)
|
|
|
|
if indexes:
|
|
|
|
if index:
|
|
|
|
await self._db[collection].create_index(indexes)
|
|
|
|
await self._db[collection].create_index(index)
|
|
|
|
log.info(f"created indexes for {collection}")
|
|
|
|
log.info(f"created index for {collection}")
|
|
|
|
return True
|
|
|
|
return True
|
|
|
|
except Exception:
|
|
|
|
except Exception:
|
|
|
|
log.exception(f"db.create_collection() for {collection}")
|
|
|
|
log.exception(f"db.create_collection() for {collection}")
|
|
|
@@ -163,10 +184,19 @@ class MongoDB:
|
|
|
|
async def create_documents(
|
|
|
|
async def create_documents(
|
|
|
|
self, collection: str, documents: List[dict], ordered: bool = True
|
|
|
|
self, collection: str, documents: List[dict], ordered: bool = True
|
|
|
|
) -> int:
|
|
|
|
) -> int:
|
|
|
|
"""insert many documents, returning the inserted count"""
|
|
|
|
"""insert many documents, returning the inserted count
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
on a partial failure (e.g. ordered=False with one duplicate key) the driver
|
|
|
|
|
|
|
|
raises BulkWriteError AFTER inserting the rest; count comes from
|
|
|
|
|
|
|
|
`e.details['nInserted']` so a caller retrying on a wrongly-reported 0 doesn't
|
|
|
|
|
|
|
|
re-insert documents that already landed.
|
|
|
|
|
|
|
|
"""
|
|
|
|
try:
|
|
|
|
try:
|
|
|
|
result = await self._db[collection].insert_many(documents, ordered=ordered)
|
|
|
|
result = await self._db[collection].insert_many(documents, ordered=ordered)
|
|
|
|
return len(result.inserted_ids)
|
|
|
|
return len(result.inserted_ids)
|
|
|
|
|
|
|
|
except BulkWriteError as e:
|
|
|
|
|
|
|
|
log.exception(f"db.create_documents() on {collection}")
|
|
|
|
|
|
|
|
return e.details.get("nInserted", 0)
|
|
|
|
except Exception:
|
|
|
|
except Exception:
|
|
|
|
log.exception(f"db.create_documents() on {collection}")
|
|
|
|
log.exception(f"db.create_documents() on {collection}")
|
|
|
|
return 0
|
|
|
|
return 0
|
|
|
@@ -346,9 +376,9 @@ class MongoDB:
|
|
|
|
async def update_document_field(self, collection: str, target: dict, updates: dict) -> bool:
|
|
|
|
async def update_document_field(self, collection: str, target: dict, updates: dict) -> bool:
|
|
|
|
"""$set one or more fields on a single document
|
|
|
|
"""$set one or more fields on a single document
|
|
|
|
|
|
|
|
|
|
|
|
returns True when a document matched, even if the write changed nothing (a
|
|
|
|
returns True when a document matched, even if `$set` to an identical value
|
|
|
|
`$set` to the identical value leaves `modified_count=0`); use `matched_count`
|
|
|
|
leaves `modified_count=0` — uses `matched_count` so an idempotent no-op isn't
|
|
|
|
so an idempotent no-op on an existing doc isn't misread as a failure.
|
|
|
|
misread as a failure.
|
|
|
|
"""
|
|
|
|
"""
|
|
|
|
try:
|
|
|
|
try:
|
|
|
|
response = await self._db[collection].update_one(target, {"$set": updates})
|
|
|
|
response = await self._db[collection].update_one(target, {"$set": updates})
|
|
|
@@ -362,9 +392,8 @@ class MongoDB:
|
|
|
|
) -> bool:
|
|
|
|
) -> bool:
|
|
|
|
"""apply raw update operators ($set/$inc/$unset/...) to a single document
|
|
|
|
"""apply raw update operators ($set/$inc/$unset/...) to a single document
|
|
|
|
|
|
|
|
|
|
|
|
returns True when a document matched, even if the operators changed nothing
|
|
|
|
returns True when a document matched, even if the operators changed nothing —
|
|
|
|
(e.g. `$set` to an identical value) — use `matched_count` so an idempotent
|
|
|
|
uses `matched_count` so an idempotent write isn't misread as a failure.
|
|
|
|
write on an existing doc isn't misread as a failure.
|
|
|
|
|
|
|
|
"""
|
|
|
|
"""
|
|
|
|
try:
|
|
|
|
try:
|
|
|
|
response = await self._db[collection].update_one(target, update_query)
|
|
|
|
response = await self._db[collection].update_one(target, update_query)
|
|
|
@@ -374,12 +403,17 @@ class MongoDB:
|
|
|
|
return False
|
|
|
|
return False
|
|
|
|
|
|
|
|
|
|
|
|
async def upsert_document_field(self, collection: str, target: dict, updates: dict) -> bool:
|
|
|
|
async def upsert_document_field(self, collection: str, target: dict, updates: dict) -> bool:
|
|
|
|
"""$set fields on a single document, creating it if absent"""
|
|
|
|
"""$set fields on a single document, creating it if absent
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
returns True when a document matched or was upserted, even if `$set` to an
|
|
|
|
|
|
|
|
identical value leaves `modified_count=0` — mirrors update_document_field's
|
|
|
|
|
|
|
|
matched_count-over-modified_count contract.
|
|
|
|
|
|
|
|
"""
|
|
|
|
try:
|
|
|
|
try:
|
|
|
|
response = await self._db[collection].update_one(
|
|
|
|
response = await self._db[collection].update_one(
|
|
|
|
target, {"$set": updates}, upsert=True
|
|
|
|
target, {"$set": updates}, upsert=True
|
|
|
|
)
|
|
|
|
)
|
|
|
|
return bool(response.modified_count or response.upserted_id)
|
|
|
|
return bool(response.matched_count or response.upserted_id)
|
|
|
|
except Exception:
|
|
|
|
except Exception:
|
|
|
|
log.exception(f"db.upsert_document_field() on {collection}")
|
|
|
|
log.exception(f"db.upsert_document_field() on {collection}")
|
|
|
|
return False
|
|
|
|
return False
|
|
|
@@ -400,9 +434,8 @@ class MongoDB:
|
|
|
|
) -> bool:
|
|
|
|
) -> bool:
|
|
|
|
"""$push a value onto an array field
|
|
|
|
"""$push a value onto an array field
|
|
|
|
|
|
|
|
|
|
|
|
returns True when a document matched (consistent with the other single-doc
|
|
|
|
returns True when a document matched; $push always mutates, so matched_count
|
|
|
|
update helpers); $push always mutates, so matched_count and modified_count
|
|
|
|
and modified_count agree here in practice.
|
|
|
|
agree here in practice.
|
|
|
|
|
|
|
|
"""
|
|
|
|
"""
|
|
|
|
try:
|
|
|
|
try:
|
|
|
|
response = await self._db[collection].update_one(target, {"$push": {array: value}})
|
|
|
|
response = await self._db[collection].update_one(target, {"$push": {array: value}})
|
|
|
@@ -417,8 +450,7 @@ class MongoDB:
|
|
|
|
) -> bool:
|
|
|
|
) -> bool:
|
|
|
|
"""$push to an array and $set a field in one update
|
|
|
|
"""$push to an array and $set a field in one update
|
|
|
|
|
|
|
|
|
|
|
|
returns True when a document matched (consistent with the other single-doc
|
|
|
|
returns True when a document matched (matched_count, not modified_count).
|
|
|
|
update helpers).
|
|
|
|
|
|
|
|
"""
|
|
|
|
"""
|
|
|
|
try:
|
|
|
|
try:
|
|
|
|
response = await self._db[collection].update_one(
|
|
|
|
response = await self._db[collection].update_one(
|
|
|
@@ -435,9 +467,8 @@ class MongoDB:
|
|
|
|
) -> bool:
|
|
|
|
) -> bool:
|
|
|
|
"""$pull a value from an array field
|
|
|
|
"""$pull a value from an array field
|
|
|
|
|
|
|
|
|
|
|
|
returns True when a document matched, even if nothing was pulled (the value
|
|
|
|
returns True when a document matched, even if the value was absent and
|
|
|
|
was absent) — use `matched_count` so a no-op pull on an existing doc isn't
|
|
|
|
nothing was pulled — uses `matched_count`, not `modified_count`.
|
|
|
|
misread as a failure.
|
|
|
|
|
|
|
|
"""
|
|
|
|
"""
|
|
|
|
try:
|
|
|
|
try:
|
|
|
|
response = await self._db[collection].update_one(target, {"$pull": {array: value}})
|
|
|
|
response = await self._db[collection].update_one(target, {"$pull": {array: value}})
|
|
|
@@ -517,10 +548,9 @@ class MongoDB:
|
|
|
|
# bulk / checks
|
|
|
|
# bulk / checks
|
|
|
|
|
|
|
|
|
|
|
|
async def bulk_write(self, collection: str, operations: list, ordered: bool = True) -> dict:
|
|
|
|
async def bulk_write(self, collection: str, operations: list, ordered: bool = True) -> dict:
|
|
|
|
"""
|
|
|
|
"""run a batch of pymongo write ops (InsertOne/UpdateOne/DeleteOne/...)
|
|
|
|
run a batch of pymongo write ops (InsertOne/UpdateOne/DeleteOne/...)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
caller builds the ops from pymongo, e.g.:
|
|
|
|
caller builds the ops, e.g.:
|
|
|
|
from pymongo import UpdateOne
|
|
|
|
from pymongo import UpdateOne
|
|
|
|
ops = [UpdateOne({'_id': i}, {'$set': {...}}, upsert=True) for i in ids]
|
|
|
|
ops = [UpdateOne({'_id': i}, {'$set': {...}}, upsert=True) for i in ids]
|
|
|
|
returns a summary dict of counts, or an empty dict on failure
|
|
|
|
returns a summary dict of counts, or an empty dict on failure
|
|
|
@@ -541,8 +571,7 @@ class MongoDB:
|
|
|
|
async def document_check_multi(
|
|
|
|
async def document_check_multi(
|
|
|
|
self, collection: str, target: dict, checks: List[str], value: str
|
|
|
|
self, collection: str, target: dict, checks: List[str], value: str
|
|
|
|
) -> bool:
|
|
|
|
) -> bool:
|
|
|
|
"""
|
|
|
|
"""return whether a doc matching target has value in any of the given arrays
|
|
|
|
return whether a doc matching target has value in any of the given arrays
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
target is a normal filter dict (merged into the query), not a bare _id
|
|
|
|
target is a normal filter dict (merged into the query), not a bare _id
|
|
|
|
"""
|
|
|
|
"""
|
|
|
@@ -563,12 +592,9 @@ Mongo = MongoDB
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
|
|
# -----------------------------------------------------------------------------
|
|
|
|
# backwards-compat module proxy
|
|
|
|
# backwards-compat module proxy: `import mongo; mongo.init(conn, db)` then
|
|
|
|
# - lets legacy call sites keep using `await mongo.get_documents(...)`
|
|
|
|
# bare `await mongo.func(...)`. needs `import mongo` — `from mongo import func`
|
|
|
|
# without an object, after a one-time `mongo.init(conn, db)`
|
|
|
|
# resolves at import time, before init() runs.
|
|
|
|
# - works with `import mongo; mongo.func(...)` (resolved at call time)
|
|
|
|
|
|
|
|
# - does NOT work with `from mongo import func` (resolved at import,
|
|
|
|
|
|
|
|
# before init runs) — switch those sites to `import mongo`
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
_default: Optional[MongoDB] = None
|
|
|
|
_default: Optional[MongoDB] = None
|
|
|
|
|
|
|
|
|
|
|
|