pymongo_client.py

# Copyright (c) 2021 Tiago Silva <https://github.com/tarsil>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
from bson.objectid import ObjectId
from pymongo import InsertOne

mongo = '<your mongo db client object>'


class MongoClient:
    """
    SlashWap version of mongo db client.
    """
    def __init__(self, doc_table: str, **kwargs):
        """
        Args:
            doc_table: str - The name of the table to query.
        """
        self.table = getattr(mongo, doc_table, None)

        if self.table is None:
            raise ValueError(f"Mongo client not found or {doc_table} not accessible.")

    def get_queryset(self, query, page=0, total_per_page=50):
        """
        Returns a cursor to a list of documents.
        Based on the page number and the number of docs per page, the result may be skipped and limited.

        Returns 2 elements in a tuple: (queryset, count)
        """
        cursor = self.table.find(query)

        count = 0 if page == 0 else self.table.count_documents(query)
        queryset = cursor.limit(total_per_page)

        return (list(queryset), count)

    def get_document(self, _id: Optional[str] = None, **query: dict):
        """
        Given a document id, obtains the data.
        """
        if _id and query:
            raise ValueError("You can only search by id or personalised query but not both.")

        if _id:
            document = self.table.find_one({'_id': ObjectId(_id)})
        else:
            document = self.table.find_one(query)
        return document

    def update_document(self, _id, data, upsert=True):
        """
        Updates the state of a document by id.
        """
        result = self.table.update_one(
            { "_id": ObjectId(_id) },
            { "$set": data }
        , upsert=upsert)
        return result

    def create_document(self, **data):
        """
        Creates a document in the database
        """
        result = self.table.insert_one(data)
        return result

    def bulk_write(self, data):
        """
        Bulk writes in to the database
        """
        result = self.table.bulk_write([
            InsertOne(document) for document in data
        ])
        return result