workercommon 0.4.1__py3-none-any.whl

workercommon/rabbitmqueue.py

@@ -0,0 +1,455 @@
+#!/usr/bin/env python3
+import sys
+
+if sys.version_info < (3, 8):
+    raise RuntimeError("At least Python 3.8 is required")
+
+import logging
+import json
+import pika.connection
+import pika.channel
+from pika import (
+    BasicProperties,
+    BlockingConnection,
+    SelectConnection,
+    URLParameters,
+)
+from pika.spec import PERSISTENT_DELIVERY_MODE, TRANSIENT_DELIVERY_MODE
+from concurrent.futures import Future
+from functools import partial
+from threading import RLock, Lock, Thread
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    List,
+    NamedTuple,
+    Optional,
+    Tuple,
+)
+from urllib.parse import quote as urlquote
+from datetime import datetime
+from collections import namedtuple
+from copy import copy
+
+
+class Parameters(NamedTuple):
+    mqparams: pika.connection.Parameters
+    queue: str
+    exchange: str
+    durable: bool
+    prefetch: int
+    routing_key: str
+
+    @classmethod
+    def of(
+        cls,
+        host: str,
+        port: int,
+        queue: str,
+        exchange: Optional[str] = None,
+        durable: bool = False,
+        username: str = "guest",
+        password: str = "guest",
+        prefetch: int = 100,
+        vhost: str = "/",
+        routing_key: str = "",
+    ) -> "Parameters":
+        cleaned_username = urlquote(username)
+        cleaned_password = urlquote(password)
+        vhost = urlquote(vhost, safe="")
+
+        parameters = URLParameters(
+            f"amqp://{cleaned_username}:{cleaned_password}@{host}:{port}/{vhost}"
+        )
+        cleaned_exchange = exchange if exchange else ""
+        return cls(parameters, queue, cleaned_exchange, durable, prefetch, routing_key)
+
+
+build_parameters = Parameters.of
+
+
+class BaseQueue(object):
+    def get_channel(self) -> pika.channel.Channel:
+        raise NotImplementedError
+
+    def get_parameters(self) -> Parameters:
+        raise NotImplementedError
+
+    @classmethod
+    def serialize_for_json(cls, obj: Any) -> str:
+        if isinstance(obj, datetime):
+            return obj.isoformat()
+
+        raise TypeError(f"Cannot serialize: {obj}")
+
+    def send_callback(self, message: Any, persist: bool = False) -> None:
+        properties = BasicProperties(
+            content_type="application/json",
+            delivery_mode=(
+                PERSISTENT_DELIVERY_MODE if persist else TRANSIENT_DELIVERY_MODE
+            ),
+        )
+
+        parameters = self.get_parameters()
+        self.get_channel().basic_publish(
+            exchange=parameters.exchange,
+            routing_key=parameters.routing_key,
+            body=json.dumps(message, default=self.serialize_for_json),
+            properties=properties,
+        )
+
+    def send(self, message: Any, persist: bool = False) -> None:
+        self.add_callback_threadsafe(
+            partial(self.send_callback, message, persist)
+        ).result(60)
+
+    def add_callback_threadsafe(self, callback: Callable[[], Any]) -> Future:
+        raise NotImplementedError
+
+    def run(self) -> None:
+        raise NotImplementedError
+
+    @staticmethod
+    def futurify(callback: Callable[[], Any]) -> Tuple[Callable[[], None], Future]:
+        future: Future = Future()
+
+        def wrapped():
+            try:
+                future.set_result(callback())
+
+            except BaseException as e:
+                future.set_exception(e)
+
+        return wrapped, future
+
+
+ToSend = namedtuple("ToSend", ["message", "persist"])
+
+
+class SendQueue(BaseQueue):
+    @classmethod
+    def build_channel(
+        cls,
+        connection: pika.connection.Connection,
+        parameters: Parameters,
+        **queue_arguments: Any,
+    ) -> pika.channel.Channel:
+        channel = connection.channel()
+        channel.exchange_declare(
+            exchange=parameters.exchange, durable=parameters.durable
+        )
+        channel.queue_declare(
+            queue=parameters.queue,
+            durable=parameters.durable,
+            arguments=queue_arguments,
+        )
+        channel.queue_bind(
+            queue=parameters.queue,
+            exchange=parameters.exchange,
+            routing_key=parameters.routing_key,
+        )
+
+        return channel
+
+    def __init__(self, parameters: Parameters, **queue_arguments: Any):
+        self.lock = RLock()
+        self.parameters = parameters
+        self.connection: Optional[BlockingConnection] = None
+        self.channel: Optional[pika.channel.Channel] = None
+        try:
+            self.connection = BlockingConnection(self.parameters.mqparams)
+            self.channel = self.build_channel(self.connection, self.parameters, **queue_arguments)
+        except BaseException:
+            self.close()
+            raise
+
+    def get_channel(self) -> pika.channel.Channel:
+        if self.channel is None:
+            raise RuntimeError("Channel is not available")
+        return self.channel
+
+    def get_parameters(self) -> Parameters:
+        return self.parameters
+
+    def get_connection(self) -> BlockingConnection:
+        with self.lock:
+            if self.connection is None:
+                raise RuntimeError("Connection is not available")
+            return self.connection
+
+    def close(self) -> None:
+        with self.lock:
+            if self.connection is not None:
+                self.connection.close()
+                self.connection = None
+            self.channel = None
+
+    @classmethod
+    def quick_send(cls, parameters: Parameters, messages_iter: Iterable[ToSend]) -> None:
+        messages = list(messages_iter)
+        if not messages:
+            return
+        queue: Optional[SendQueue] = None
+        try:
+            queue = cls(parameters)
+            while messages:
+                msg = messages.pop()
+                try:
+                    queue.send(msg.message, msg.persist)
+                except BaseException:
+                    messages.insert(0, msg)
+        except BaseException:
+            logging.exception(f"Failed to send messages: {messages}")
+            raise
+        finally:
+            if queue is not None:
+                queue.close()
+
+    def add_callback_threadsafe(self, callback: Callable[[], Any]) -> Future:
+        wrapped, future = self.futurify(callback)
+        try:
+            self.get_connection().add_callback_threadsafe(wrapped)
+        except BaseException as e:
+            future.set_exception(e)
+        return future
+
+    def run(self) -> None:
+        logging.info(f"Starting data events for {self.get_connection()}")
+        while True:
+            self.get_connection().process_data_events(None)
+
+
+class SendQueueWrapper(object):
+    "Wraps a SendQueue and recreates on error."
+    TRIES = 3
+
+    def __init__(self, parameters: Parameters, **queue_arguments: Any):
+        self.lock = Lock()
+        self.parameters = parameters
+        self.queue_arguments = queue_arguments
+        self.queue: Optional[SendQueue] = SendQueue(self.parameters, **self.queue_arguments)
+        self.queue_thread: Optional[Thread] = self.create_queue_thread(self.queue)
+        self.closed = False
+
+    def close(self) -> None:
+        with self.lock:
+            self.closed = True
+            if self.queue is not None:
+                self.queue.close()
+                self.queue = None
+
+            self.join_queue_thread()
+
+    @staticmethod
+    def create_queue_thread(queue: SendQueue) -> Thread:
+        queue_thread = Thread(target=queue.run)
+        queue_thread.start()
+        return queue_thread
+
+    def join_queue_thread(self) -> None:
+        "Must be called with self.lock held!"
+        if self.queue_thread is not None:
+            try:
+                self.queue_thread.join(10)
+            except BaseException:
+                logging.exception(f"Failed to wait for queue thread {self.queue_thread}")
+            self.queue_thread = None
+
+    def send(self, message: Any, persist: bool = False) -> None:
+        for i in range(self.TRIES):
+            with self.lock:
+                if self.closed:
+                    raise RuntimeError("Queue is closed")
+                try:
+                    if self.queue is None:
+                        self.queue = SendQueue(self.parameters, **self.queue_arguments)
+                        self.join_queue_thread()
+                        self.queue_thread = self.create_queue_thread(self.queue)
+                        logging.info("Re-established broken connection")
+
+                    self.queue.send(message, persist)
+                    break
+
+                except BaseException:
+                    logging.exception("Failed to send.  Attempting to reconnect")
+                    if self.queue is not None:
+                        try:
+                            self.queue.close()
+                        except BaseException:
+                            logging.warning("Failed to close queue")
+                        self.queue = None
+
+                    self.join_queue_thread()
+
+        else:
+            logging.error(f"Failed to send after {self.TRIES} tries: {message}")
+
+
+Callback = Callable[
+    [pika.channel.Channel, pika.spec.Basic.Deliver, pika.spec.BasicProperties, str], None
+]
+DisconnectCallback = Callable[[bool, List[ToSend]], None]
+
+
+class ReceiveSendQueue(BaseQueue):
+    def __init__(
+        self,
+        parameters: Parameters,
+        on_message: Callback,
+        on_disconnect: DisconnectCallback,
+        **queue_arguments: Any,
+    ):
+        self.lock = RLock()
+        self.outbox: List[ToSend] = []
+        self.parameters = parameters
+        self.on_message = on_message
+        self.on_disconnect = on_disconnect
+        self.terminated = False
+        self.queue_arguments = queue_arguments
+        self.connection: Optional[pika.connection.Connection] = None
+        self.channel: Optional[pika.channel.Channel] = None
+        self.init()
+
+    def add_to_outbox(self, items: Iterable[ToSend]) -> None:
+        self.outbox.extend(items)
+
+    def get_outbox(self) -> List[ToSend]:
+        with self.lock:
+            return copy(self.outbox)
+
+    def init(self):
+        with self.lock:
+            try:
+                self.terminated = False
+                self.connecion = None
+                self.channel = None
+                self.connection = SelectConnection(
+                    parameters=self.parameters.mqparams,
+                    on_open_callback=self.on_connection_open,
+                    on_close_callback=self.on_connection_closed,
+                )
+            except BaseException:
+                logging.execption("Failed to initialize connection")
+                self.close()
+                raise
+
+    def recovery(self, channel: pika.channel.Channel) -> None:
+        with self.lock:
+            while self.outbox:
+                to_send = self.outbox.pop(0)
+                self.send(to_send.message, to_send.persist)
+
+    def get_channel(self) -> pika.channel.Channel:
+        with self.lock:
+            if self.channel is None:
+                raise RuntimeError("Channel is not available")
+            return self.channel
+
+    def get_parameters(self) -> Parameters:
+        return self.parameters
+
+    def get_connection(self) -> SelectConnection:
+        with self.lock:
+            if self.connection is None:
+                raise RuntimeError("Connection is not available")
+            return self.connection
+
+    def send(self, message: Any, persist: bool = False) -> None:
+        with self.lock:
+            self.outbox.append(ToSend(message, persist))
+        super().send(message, persist)
+        with self.lock:
+            del self.outbox[-1]
+
+    def close(self) -> None:
+        try:
+            with self.lock:
+                if self.outbox:
+                    logging.info(f"Trying to empty outbox before closing: {self.outbox}")
+                    self.recovery(self.channel)
+        except BaseException:
+            logging.exception(f"Failed to empty outbox: {self.outbox}")
+
+        connection: Optional[pika.connection.Connection] = None
+        with self.lock:
+            self.terminated = True
+            self.channel = None
+            connection = self.connection
+            self.connection = None
+
+        if connection is not None:
+            try:
+                connection.ioloop.stop()
+            except BaseException:
+                logging.exception("Failed to stop ioloop")
+            finally:
+                try:
+                    connection.close()
+                except BaseException:
+                    logging.exception("Failed to close connection")
+
+        with self.lock:
+            if self.outbox:
+                logging.warning(f"Outbound items not sent: {self.outbox}")
+
+    def on_connection_closed(
+        self, connection: pika.connection.Connection, error: Any
+    ) -> None:
+        with self.lock:
+            if not self.terminated:
+                logging.warning(
+                    f"Flagging to set up a new connection because the old one ({connection}) was lost: {error}"
+                )
+                try:
+                    connection.close()
+                finally:
+                    try:
+                        connection.ioloop.stop()
+                    finally:
+                        self.on_disconnect(True, copy(self.outbox))
+            else:
+                self.on_disconnect(False, copy(self.outbox))
+
+    def on_connection_open(self, connection: pika.connection.Connection) -> None:
+        self.channel = connection.channel(on_open_callback=self.on_channel_open)
+
+    def on_channel_open(self, channel: pika.channel.Channel) -> None:
+        channel.basic_qos(prefetch_count=self.parameters.prefetch)
+        channel.queue_declare(
+            queue=self.parameters.queue,
+            durable=self.parameters.durable,
+            callback=lambda _: self.on_queue_declare(channel),
+            arguments=self.queue_arguments,
+        )
+
+    def on_queue_declare(self, channel: pika.channel.Channel) -> None:
+        channel.exchange_declare(
+            exchange=self.parameters.exchange,
+            durable=self.parameters.durable,
+            callback=lambda _: channel.queue_bind(
+                queue=self.parameters.queue,
+                exchange=self.parameters.exchange,
+                routing_key=self.parameters.routing_key,
+                callback=lambda _: self.on_channel_ready(channel),
+            ),
+        )
+
+    def on_channel_ready(self, channel: pika.channel.Channel) -> None:
+        self.recovery(channel)
+        channel.basic_consume(
+            on_message_callback=self.on_message,
+            queue=self.parameters.queue,
+            auto_ack=False,
+        )
+
+    def add_callback_threadsafe(self, callback: Callable[[], Any]) -> Future:
+        wrapped, future = self.futurify(callback)
+        try:
+            self.get_connection().ioloop.add_callback_threadsafe(callback)
+        except BaseException as e:
+            future.set_exception(e)
+        return future
+
+    def run(self) -> None:
+        self.get_connection().ioloop.start()

workercommon/test.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python3
+import sys
+
+if sys.version_info < (3, 8):
+    raise RuntimeError("At least Python 3.8 is required")
+import unittest
+
+from os import environ
+from getpass import getpass
+import logging
+from database import *
+from typing import List
+
+
+class DBTesting(object):
+    DB_PASSWORD: List[str] = []
+
+    @classmethod
+    def get_password(cls) -> str:
+        if not cls.DB_PASSWORD:
+            cls.DB_PASSWORD.append(getpass("Password: "))
+        return cls.DB_PASSWORD[0]
+
+    @classmethod
+    def build_connection(cls) -> Connection:
+        host = environ.get("DB_HOST", None)
+
+        port_number = 5432
+        if port := environ.get("DB_PORT", ""):
+            port_number = int(port)
+
+        username = environ["DB_USER"]
+        password = cls.get_password()
+        database = environ["DB_DATABASE"]
+
+        return PgConnection(host, port_number, username, password, database)
+
+
+class TestConnect(unittest.TestCase, DBTesting):
+    def test_connect(self):
+        self.build_connection().close()
+
+
+class TestOperations(unittest.TestCase, DBTesting):
+    def setUp(self):
+        self.connection = self.build_connection()
+
+    def tearDown(self):
+        self.connection.close()
+        self.connection = None
+
+    def test_cursor(self):
+        cursor = self.connection.cursor()
+        self.assertIsNotNone(cursor)
+        with self.assertRaises(RuntimeError):
+            cursor.get()
+
+        with cursor:
+            self.assertIsNotNone(cursor.get())
+
+        with self.assertRaises(RuntimeError):
+            cursor.get()
+
+    def test_direct_table(self):
+        with self.connection.cursor() as cursor:
+            cursor.execute("CREATE TEMPORARY TABLE test(value1 INTEGER, value2 TEXT)")
+            cursor.execute("INSERT INTO test(value1, value2) VALUES(%s, %s)", 1, "foo")
+
+            # Basic functionality
+            row = cursor.execute_direct(
+                "SELECT value1, value2 FROM test LIMIT 1"
+            ).fetchone()
+            self.assertIsNotNone(row)
+            self.assertTrue(cursor.columns)
+            self.assertEqual(row[0], 1)
+            self.assertEqual(row[1], "foo")
+
+            # Advanced functionality
+            rows = list(cursor.execute("SELECT value1, value2 FROM test LIMIT 1"))
+            self.assertEqual(len(rows), 1)
+            dict_row = rows[0]
+            self.assertEqual(dict_row["value1"], 1)
+            self.assertEqual(dict_row["value2"], "foo")
+
+    def test_table(self):
+        table_type = PgBaseTable("table_test", "key", ["key", "value1", "value2"])
+        with self.connection.cursor() as cursor:
+            table_cursor = TableCursor(cursor, table_type)
+
+            def verify2():
+                row = table_cursor[2]
+                self.assertTrue(row)
+                self.assertEqual(row["key"], 2)
+                self.assertEqual(row["value1"], 30)
+                self.assertEqual(row["value2"], "brown")
+
+            cursor.execute(
+                "CREATE TEMPORARY TABLE table_test(key SERIAL PRIMARY KEY NOT NULL, value1 INTEGER, value2 TEXT)"
+            )
+            # Test insert and get
+            self.assertEqual(table_type.insert(cursor, value1=100, value2="cow"), 1)
+
+            row = table_type.select_by_id(cursor, 1)
+            self.assertTrue(row)
+            self.assertEqual(row["key"], 1)
+            self.assertEqual(row["value1"], 100)
+            self.assertEqual(row["value2"], "cow")
+
+            # Test insert again and get
+            self.assertEqual(table_type.insert(cursor, value1=30, value2="brown"), 2)
+
+            verify2()
+
+            # Test selecting by other column values
+            rows = list(table_cursor.where(100, value1=100))
+            self.assertEqual(len(rows), 1)
+            row = rows[0]
+            self.assertEqual(row["key"], 1)
+            self.assertEqual(row["value1"], 100)
+            self.assertEqual(row["value2"], "cow")
+
+            # Test custom selecting
+            rows = list(table_cursor.custom_where("value1 > %s", 50))
+            self.assertEqual(len(rows), 1)
+            row = rows[0]
+            self.assertEqual(row["key"], 1)
+            self.assertEqual(row["value1"], 100)
+            self.assertEqual(row["value2"], "cow")
+
+            # Test selecting by other column values where a None is present
+            self.assertEqual(table_type.insert(cursor, value1=None, value2="skittle"), 3)
+            rows = list(table_cursor.where(100, value1=None))
+            self.assertEqual(len(rows), 1)
+            row = rows[0]
+            self.assertEqual(row["key"], 3)
+            self.assertIsNone(row["value1"])
+            self.assertEqual(row["value2"], "skittle")
+
+            # Test upserting on an alternate column
+            self.assertEqual(
+                table_cursor.upsert_on_column("value1", None, value2="1234"), 3
+            )
+            rows = list(table_cursor.where(100, value1=None))
+            self.assertEqual(len(rows), 1)
+            row = rows[0]
+            self.assertEqual(row["key"], 3)
+            self.assertIsNone(row["value1"])
+            self.assertEqual(row["value2"], "1234")
+
+            # Test upserting
+            table_cursor[1] = {"value1": 200, "value2": "quack"}
+            self.assertEqual(table_type.upsert(cursor, 1, value1=200, value2="quack"), 1)
+
+            row = table_type.select_by_id(cursor, 1)
+            self.assertTrue(row)
+            self.assertEqual(row["key"], 1)
+            self.assertEqual(row["value1"], 200)
+            self.assertEqual(row["value2"], "quack")
+
+            verify2()
+
+            # Test deletion
+            self.assertTrue(table_type.delete(cursor, 1))
+            with self.assertRaises(KeyError):
+                table_type.select_by_id(cursor, 1)
+            verify2()
+
+            self.assertFalse(table_type.delete(cursor, 1))
+
+            with self.assertRaises(KeyError):
+                del table_cursor[1]
+
+            del table_cursor[2]
+
+
+logging.basicConfig(level=logging.DEBUG, stream=sys.stderr)
+unittest.main()