lorrystream 0.0.2__py3-none-any.whl

lorrystream/__init__.py

@@ -0,0 +1 @@
+from .cmd import parse_launch  # noqa: F401

lorrystream/cli.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 logging
+import typing as t
+
+import click
+
+from lorrystream import parse_launch
+from lorrystream.core import run_single
+from lorrystream.util.about import AboutReport
+from lorrystream.util.aio import make_sync
+from lorrystream.util.cli import boot_click, docstring_format_verbatim
+
+logger = logging.getLogger(__name__)
+
+
+def help_relay():
+    """
+    Import and export data into/from InfluxDB
+
+    SOURCE can be a file or a URL.
+    TARGET can be a file or a URL.
+
+    Synopsis
+    ========
+
+    # Relay messages from MQTT to CrateDB.
+    lorry relay \\
+        "mqtt://localhost/testdrive/#" \\
+        "crate://localhost/testdrive/data"
+
+    # Relay messages from AMQP to MQTT.
+    lorry relay \\
+        "amqp://localhost/testdrive/demo" \\
+        "mqtt://localhost/testdrive/demo"
+
+    """  # noqa: E501
+
+
+def help_launch():
+    """
+    Launch a LorryStream pipeline.
+
+    """  # noqa: E501
+
+
+@click.group()
+@click.version_option(package_name="lorrystream")
+@click.option("--verbose", is_flag=True, required=False, help="Turn on logging")
+@click.option("--debug", is_flag=True, required=False, help="Turn on logging with debug level")
+@click.pass_context
+def cli(ctx: click.Context, verbose: bool, debug: bool):
+    verbose = True
+    return boot_click(ctx, verbose, debug)
+
+
+@cli.command("info", help="Report about platform information")
+def info():
+    AboutReport.platform()
+
+
+@cli.command(
+    "launch",
+    help=docstring_format_verbatim(help_launch.__doc__),
+    context_settings={"max_content_width": 120},
+)
+@click.argument("command", nargs=-1)
+@click.pass_context
+def launch(ctx: click.Context, command: t.Tuple[str]):
+    logger.info("Starting")
+    command_str = " ".join(command)
+    parse_launch(command_str)
+
+
+@cli.command(
+    "relay",
+    help=docstring_format_verbatim(help_relay.__doc__),
+    context_settings={"max_content_width": 120},
+)
+@click.argument("source", type=str, required=True)
+@click.argument("sink", type=str, required=False)
+@click.pass_context
+@make_sync
+async def relay(ctx: click.Context, source: str, sink: str):
+    logger.info("Starting")
+    await run_single(source, sink)

lorrystream/cmd.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2013-2023, The Kotori developers and contributors.
+# Distributed under the terms of the LGPLv3 license, see LICENSE.
+
+import asyncio
+import logging
+import re
+import typing as t
+
+from lorrystream.core import run_single
+
+logger = logging.getLogger(__name__)
+
+
+def get_location(element: t.Union[str, None]) -> str:
+    if element is None:
+        raise ValueError("Unable to decode location from empty element")
+    matches = re.match(".*location=(.*)", element)
+    if not matches:
+        raise ValueError(f"Unable to decode location from element: {element}")
+    location = matches.group(1)
+    return location
+
+
+def parse_launch(command: str):
+    logger.info(f"Running command: {command}")
+    elements = map(str.strip, command.split("!"))
+    source = None
+    sink = None
+    for element in elements:
+        name, rest = element.split(" ")
+        if name.endswith("src"):
+            source = element
+        elif name.endswith("sink"):
+            sink = element
+    source_uri = get_location(source)
+    sink_uri = get_location(sink)
+    logger.info(f"Source: {source_uri}, Sink: {sink_uri}")
+    asyncio.run(run_single(source_uri, sink_uri))

lorrystream/core.py

@@ -0,0 +1,215 @@
+# Copyright (c) 2013-2023, The Kotori developers and contributors.
+# Distributed under the terms of the LGPLv3 license, see LICENSE.
+
+import asyncio
+import json
+import logging
+import operator
+import typing as t
+
+from streamz import Sink, Source, Stream
+from streamz.batch import Batch
+
+from lorrystream.exceptions import InvalidContentTypeError, InvalidSinkError, InvalidSourceError
+from lorrystream.model import Channel, Packet, SinkInputType, StreamAddress
+from lorrystream.streamz.model import BusMessage
+from lorrystream.util.data import get_sqlalchemy_dialects
+
+logger = logging.getLogger(__name__)
+
+
+async def run_single(source_uri: str, sink_uri: str):
+    """
+    Create a single channel and run it on the engine, blocking forever.
+
+    :param source_uri: Source element URI
+    :param sink_uri: Sink element URI
+    :return:
+    """
+
+    # Create channel.
+    channel = ChannelFactory(source=source_uri, sink=sink_uri).channel()
+
+    # Run channel with engine.
+    engine = Engine()
+    engine.register(channel)
+    await engine.run_forever()
+
+
+async def run_channels(*channels: Channel):
+    """
+    Obtain multiple channel objects and run them on the engine, blocking forever.
+
+    :param channels: List of channel objects
+    :return:
+    """
+    engine = Engine()
+    for channel in channels:
+        engine.register(channel)
+    await engine.run_forever()
+
+
+class Decoders:
+    @staticmethod
+    def decode_mqtt_json(msg):
+        """
+        Decode `MQTTMessage` object into `Packet` object.
+        """
+        return Packet(payload=json.loads(msg.payload), busmsg=msg)
+
+    @staticmethod
+    def decode_mqtt(msg):
+        """
+        Decode `MQTTMessage` object into `Packet` object.
+        """
+        return Packet(payload=json.loads(msg.payload), busmsg=msg)
+
+    @staticmethod
+    def decode_busmessage(busmsg: BusMessage):
+        """
+        Decode `BusMessage` object into `Packet` object.
+        """
+        return Packet(payload=busmsg.data.payload, busmsg=busmsg)
+
+    @staticmethod
+    def decode_json(packet: Packet):
+        """
+        Decode `MQTTMessage` object into `Packet` object.
+        """
+        packet.payload = json.loads(packet.payload)
+        return packet
+
+    @staticmethod
+    def decode_busmessage_json(busmsg: BusMessage):
+        """
+        Decode `BusMessage` object into `Packet` object.
+        """
+        if busmsg.data.payload is None:
+            return None
+        return Packet(payload=json.loads(busmsg.data.payload), busmsg=busmsg)
+
+
+class ChannelFactory:
+    def __init__(self, source: str, sink: SinkInputType):
+        self.source_element: Source = None
+        self.sink_element: t.Union[Sink, t.Callable] = None
+        self.transformers: t.List[t.Callable] = []
+        self.pipeline: t.Union[Batch, Stream] = None
+
+        # FIXME: Obtain parameters from user.
+        self.batch_size = 2
+        self.timeout = 0.25
+
+        # When no sink is specified, use STDOUT.
+        # TODO: How to use `sys.stdout.buffer.write` instead?
+        if sink is None:
+            sink = print
+
+        self.source_address = StreamAddress.from_url(source)
+        self.sink_address = StreamAddress.resolve_sink(sink)
+
+        # Select source element for pipeline.
+        self.select_source(source)
+
+        # Create first half of pipeline.
+        self.mkpipeline()
+
+        # Select sink element for pipeline.
+        self.select_sink(sink)
+
+    def select_source(self, location: str):
+        """
+        Select source element of pipeline.
+        """
+        uri = self.source_address.uri
+        if uri.scheme.startswith("amqp"):
+            logger.info("Subscribing to AMQP")
+            self.source_element = Stream.from_amqp(self.source_address)
+            self.transformers.append(Decoders.decode_busmessage)
+
+        elif uri.scheme.startswith("mqtt"):
+            self.source_element = Stream.from_mqtt_plus(uri)
+            self.transformers.append(Decoders.decode_busmessage)
+
+        else:
+            raise InvalidSourceError(f"Source scheme unknown: {uri.scheme}")
+
+        if "content-type" in self.source_address.options:
+            source_content_type = self.source_address.options["content-type"]
+            if source_content_type == "json":
+                self.transformers.append(Decoders.decode_json)
+            else:
+                raise InvalidContentTypeError(f"Invalid content type for source '{uri}': {source_content_type}")
+
+    def mkpipeline(self):
+        """
+        n: int
+            Maximum partition size
+        timeout: int or float, optional
+            Number of seconds after which a partition will be emitted,
+            even if its size is less than ``n``. If ``None`` (default),
+            a partition will be emitted only when its size reaches ``n``.
+        """
+        self.pipeline = self.source_element.partition(n=self.batch_size, timeout=self.timeout).to_batch()
+        for transformer in self.transformers:
+            self.pipeline = self.pipeline.map(transformer)
+
+    def select_sink(self, location: SinkInputType):
+        """
+        Select sink element of pipeline.
+        """
+        db_dialects = get_sqlalchemy_dialects()
+        uri = self.sink_address.uri
+
+        if uri.scheme.startswith("callable"):
+            self.sink_element = self.pipeline.stream.sink(location)
+
+        elif uri.scheme in db_dialects:
+            # TODO: Weave in more sophisticated transformations here,
+            #       like topic/topology/storage convergence from Kotori.
+            self.pipeline = self.pipeline.map(operator.attrgetter("payload")).to_dataframe()
+
+            self.sink_element = self.pipeline.stream.dataframe_to_sql(dburi=str(self.sink_address.uri))
+        else:
+            raise InvalidSinkError(f"Invalid sink location: {location}. Scheme unknown: {uri.scheme}.")
+
+    def transform(self, thing: t.Callable) -> "ChannelFactory":
+        self.transformers.append(thing)
+        return self
+
+    def channel(self) -> Channel:
+        """
+        Produce `Channel` object from elements.
+        """
+        return Channel(source=self.source_element, pipeline=self.pipeline, sink=self.sink_element)
+
+
+class Engine:
+    def __init__(self):
+        self.channels: t.Dict[str, Channel] = {}
+        self._terminate_event = asyncio.Event()
+
+    def register(self, channel: Channel):
+        logger.info(f"Registering channel: {channel}")
+        # FIXME: The registry is currently a bit naive.
+        name = str(channel.source)
+        self.channels[name] = channel
+
+    def start(self):
+        for name, channel in self.channels.items():
+            logger.info(f"Starting channel {name}")
+            channel.source.start()
+
+    def stop(self):
+        for channel in self.channels.values():
+            logger.info(f"Stopping channel {channel}")
+            channel.source.stop()
+
+    def terminate(self):
+        self._terminate_event.set()
+
+    async def run_forever(self):
+        # Start engine.
+        self.start()
+        # Wait forever.
+        await self._terminate_event.wait()

lorrystream/exceptions.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2013-2023, The Kotori developers and contributors.
+# Distributed under the terms of the LGPLv3 license, see LICENSE.
+
+
+class InvalidSourceError(Exception):
+    pass
+
+
+class InvalidContentTypeError(Exception):
+    pass
+
+
+class InvalidSinkError(Exception):
+    pass

lorrystream/model.py

@@ -0,0 +1,132 @@
+# Copyright (c) 2013-2023, The Kotori developers and contributors.
+# Distributed under the terms of the LGPLv3 license, see LICENSE.
+
+import dataclasses
+import operator
+import typing as t
+from urllib.parse import parse_qs, urlparse
+
+import funcy
+from streamz import Sink, Source
+
+from lorrystream.exceptions import InvalidSinkError
+from lorrystream.streamz.model import URL, BusMessage
+from lorrystream.util.data import asbool, split_list
+
+SinkInputType = t.Union[str, t.Callable, None]
+
+
+@dataclasses.dataclass
+class Channel:
+    source: Source
+    pipeline: t.Any
+    sink: t.Union[t.Callable, Sink]
+
+    def tap(self, callback: t.Callable):
+        self.pipeline.stream.sink(callback)
+
+
+@dataclasses.dataclass
+class Packet:
+    payload: t.Any
+    busmsg: t.Optional[BusMessage] = None
+
+    @staticmethod
+    def payloads(data):
+        return list(map(operator.attrgetter("payload"), data))
+
+
+class ConnectionString:
+    """
+    Helper class to support ``IoAccessor.export()``.
+    """
+
+    def __init__(self, url):
+        self.url_raw = url
+        self.url = urlparse(url)
+
+    def get_database(self):
+        # Try to get database name from query parameter.
+        database = self.get_query_param("database")
+
+        # Try to get database name from URL path.
+        if database is None:
+            if self.url.path.startswith("/"):
+                database = self.url.path[1:]
+
+        return database or "dwd"
+
+    def get_table(self):
+        return self.get_query_param("table") or "weather"
+
+    def get_path(self):
+        return self.url.path or self.url.netloc
+
+    def get_query_param(self, name):
+        query = parse_qs(self.url.query)
+        try:
+            return query[name][0]
+        except (KeyError, IndexError):
+            return None
+
+
+@dataclasses.dataclass
+class StreamAddress:
+    uri: URL
+    options: t.Dict[str, str] = dataclasses.field(default_factory=dict)
+
+    @classmethod
+    def from_url(cls, url: str):
+        uri = URL(url)
+        options = cls.extract_options(uri)
+        return cls(uri=uri, options=options)
+
+    @staticmethod
+    def extract_options(uri: URL) -> t.Dict[str, str]:
+        # Separate URI query parameters used by LorryStream.
+        control_option_names = [
+            # General options.
+            "content-type",
+            "reconnect",
+            # AMQP options.
+            "exchange",
+            "exchange-type",
+            "queue",
+            "routing-key",
+            "setup",
+        ]
+        list_options = [
+            "setup",
+        ]
+        boolean_options = [
+            "reconnect",
+        ]
+        options = funcy.project(uri.query_params, control_option_names)
+        query_params = funcy.omit(uri.query_params, control_option_names)
+        uri.query_params = query_params
+
+        for list_option in list_options:
+            if list_option in options:
+                options[list_option] = split_list(options[list_option])
+
+        for boolean_option in boolean_options:
+            if boolean_option in options:
+                options[boolean_option] = asbool(options[boolean_option])
+
+        return options
+
+    @classmethod
+    def resolve_sink(cls, sink: SinkInputType):
+        """
+        Select sink element of pipeline.
+        """
+
+        if callable(sink):
+            url = f"callable://{sink.__name__}"
+        elif isinstance(sink, str):
+            url = sink
+        else:
+            raise InvalidSinkError(f"Invalid sink: {sink}")
+
+        uri = URL(url)
+        return cls(uri=uri)