planetmint/planetmint/backend/localmongodb/connection.py

# Copyright © 2020 Interplanetary Database Association e.V.,
# Planetmint and IPDB software contributors.
# SPDX-License-Identifier: (Apache-2.0 AND CC-BY-4.0)
# Code is Apache-2.0 and docs are CC-BY-4.0

import logging
from ssl import CERT_REQUIRED
import pymongo

from planetmint.config import Config
from planetmint.backend.exceptions import DuplicateKeyError, OperationError, ConnectionError
from planetmint.transactions.common.exceptions import ConfigurationError
from planetmint.utils import Lazy
from planetmint.backend.connection import DBConnection, _kwargs_parser

logger = logging.getLogger(__name__)


class LocalMongoDBConnection(DBConnection):
    def __init__(self, host: str = None, port: int = None, login: str = None, password: str = None, **kwargs):
        """Create a new Connection instance.

        Args:
            replicaset (str, optional): the name of the replica set to
                                        connect to.
            **kwargs: arbitrary keyword arguments provided by the
                configuration's ``database`` settings
        """

        super().__init__(host=host, port=port, login=login, password=password, **kwargs)

        dbconf = Config().get()["database"]
        self.dbname = _kwargs_parser(key="name", kwargs=kwargs) or dbconf["name"]
        self.replicaset = _kwargs_parser(key="replicaset", kwargs=kwargs) or dbconf["replicaset"]
        self.ssl = _kwargs_parser(key="ssl", kwargs=kwargs) or dbconf["ssl"]

        self.ca_cert = _kwargs_parser(key="ca_cert", kwargs=kwargs) or dbconf["ca_cert"]
        self.certfile = _kwargs_parser(key="certfile", kwargs=kwargs) or dbconf["certfile"]
        self.keyfile = _kwargs_parser(key="keyfile", kwargs=kwargs) or dbconf["keyfile"]
        self.keyfile_passphrase = (
            _kwargs_parser(key="keyfile_passphrase", kwargs=kwargs) or dbconf["keyfile_passphrase"]
        )
        self.crlfile = _kwargs_parser(key="crlfile", kwargs=kwargs) or dbconf["crlfile"]
        self.max_tries = _kwargs_parser(key="max_tries", kwargs=kwargs)
        self.connection_timeout = _kwargs_parser(key="connection_timeout", kwargs=kwargs)
        self.__conn = None
        self.connect()

        if not self.ssl:
            self.ssl = False
        if not self.keyfile_passphrase:
            self.keyfile_passphrase = None

    @property
    def db(self):
        return self.connect()[self.dbname]

    def query(self):
        return Lazy()

    def collection(self, name):
        """Return a lazy object that can be used to compose a query.

        Args:
            name (str): the name of the collection to query.
        """
        return self.query()[self.dbname][name]

    def run(self, query):
        try:
            try:
                return query.run(self.connect())
            except pymongo.errors.AutoReconnect:
                logger.warning("Lost connection to the database, " "retrying query.")
                return query.run(self.connect())
        except pymongo.errors.AutoReconnect as exc:
            raise ConnectionError from exc
        except pymongo.errors.DuplicateKeyError as exc:
            raise DuplicateKeyError from exc
        except pymongo.errors.OperationFailure as exc:
            print(f"DETAILS: {exc.details}")
            raise OperationError from exc

    def connect(self):
        """Try to connect to the database.

        Raises:
            :exc:`~ConnectionError`: If the connection to the database
                fails.
            :exc:`~AuthenticationError`: If there is a OperationFailure due to
                Authentication failure after connecting to the database.
            :exc:`~ConfigurationError`: If there is a ConfigurationError while
                connecting to the database.
        """
        if self.__conn:
            return self._conn
        try:
            # FYI: the connection process might raise a
            # `ServerSelectionTimeoutError`, that is a subclass of
            # `ConnectionFailure`.
            # The presence of ca_cert, certfile, keyfile, crlfile implies the
            # use of certificates for TLS connectivity.
            if self.ca_cert is None or self.certfile is None or self.keyfile is None or self.crlfile is None:
                client = pymongo.MongoClient(
                    self.host,
                    self.port,
                    replicaset=self.replicaset,
                    serverselectiontimeoutms=self.connection_timeout,
                    ssl=self.ssl,
                    **MONGO_OPTS,
                )
                if self.login is not None and self.password is not None:
                    client[self.dbname].authenticate(self.login, self.password)
            else:
                logger.info("Connecting to MongoDB over TLS/SSL...")
                client = pymongo.MongoClient(
                    self.host,
                    self.port,
                    replicaset=self.replicaset,
                    serverselectiontimeoutms=self.connection_timeout,
                    ssl=self.ssl,
                    ssl_ca_certs=self.ca_cert,
                    ssl_certfile=self.certfile,
                    ssl_keyfile=self.keyfile,
                    ssl_pem_passphrase=self.keyfile_passphrase,
                    ssl_crlfile=self.crlfile,
                    ssl_cert_reqs=CERT_REQUIRED,
                    **MONGO_OPTS,
                )
                if self.login is not None:
                    client[self.dbname].authenticate(self.login, mechanism="MONGODB-X509")
            self.__conn = client
            return client

        except (pymongo.errors.ConnectionFailure, pymongo.errors.OperationFailure) as exc:
            logger.info("Exception in connect(): {}".format(exc))
            raise ConnectionError(str(exc)) from exc
        except pymongo.errors.ConfigurationError as exc:
            raise ConfigurationError from exc

    def close(self):
        try:
            self.__conn.close()
            self.__conn = None
        except Exception as exc:
            logger.info("Exception in planetmint.backend.localmongodb.close(): {}".format(exc))
            raise ConnectionError(str(exc)) from exc


MONGO_OPTS = {
    "socketTimeoutMS": 20000,
}