lorrystream 0.0.2__py3-none-any.whl

lorrystream/streamz/sinks.py

@@ -0,0 +1,114 @@
+# Copyright (c) 2013-2023, The Kotori developers and contributors.
+# Distributed under the terms of a BSD-3-Clause license, see LICENSE.
+
+import logging
+import re
+import typing as t
+
+import pandas as pd
+import sqlalchemy as sa
+from sqlalchemy import Engine
+from streamz import Sink, Stream
+
+from lorrystream.exceptions import InvalidSinkError
+
+logger = logging.getLogger(__name__)
+
+
+@Stream.register_api()
+class dataframe_to_sql(Sink):
+    """
+    Store data into SQLAlchemy-compatible database.
+
+    Requires ``sqlalchemy``
+
+    :param dburi: str
+        SQLAlchemy connection URI.
+    :param engine_options:
+        Propagated to SQLAlchemy's ``create_engine(**kwargs)``.
+    """
+
+    def __init__(self, upstream, dburi, engine_options=None, **kwargs):
+        self.dburi = dburi
+        self.engine: t.Union[Engine, None] = None
+        self.engine_options = engine_options or {}
+        self.table_name = None
+        self.if_exists = "append"
+        self.method = None
+        self.chunksize = 10_000
+        super().__init__(upstream, ensure_io_loop=True, **kwargs)
+
+    def update(self, x, who=None, metadata=None):
+        """
+        Store packets into database.
+        """
+        df: pd.DataFrame = x
+        logger.info(f"to_sql.update: who={who}, metadata={metadata}")
+        df.info()
+        print(df)  # noqa: T201
+
+        if self.engine is None:
+            logger.info(f"Connecting to {self.dburi}")
+            # TODO: Improve.
+            matches = re.match(r"^.*table=([\w-]*)", self.dburi)
+            if matches:
+                self.table_name = matches.group(1)
+            else:
+                raise InvalidSinkError("Unable to obtain table name")
+            matches = re.match(r"^.*if_exists=(\w*)", self.dburi)
+            if matches:
+                self.if_exists = matches.group(1)
+            dburi = re.sub(r"\?.*", "", self.dburi)
+            logger.info(f"Effective dburi: {dburi}")
+            logger.info(f"Writing to table: {self.table_name}, if_exists={self.if_exists}")
+            self.engine = sa.create_engine(dburi, **self.engine_options)
+
+            # Use CrateDB bulk operations endpoint for improved efficiency.
+            if self.dburi.startswith("crate"):
+                self.method = self.insert_bulk
+
+        df.to_sql(
+            name=self.table_name,
+            con=self.engine,
+            if_exists=self.if_exists,
+            index=False,
+            chunksize=self.chunksize,
+            method=self.method,
+        )
+
+    @staticmethod
+    def insert_bulk(pd_table, conn, keys, data_iter):
+        """
+        A fast insert method for pandas and Dask, using CrateDB's "bulk operations" endpoint.
+
+        The idea is to break out of SQLAlchemy, compile the insert statement, and use the raw
+        DBAPI connection client, in order to invoke a request using `bulk_parameters`::
+
+            cursor.execute(sql=sql, bulk_parameters=data)
+
+        - https://crate.io/docs/crate/reference/en/5.2/interfaces/http.html#bulk-operations
+        """
+
+        """
+        # Vanilla
+        # This is the regular implementation, using SQLAlchemy.
+        data = [dict(zip(keys, row)) for row in data_iter]
+        conn.execute(pd_table.table.insert(), data)
+        """
+
+        # Bulk
+        sql = str(pd_table.table.insert().compile(bind=conn))
+        data = list(data_iter)
+
+        logger.info(f"Bulk SQL:     {sql}")
+        logger.info(f"Bulk records: {len(data)}")
+
+        cursor = conn._dbapi_connection.cursor()
+        cursor.execute(sql=sql, bulk_parameters=data)
+        cursor.close()
+
+    def destroy(self):
+        if self.engine is not None:
+            logger.info(f"Disconnecting from {self.dburi}")
+            self.engine.dispose()
+        super().destroy()

lorrystream/streamz/sources.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2013-2024, The Kotori developers and contributors.
+# Distributed under the terms of a BSD-3-Clause license, see LICENSE.
+import logging
+import queue
+import typing as t
+from collections import OrderedDict
+
+from streamz import Stream, from_mqtt, from_q
+
+from lorrystream.model import StreamAddress
+from lorrystream.streamz.amqp import AMQPAdapter, ReconnectingAMQPAdapter
+from lorrystream.streamz.model import URL, BusMessage
+from lorrystream.util.aio import AsyncThreadTask
+
+logger = logging.getLogger(__name__)
+
+
+@Stream.register_api()
+class FromAmqp(from_q):
+    """Read from AMQP source
+
+    See https://en.wikipedia.org/wiki/AMQP for a description of the protocol
+    and its uses. Requires the ``pika`` package.
+
+    - https://www.rabbitmq.com/uri-spec.html
+    - https://www.rabbitmq.com/uri-query-parameters.html
+
+    TODO: See also ``sinks.to_amqp``.
+
+    :param uri: str
+    :param reconnect: bool
+    """
+
+    def __init__(self, address: StreamAddress, **kwargs):
+        self.address = address
+        self.stopped = True
+        consumer_class: t.Callable = AMQPAdapter
+        if self.address.options.get("reconnect", True):
+            consumer_class = ReconnectingAMQPAdapter
+        self.consumer = consumer_class(address=address, on_message=self._on_message)
+        super().__init__(q=queue.Queue(), **kwargs)
+
+    def _delivery_to_dict(self, basic_deliver):
+        deliver_attrs = ["consumer_tag", "delivery_tag", "redelivered", "exchange", "routing_key"]
+        items = OrderedDict()
+        for attr in deliver_attrs:
+            items[attr] = getattr(basic_deliver, attr)
+        return items
+
+    def _properties_to_dict(self, properties):
+        items = OrderedDict()
+        for key, value in properties.__dict__.items():
+            if getattr(properties.__class__, key, None) != value:
+                items[key] = value
+        return items
+
+    def _on_message(self, channel, basic_deliver, properties, body):
+        busmsg = BusMessage.from_amqp_pika(self, channel, basic_deliver, properties, body)
+        self.q.put(busmsg)
+
+    async def run(self):
+        AsyncThreadTask(self.consumer.run).run()
+        await super().run()
+
+    def stop(self):
+        # TODO: Does not work with CTRL+C yet. Validate if it works on other occasions at least.
+        if not self.stopped:
+            self.consumer.stop()
+            self.consumer = None
+            self.stopped = True
+        super().stop()
+
+
+@Stream.register_api()
+class FromMqttPlus(from_mqtt):
+    def __init__(self, uri: URL, client_kwargs=None, **kwargs):
+        host = uri.host or "localhost"
+        port = uri.port or uri.default_port
+        topic = uri.path_unquoted.lstrip("/")
+        logger.info(f"Subscribing to MQTT topic '{topic}' at broker on '{host}")
+        super().__init__(host, port, topic, keepalive=60, client_kwargs=client_kwargs, **kwargs)
+
+    def _on_message(self, client, userdata, msg):
+        busmsg = BusMessage.from_mqtt_paho(self, client, userdata, msg)
+        self.q.put(busmsg)
+
+
+from_amqp = FromAmqp
+from_mqtt_plus = FromMqttPlus

lorrystream/util/about.py

@@ -0,0 +1,120 @@
+# Copyright (c) 2013-2023, The Kotori developers and contributors.
+# Distributed under the terms of the LGPLv3 license, see LICENSE.
+
+import platform
+import sys
+import textwrap
+import typing as t
+from collections import abc
+
+from colorama import Fore, Style
+
+string_or_list = t.Union[str, t.List[str], t.Iterable[str]]
+
+
+def section(text: str):
+    return armor(text, title)
+
+
+def subsection(text: str):
+    return armor(text, subtitle)
+
+
+def armor(text: str, formatter: t.Callable):
+    guard = "=" * len(text)
+    print(guard)
+    print(formatter(text))
+    print(guard)
+
+
+def title(text: str):
+    return Fore.CYAN + Style.BRIGHT + text + Style.RESET_ALL
+
+
+def subtitle(text: str):
+    return Fore.YELLOW + Style.BRIGHT + text + Style.RESET_ALL
+
+
+def text_list(data: string_or_list):
+    if isinstance(data, str):
+        return data
+    elif isinstance(data, t.List):
+        return ", ".join(data)
+    else:
+        raise ValueError("Unable to flatten list")
+
+
+def bullet_list(data: t.List[str]):
+    data.sort()
+    lines = [f"- {text_list(line)}" for line in data]
+    return "\n".join(lines)
+
+
+def wrap_list(data: t.List[str], **kwargs):
+    data = [item for item in data if item is not None]
+    data.sort()
+    text = ", ".join(data)
+    return "\n".join(textwrap.wrap(text, width=80, **kwargs))
+
+
+def bullet_item(data: string_or_list, label: str = None):
+    if data is None:
+        return None
+    if label is None:
+        label = ""
+    else:
+        label = f"{label}: "
+    if isinstance(data, (list, tuple, abc.KeysView, abc.ValuesView)):
+        text = wrap_list(list(data), subsequent_indent="  ")
+    else:
+        text = data
+    return f"- {label}{text}"
+
+
+class AboutReport:
+    @staticmethod
+    def platform():
+        section("Platform")
+        print()
+
+        subsection("Python")
+        print(bullet_item(platform.platform()))
+        print()
+
+        # SQLAlchemy
+        from importlib.metadata import entry_points
+
+        import sqlalchemy.dialects
+
+        subsection("SQLAlchemy")
+        print(bullet_item(sqlalchemy.dialects.__all__, label="Dialects built-in"))
+        dialects: t.List[str]
+        if sys.version_info >= (3, 10):
+            eps = entry_points(group="sqlalchemy.dialects")
+            dialects = [dialect.name for dialect in eps]
+        else:
+            dialects = []
+        print(bullet_item(dialects, label="Dialects 3rd-party"))
+        print(bullet_item(sqlalchemy.dialects.plugins.impls.keys(), label="Plugins"))
+        print()
+
+        # fsspec
+        try:
+            import fsspec
+
+            subsection("fsspec protocols")
+            print(bullet_item(fsspec.available_protocols()))
+            print()
+
+            subsection("fsspec compressions")
+            print(bullet_item(fsspec.available_compressions()))
+            print()
+        except ImportError:
+            pass
+
+        # pandas
+        subsection("pandas module versions")
+        import pandas
+
+        pandas.show_versions(as_json=False)
+        print()

lorrystream/util/aio.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2013-2023, The Kotori developers and contributors.
+# Distributed under the terms of the LGPLv3 license, see LICENSE.
+
+import asyncio
+import functools
+
+
+def make_sync(func):
+    """
+    Click entrypoint decorator for wrapping asynchronous functions.
+
+    https://github.com/pallets/click/issues/2033
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        return asyncio.run(func(*args, **kwargs))
+
+    return wrapper
+
+
+def tornado_asyncio_run_forever():
+    from tornado.platform.asyncio import AsyncIOMainLoop
+
+    AsyncIOMainLoop().install()
+
+    def run_asyncio_loop():
+        loop = asyncio.get_event_loop()
+        try:
+            loop.run_forever()
+        except KeyboardInterrupt:
+            pass
+        finally:
+            loop.close()
+
+    run_asyncio_loop()
+
+
+class AsyncThreadTask:
+    """
+    Starts a function using `asyncio.to_thread()`, with a new asyncio I/O loop.
+    """
+
+    def __init__(self, func, *args, **kwargs):
+        self.func = func
+        self.args = args
+        self.kwargs = kwargs
+        self.loop = asyncio.new_event_loop()
+
+    def run(self):
+        return self.launch()
+
+    def launch(self):
+        coro = asyncio.to_thread(self._thread)
+        return asyncio.create_task(coro)
+
+    def _thread(self):
+        asyncio.set_event_loop(self.loop)
+        return self.func(*self.args, **self.kwargs)

lorrystream/util/cli.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2013-2023, The Kotori developers and contributors.
+# Distributed under the terms of the LGPLv3 license, see LICENSE.
+
+import logging
+import textwrap
+import typing as t
+
+import click
+
+from lorrystream.util.common import setup_logging
+
+logger = logging.getLogger(__name__)
+
+
+def boot_click(ctx: click.Context, verbose: bool = False, debug: bool = False):
+    """
+    Bootstrap the CLI application.
+    """
+
+    # Adjust log level according to `verbose` / `debug` flags.
+    log_level = logging.WARNING
+    if verbose:
+        log_level = logging.INFO
+    if debug:
+        log_level = logging.DEBUG
+
+    # Setup logging, according to `verbose` / `debug` flags.
+    setup_logging(level=log_level)
+
+
+def to_list(x: t.Any, default: t.List[t.Any] = None) -> t.List[t.Any]:
+    if not isinstance(default, t.List):
+        raise ValueError("Default value is not a list")
+    if x is None:
+        return default
+    if not isinstance(x, t.Iterable) or isinstance(x, str):
+        return [x]
+    elif isinstance(x, list):
+        return x
+    else:
+        return list(x)
+
+
+def docstring_format_verbatim(text: t.Optional[str]) -> str:
+    """
+    Format docstring to be displayed verbatim as a help text by Click.
+
+    - https://click.palletsprojects.com/en/8.1.x/documentation/#preventing-rewrapping
+    - https://github.com/pallets/click/issues/56
+    """
+    text = text or ""
+    text = textwrap.dedent(text)
+    lines = [line if line.strip() else "\b" for line in text.splitlines()]
+    return "\n".join(lines)

lorrystream/util/common.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2013-2023, The Kotori developers and contributors.
+# Distributed under the terms of the LGPLv3 license, see LICENSE.
+
+import logging
+import shlex
+import subprocess
+import sys
+from textwrap import dedent
+
+import colorlog
+from colorlog.escape_codes import escape_codes
+
+logger = logging.getLogger(__name__)
+
+
+SQLALCHEMY_LOGGING = True
+
+
+def setup_logging_basic(level=logging.INFO):
+    log_format = "%(asctime)-15s [%(name)-24s] %(levelname)-8s: %(message)s"
+    logging.basicConfig(format=log_format, stream=sys.stderr, level=level)
+
+
+def setup_logging(level=logging.INFO):
+    reset = escape_codes["reset"]
+    log_format = f"%(asctime)-15s [%(name)-28s] %(log_color)s%(levelname)-8s:{reset} %(message)s"
+
+    handler = colorlog.StreamHandler()
+    handler.setFormatter(colorlog.ColoredFormatter(log_format))
+
+    logging.basicConfig(format=log_format, level=level, handlers=[handler])
+
+    # Enable SQLAlchemy logging.
+    if SQLALCHEMY_LOGGING:
+        logging.getLogger("sqlalchemy").setLevel(level)
+
+
+def get_version(appname):
+    from importlib.metadata import PackageNotFoundError, version  # noqa
+
+    try:
+        return version(appname)
+    except PackageNotFoundError:  # pragma: no cover
+        return "unknown"
+
+
+def run_command(command: str):
+    """
+    https://stackoverflow.com/a/48813330
+    """
+    command = dedent(command).strip()
+    cmd = shlex.split(command)
+
+    try:
+        output = subprocess.check_output(cmd, stderr=subprocess.STDOUT, universal_newlines=True)  # noqa: S603
+    except subprocess.CalledProcessError as exc:
+        logger.error(f"Command failed (exit code {exc.returncode}). The command was:\n{command}")
+        logger.error(exc.output)
+        raise
+    else:
+        if output:
+            logger.info(f"Command output:\n{output}")

lorrystream/util/data.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2013-2023, The Kotori developers and contributors.
+# Distributed under the terms of the LGPLv3 license, see LICENSE.
+
+import json
+import sys
+import typing as t
+
+import pandas as pd
+
+
+def jd(data: t.Any):
+    """
+    Pretty-print JSON with indentation.
+    """
+    print(json.dumps(data, indent=2))  # noqa: T201
+
+
+def chunker(seq: pd.DataFrame, chunksize: int):
+    """
+    Chunks generator function for iterating pandas Dataframes and Series.
+
+    https://stackoverflow.com/a/61798585
+    :return:
+    """
+    for pos in range(0, len(seq), chunksize):
+        yield seq.iloc[pos : pos + chunksize]
+
+
+def df_convert_datetimes(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Convert all datetime columns to ISO format.
+
+    :param df:        df[Columns.FROM_DATE] = df[Columns.FROM_DATE].dt.tz_localize(self.tz)
+    :return:
+    """
+    df = df.copy(deep=True)
+
+    date_columns = list(df.select_dtypes(include=[pd.DatetimeTZDtype]).columns)
+    # date_columns.extend([Columns.FROM_DATE.value, Columns.TO_DATE.value, Columns.DATE.value])  # noqa: ERA001
+    date_columns_set = set(date_columns)
+    for date_column in date_columns_set:
+        if date_column in df:
+            df[date_column] = df[date_column].apply(lambda d: d.isoformat() if pd.notna(d) else None)
+
+    return df
+
+
+def get_sqlalchemy_dialects() -> t.List[str]:
+    """
+    Return list of available SQLAlchemy dialects.
+
+    :return:
+    """
+    from importlib.metadata import entry_points
+
+    import sqlalchemy.dialects
+
+    builtins = sqlalchemy.dialects.__all__
+    more: t.List
+    if sys.version_info >= (3, 10):
+        eps = entry_points(group="sqlalchemy.dialects")
+        more = [dialect.name for dialect in eps]
+    else:
+        eps = entry_points()
+        more = [dialect[0].name for dialect in eps.values() if dialect[0].group == "sqlalchemy.dialects"]
+
+    return sorted(list(builtins) + list(more))
+
+
+def asbool(obj: t.Any) -> bool:
+    # from sqlalchemy.util.langhelpers
+    # from paste.deploy.converters
+    if isinstance(obj, str):
+        obj = obj.strip().lower()
+        if obj in ["true", "yes", "on", "y", "t", "1"]:
+            return True
+        elif obj in ["false", "no", "off", "n", "f", "0"]:
+            return False
+        else:
+            raise ValueError("String is not true/false: %r" % obj)
+    return bool(obj)
+
+
+def split_list(value: str, delimiter: str = ",") -> t.List[str]:
+    if value is None:
+        return []
+    return [c.strip() for c in value.split(delimiter)]