GNXSOFT.COM

gnx-react/venv/lib/python3.12/site-packages/kombu/transport/mongodb.py

@@ -0,0 +1,534 @@
+# copyright: (c) 2010 - 2013 by Flavio Percoco Premoli.
+# license: BSD, see LICENSE for more details.
+
+"""MongoDB transport module for kombu.
+
+Features
+========
+* Type: Virtual
+* Supports Direct: Yes
+* Supports Topic: Yes
+* Supports Fanout: Yes
+* Supports Priority: Yes
+* Supports TTL: Yes
+
+Connection String
+=================
+ *Unreviewed*
+
+Transport Options
+=================
+
+* ``connect_timeout``,
+* ``ssl``,
+* ``ttl``,
+* ``capped_queue_size``,
+* ``default_hostname``,
+* ``default_port``,
+* ``default_database``,
+* ``messages_collection``,
+* ``routing_collection``,
+* ``broadcast_collection``,
+* ``queues_collection``,
+* ``calc_queue_size``,
+"""
+
+from __future__ import annotations
+
+import datetime
+from queue import Empty
+
+import pymongo
+from pymongo import MongoClient, errors, uri_parser
+from pymongo.cursor import CursorType
+
+from kombu.exceptions import VersionMismatch
+from kombu.utils.compat import _detect_environment
+from kombu.utils.encoding import bytes_to_str
+from kombu.utils.json import dumps, loads
+from kombu.utils.objects import cached_property
+from kombu.utils.url import maybe_sanitize_url
+
+from . import virtual
+from .base import to_rabbitmq_queue_arguments
+
+E_SERVER_VERSION = """\
+Kombu requires MongoDB version 1.3+ (server is {0})\
+"""
+
+E_NO_TTL_INDEXES = """\
+Kombu requires MongoDB version 2.2+ (server is {0}) for TTL indexes support\
+"""
+
+
+class BroadcastCursor:
+    """Cursor for broadcast queues."""
+
+    def __init__(self, cursor):
+        self._cursor = cursor
+        self._offset = 0
+        self.purge(rewind=False)
+
+    def get_size(self):
+        return self._cursor.collection.count_documents({}) - self._offset
+
+    def close(self):
+        self._cursor.close()
+
+    def purge(self, rewind=True):
+        if rewind:
+            self._cursor.rewind()
+
+        # Fast-forward the cursor past old events
+        self._offset = self._cursor.collection.count_documents({})
+        self._cursor = self._cursor.skip(self._offset)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        while True:
+            try:
+                msg = next(self._cursor)
+            except pymongo.errors.OperationFailure as exc:
+                # In some cases tailed cursor can become invalid
+                # and have to be reinitalized
+                if 'not valid at server' in str(exc):
+                    self.purge()
+
+                    continue
+
+                raise
+            else:
+                break
+
+        self._offset += 1
+
+        return msg
+    next = __next__
+
+
+class Channel(virtual.Channel):
+    """MongoDB Channel."""
+
+    supports_fanout = True
+
+    # Mutable container. Shared by all class instances
+    _fanout_queues = {}
+
+    # Options
+    ssl = False
+    ttl = False
+    connect_timeout = None
+    capped_queue_size = 100000
+    calc_queue_size = True
+
+    default_hostname = '127.0.0.1'
+    default_port = 27017
+    default_database = 'kombu_default'
+
+    messages_collection = 'messages'
+    routing_collection = 'messages.routing'
+    broadcast_collection = 'messages.broadcast'
+    queues_collection = 'messages.queues'
+
+    from_transport_options = (virtual.Channel.from_transport_options + (
+        'connect_timeout', 'ssl', 'ttl', 'capped_queue_size',
+        'default_hostname', 'default_port', 'default_database',
+        'messages_collection', 'routing_collection',
+        'broadcast_collection', 'queues_collection',
+        'calc_queue_size',
+    ))
+
+    def __init__(self, *vargs, **kwargs):
+        super().__init__(*vargs, **kwargs)
+
+        self._broadcast_cursors = {}
+
+        # Evaluate connection
+        self.client
+
+    # AbstractChannel/Channel interface implementation
+
+    def _new_queue(self, queue, **kwargs):
+        if self.ttl:
+            self.queues.update_one(
+                {'_id': queue},
+                {
+                    '$set': {
+                        '_id': queue,
+                        'options': kwargs,
+                        'expire_at': self._get_queue_expire(
+                            kwargs, 'x-expires'
+                        ),
+                    },
+                },
+                upsert=True)
+
+    def _get(self, queue):
+        if queue in self._fanout_queues:
+            try:
+                msg = next(self._get_broadcast_cursor(queue))
+            except StopIteration:
+                msg = None
+        else:
+            msg = self.messages.find_one_and_delete(
+                {'queue': queue},
+                sort=[('priority', pymongo.ASCENDING)],
+            )
+
+        if self.ttl:
+            self._update_queues_expire(queue)
+
+        if msg is None:
+            raise Empty()
+
+        return loads(bytes_to_str(msg['payload']))
+
+    def _size(self, queue):
+        # Do not calculate actual queue size if requested
+        # for performance considerations
+        if not self.calc_queue_size:
+            return super()._size(queue)
+
+        if queue in self._fanout_queues:
+            return self._get_broadcast_cursor(queue).get_size()
+
+        return self.messages.count_documents({'queue': queue})
+
+    def _put(self, queue, message, **kwargs):
+        data = {
+            'payload': dumps(message),
+            'queue': queue,
+            'priority': self._get_message_priority(message, reverse=True)
+        }
+
+        if self.ttl:
+            data['expire_at'] = self._get_queue_expire(queue, 'x-message-ttl')
+            msg_expire = self._get_message_expire(message)
+            if msg_expire is not None and (
+                data['expire_at'] is None or msg_expire < data['expire_at']
+            ):
+                data['expire_at'] = msg_expire
+
+        self.messages.insert_one(data)
+
+    def _put_fanout(self, exchange, message, routing_key, **kwargs):
+        self.broadcast.insert_one({'payload': dumps(message),
+                                  'queue': exchange})
+
+    def _purge(self, queue):
+        size = self._size(queue)
+
+        if queue in self._fanout_queues:
+            self._get_broadcast_cursor(queue).purge()
+        else:
+            self.messages.delete_many({'queue': queue})
+
+        return size
+
+    def get_table(self, exchange):
+        localRoutes = frozenset(self.state.exchanges[exchange]['table'])
+        brokerRoutes = self.routing.find(
+            {'exchange': exchange}
+        )
+
+        return localRoutes | frozenset(
+            (r['routing_key'], r['pattern'], r['queue'])
+            for r in brokerRoutes
+        )
+
+    def _queue_bind(self, exchange, routing_key, pattern, queue):
+        if self.typeof(exchange).type == 'fanout':
+            self._create_broadcast_cursor(
+                exchange, routing_key, pattern, queue)
+            self._fanout_queues[queue] = exchange
+
+        lookup = {
+            'exchange': exchange,
+            'queue': queue,
+            'routing_key': routing_key,
+            'pattern': pattern,
+        }
+
+        data = lookup.copy()
+
+        if self.ttl:
+            data['expire_at'] = self._get_queue_expire(queue, 'x-expires')
+
+        self.routing.update_one(lookup, {'$set': data}, upsert=True)
+
+    def queue_delete(self, queue, **kwargs):
+        self.routing.delete_many({'queue': queue})
+
+        if self.ttl:
+            self.queues.delete_one({'_id': queue})
+
+        super().queue_delete(queue, **kwargs)
+
+        if queue in self._fanout_queues:
+            try:
+                cursor = self._broadcast_cursors.pop(queue)
+            except KeyError:
+                pass
+            else:
+                cursor.close()
+
+                self._fanout_queues.pop(queue)
+
+    # Implementation details
+
+    def _parse_uri(self, scheme='mongodb://'):
+        # See mongodb uri documentation:
+        # https://docs.mongodb.org/manual/reference/connection-string/
+        client = self.connection.client
+        hostname = client.hostname
+
+        if hostname.startswith('srv://'):
+            scheme = 'mongodb+srv://'
+            hostname = 'mongodb+' + hostname
+
+        if not hostname.startswith(scheme):
+            hostname = scheme + hostname
+
+        if not hostname[len(scheme):]:
+            hostname += self.default_hostname
+
+        if client.userid and '@' not in hostname:
+            head, tail = hostname.split('://')
+
+            credentials = client.userid
+            if client.password:
+                credentials += ':' + client.password
+
+            hostname = head + '://' + credentials + '@' + tail
+
+        port = client.port if client.port else self.default_port
+
+        # We disable validating and normalization parameters here,
+        # because pymongo will validate and normalize parameters later in __init__ of MongoClient
+        parsed = uri_parser.parse_uri(hostname, port, validate=False)
+
+        dbname = parsed['database'] or client.virtual_host
+
+        if dbname in ('/', None):
+            dbname = self.default_database
+
+        options = {
+            'auto_start_request': True,
+            'ssl': self.ssl,
+            'connectTimeoutMS': (int(self.connect_timeout * 1000)
+                                 if self.connect_timeout else None),
+        }
+        options.update(parsed['options'])
+        options = self._prepare_client_options(options)
+
+        if 'tls' in options:
+            options.pop('ssl')
+
+        return hostname, dbname, options
+
+    def _prepare_client_options(self, options):
+        if pymongo.version_tuple >= (3,):
+            options.pop('auto_start_request', None)
+            if isinstance(options.get('readpreference'), int):
+                modes = pymongo.read_preferences._MONGOS_MODES
+                options['readpreference'] = modes[options['readpreference']]
+        return options
+
+    def prepare_queue_arguments(self, arguments, **kwargs):
+        return to_rabbitmq_queue_arguments(arguments, **kwargs)
+
+    def _open(self, scheme='mongodb://'):
+        hostname, dbname, conf = self._parse_uri(scheme=scheme)
+
+        conf['host'] = hostname
+
+        env = _detect_environment()
+        if env == 'gevent':
+            from gevent import monkey
+            monkey.patch_all()
+        elif env == 'eventlet':
+            from eventlet import monkey_patch
+            monkey_patch()
+
+        mongoconn = MongoClient(**conf)
+        database = mongoconn[dbname]
+
+        version_str = mongoconn.server_info()['version']
+        version_str = version_str.split('-')[0]
+        version = tuple(map(int, version_str.split('.')))
+
+        if version < (1, 3):
+            raise VersionMismatch(E_SERVER_VERSION.format(version_str))
+        elif self.ttl and version < (2, 2):
+            raise VersionMismatch(E_NO_TTL_INDEXES.format(version_str))
+
+        return database
+
+    def _create_broadcast(self, database):
+        """Create capped collection for broadcast messages."""
+        if self.broadcast_collection in database.list_collection_names():
+            return
+
+        database.create_collection(self.broadcast_collection,
+                                   size=self.capped_queue_size,
+                                   capped=True)
+
+    def _ensure_indexes(self, database):
+        """Ensure indexes on collections."""
+        messages = database[self.messages_collection]
+        messages.create_index(
+            [('queue', 1), ('priority', 1), ('_id', 1)], background=True,
+        )
+
+        database[self.broadcast_collection].create_index([('queue', 1)])
+
+        routing = database[self.routing_collection]
+        routing.create_index([('queue', 1), ('exchange', 1)])
+
+        if self.ttl:
+            messages.create_index([('expire_at', 1)], expireAfterSeconds=0)
+            routing.create_index([('expire_at', 1)], expireAfterSeconds=0)
+
+            database[self.queues_collection].create_index(
+                [('expire_at', 1)], expireAfterSeconds=0)
+
+    def _create_client(self):
+        """Actually creates connection."""
+        database = self._open()
+        self._create_broadcast(database)
+        self._ensure_indexes(database)
+
+        return database
+
+    @cached_property
+    def client(self):
+        return self._create_client()
+
+    @cached_property
+    def messages(self):
+        return self.client[self.messages_collection]
+
+    @cached_property
+    def routing(self):
+        return self.client[self.routing_collection]
+
+    @cached_property
+    def broadcast(self):
+        return self.client[self.broadcast_collection]
+
+    @cached_property
+    def queues(self):
+        return self.client[self.queues_collection]
+
+    def _get_broadcast_cursor(self, queue):
+        try:
+            return self._broadcast_cursors[queue]
+        except KeyError:
+            # Cursor may be absent when Channel created more than once.
+            # _fanout_queues is a class-level mutable attribute so it's
+            # shared over all Channel instances.
+            return self._create_broadcast_cursor(
+                self._fanout_queues[queue], None, None, queue,
+            )
+
+    def _create_broadcast_cursor(self, exchange, routing_key, pattern, queue):
+        if pymongo.version_tuple >= (3, ):
+            query = {
+                'filter': {'queue': exchange},
+                'cursor_type': CursorType.TAILABLE,
+            }
+        else:
+            query = {
+                'query': {'queue': exchange},
+                'tailable': True,
+            }
+
+        cursor = self.broadcast.find(**query)
+        ret = self._broadcast_cursors[queue] = BroadcastCursor(cursor)
+        return ret
+
+    def _get_message_expire(self, message):
+        value = message.get('properties', {}).get('expiration')
+        if value is not None:
+            return self.get_now() + datetime.timedelta(milliseconds=int(value))
+
+    def _get_queue_expire(self, queue, argument):
+        """Get expiration header named `argument` of queue definition.
+
+        Note:
+        ----
+            `queue` must be either queue name or options itself.
+        """
+        if isinstance(queue, str):
+            doc = self.queues.find_one({'_id': queue})
+
+            if not doc:
+                return
+
+            data = doc['options']
+        else:
+            data = queue
+
+        try:
+            value = data['arguments'][argument]
+        except (KeyError, TypeError):
+            return
+
+        return self.get_now() + datetime.timedelta(milliseconds=value)
+
+    def _update_queues_expire(self, queue):
+        """Update expiration field on queues documents."""
+        expire_at = self._get_queue_expire(queue, 'x-expires')
+
+        if not expire_at:
+            return
+
+        self.routing.update_many(
+            {'queue': queue}, {'$set': {'expire_at': expire_at}})
+        self.queues.update_many(
+            {'_id': queue}, {'$set': {'expire_at': expire_at}})
+
+    def get_now(self):
+        """Return current time in UTC."""
+        return datetime.datetime.utcnow()
+
+
+class Transport(virtual.Transport):
+    """MongoDB Transport."""
+
+    Channel = Channel
+
+    can_parse_url = True
+    polling_interval = 1
+    default_port = Channel.default_port
+    connection_errors = (
+        virtual.Transport.connection_errors + (errors.ConnectionFailure,)
+    )
+    channel_errors = (
+        virtual.Transport.channel_errors + (
+            errors.ConnectionFailure,
+            errors.OperationFailure)
+    )
+    driver_type = 'mongodb'
+    driver_name = 'pymongo'
+
+    implements = virtual.Transport.implements.extend(
+        exchange_type=frozenset(['direct', 'topic', 'fanout']),
+    )
+
+    def driver_version(self):
+        return pymongo.version
+
+    def as_uri(self, uri: str, include_password=False, mask='**') -> str:
+        if not uri:
+            return 'mongodb://'
+        if include_password:
+            return uri
+
+        if ',' not in uri:
+            return maybe_sanitize_url(uri)
+
+        uri1, remainder = uri.split(',', 1)
+        return ','.join([maybe_sanitize_url(uri1), remainder])