Merge pull request #102 from Point72/ac/filedrop

Add filedrop gateway module

Author: timkpaine

csp_gateway/server/modules/__init__.py

@@ -1,4 +1,5 @@
 from .controls import *
+from .filedrop import *
 from .initializer import *
 from .io import *
 from .kafka import *

csp_gateway/server/modules/filedrop/__init__.py

@@ -0,0 +1,10 @@
+import sys
+
+if sys.platform.startswith("linux"):
+    try:
+        from .adapter import *
+        from .filedrop import *
+    except ImportError:
+        pass
+else:
+    pass

csp_gateway/server/modules/filedrop/adapter.py

@@ -0,0 +1,223 @@
+import csv
+import logging
+from dataclasses import dataclass
+from datetime import datetime
+from enum import Enum, auto
+from typing import Any, Dict, List, Optional, TypeVar, get_args, get_origin
+
+import orjson
+import pyarrow.parquet as pq
+from csp import ts
+from csp.impl.pushadapter import PushInputAdapter
+from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer
+from csp.impl.wiring import py_push_adapter_def
+from pydantic import TypeAdapter
+from watchdog.events import FileSystemEvent, FileSystemEventHandler
+from watchdog.observers import Observer
+
+__all__ = (
+    "FileDropType",
+    "FileDropAdapterConfiguration",
+    "filedrop_adapter_def",
+)
+
+
+log = logging.getLogger(__name__)
+
+
+class FileDropType(Enum):
+    CSV = auto()
+    JSON = auto()
+    PARQUET = auto()
+
+
+@dataclass
+class FileDropAdapterConfiguration:
+    """Configuration for the filedrop push adapter"""
+
+    # Path to the directory to monitor
+    dir_path: str
+    # Format of files to expect to load properly i.e parquet, json, csv
+    filedrop_type: FileDropType
+    # Map the data fields from the file to the fields of the structs
+    field_map: Dict[str, str]
+    # List of extensions to filter, empty list means all extensions are allowed
+    extensions: List[str]
+    # Extra args to the type adapter deserializer
+    type_adapter_args: Dict[str, Any]
+
+
+class FileReaderBase:
+    """The base file reader that reads data from files and generates structs"""
+
+    def __init__(self, config: FileDropAdapterConfiguration, ts_typ: object, deserializer: Optional[object] = None):
+        self.field_map = config.field_map
+        self.extensions = config.extensions
+        if hasattr(config, "type_adapter_args"):
+            self.context = config.type_adapter_args
+        else:
+            self.context = {}
+        if not deserializer:
+            normalized_type = ContainerTypeNormalizer.normalize_type(ts_typ)
+            type_adapter = TypeAdapter(normalized_type)
+            if get_origin(normalized_type) is list:
+
+                def deserialize_tick(data, type_adapter=type_adapter, apply_field_map=self.apply_field_map, context=self.context):
+                    data = [apply_field_map(d) for d in data]
+                    return type_adapter.validate_python(data, context=context)
+            elif get_origin(normalized_type) is dict:
+                key_type, inner_type = get_args(normalized_type)
+
+                def deserialize_tick(data, type_adapter=type_adapter, apply_field_map=self.apply_field_map, context=self.context):
+                    data = {k: apply_field_map(d) for k, d in data.items()}
+                    return type_adapter.validate_python(data, context=context)
+            else:
+
+                def deserialize_tick(data, type_adapter=type_adapter, apply_field_map=self.apply_field_map, context=self.context):
+                    return type_adapter.validate_python(apply_field_map(data), context=context)
+
+            self.deserializer = deserialize_tick
+        else:
+            self.deserializer = deserializer
+
+    def read(self, src_path: str) -> object:
+        """Generator to return stucts from a filepath"""
+
+        should_read = True
+        if self.extensions:
+            if not any([src_path.endswith(suffix) for suffix in self.extensions]):
+                should_read = False
+        if should_read:
+            dicts = self.read_impl(src_path)
+            structs = [self.deserializer(d) for d in dicts]
+            for s in structs:
+                yield s
+
+    def read_impl(self, src_path: str) -> List[dict]:
+        """File type specific implementation"""
+
+        raise Exception(f"read not implemented for {self}")
+
+    def apply_field_map(self, data: dict) -> dict:
+        """Convert the keys in the data to the field names of the struct"""
+
+        if self.field_map:
+            new_data = {}
+            for k, v in data.items():
+                new_k = self.field_map.get(k, k)
+                new_data[new_k] = v
+            return new_data
+        else:
+            return data
+
+
+class FileReaderCsv(FileReaderBase):
+    """File reader for json file type"""
+
+    def read_impl(self, src_path: str) -> List[dict]:
+        data = []
+        with open(src_path, "r") as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                data.append(row)
+        return data
+
+
+class FileReaderJson(FileReaderBase):
+    """File reader for json file type"""
+
+    def read_impl(self, src_path: str) -> List[dict]:
+        with open(src_path, "rb") as f:
+            data = orjson.loads(f.read())
+        if isinstance(data, list):
+            res = data
+        else:
+            res = [data]
+        return res
+
+
+class FileReaderParquet(FileReaderBase):
+    """File reader for parquet file type"""
+
+    def read_impl(self, src_path: str) -> List[dict]:
+        table = pq.read_table(src_path)
+        return table.to_pylist()
+
+
+class EventHandlerCustom(FileSystemEventHandler):
+    def __init__(self, adapter: PushInputAdapter, file_reader: FileReaderBase):
+        self.file_reader = file_reader
+        self.adapter = adapter
+        self._created = set()
+        self._opened = set()
+        self._modified = set()
+        self._closed = set()
+
+    def on_created(self, event: FileSystemEvent):
+        self._created.add(event.src_path)
+
+    def on_opened(self, event: FileSystemEvent):
+        if event.src_path in self._created:
+            self._created.remove(event.src_path)
+            self._opened.add(event.src_path)
+
+    def on_modified(self, event: FileSystemEvent):
+        if event.src_path in self._opened:
+            self._opened.remove(event.src_path)
+            self._modified.add(event.src_path)
+
+    def on_closed(self, event: FileSystemEvent):
+        if event.src_path in self._modified:
+            self._modified.remove(event.src_path)
+            file_path = event.src_path
+            try:
+                for data in self.file_reader.read(file_path):
+                    self.adapter.push_tick(data)
+            except Exception as e:
+                log.error(f"Failed to read data from {file_path} with exception: {e}, skipping")
+
+
+T = TypeVar("T")
+
+
+class _FileDropImpl(PushInputAdapter):
+    FILEREADER_MAP = {
+        FileDropType.CSV: FileReaderCsv,
+        FileDropType.JSON: FileReaderJson,
+        FileDropType.PARQUET: FileReaderParquet,
+    }
+
+    def __init__(self, config: FileDropAdapterConfiguration, ts_typ: T, deserializer: Optional[object] = None):
+        # NOTE ts_typ is assumed to be a List["Y"] type where "Y" is the actual type
+        self.dir_path = config.dir_path
+        self.observer = Observer()
+        reader = self.FILEREADER_MAP[config.filedrop_type]
+        file_reader = reader(config, ts_typ[0], deserializer)
+        self.event_handler = EventHandlerCustom(self, file_reader)
+        self.observer.schedule(self.event_handler, self.dir_path, recursive=False)
+
+    def start(self, starttime: datetime, endtime: datetime):
+        self.observer.start()
+        self.observer_started = True
+
+    def stop(self):
+        if self.observer_started:
+            self.observer.stop()
+            self.observer.join()
+
+
+# NOTE: Ref: https://github.com/Point72/csp/issues/569
+# Instead of passing the actual type T we want as ts_typ, we have to pass in List[T]
+# This is due to a type normalization bug in csp, where it converts Dict[K, T] to dict, i.e removing the type annotation hints
+# The type annotations are needed to properly deserialize the raw data using the pydantic type adapters.
+# Further we can only take the object type as output due to issues caused by using List[T] instead of T and bypassing
+# the type normalization process in csp. The user of the adapter, needs to keep track of the type they pass in ts_type,
+# and cast the ts[object] to the required ts[T], they want using csp.appy
+filedrop_adapter_def = py_push_adapter_def(
+    "filedrop_adapter_def",
+    _FileDropImpl,
+    ts[object],
+    config=FileDropAdapterConfiguration,
+    ts_typ=List[T],
+    deserializer=Optional[object],
+)

csp_gateway/server/modules/filedrop/filedrop.py

@@ -0,0 +1,122 @@
+import logging
+import os
+from typing import Dict, List, TypeVar, get_args, get_origin
+
+import csp
+from ccflow import BaseModel
+from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer
+from csp.impl.types.tstype import isTsType
+from pydantic import Field
+
+from csp_gateway.server import (
+    GatewayChannels,
+    GatewayModule,
+)
+
+from .adapter import FileDropAdapterConfiguration, FileDropType, filedrop_adapter_def
+
+__all__ = (
+    "ReadFileDropConfiguration",
+    "ReadFileDrop",
+)
+
+log = logging.getLogger(__name__)
+
+
+class ReadFileDropConfiguration(BaseModel):
+    """The configuration of a filedrop adapter for a directory and filetype"""
+
+    channel_name: str = Field(description="Name of the channel to send the structs to")
+    filedrop_type: FileDropType = Field(description="The type of files to expect and accordingly read i.e. parquet, json, etc")
+    field_map: Dict[str, str] = Field(
+        default={}, description="A map to convert the keys in the data to the field names in the struct type of the channel"
+    )
+    extensions: List[str] = Field(default=[], description="List of extensions to decide which files to read, empty list means all extensions")
+    subscribe_with_struct_id: bool = Field(
+        default=False,
+        description=("If False, replaces the id field on GatewayStructs from files with one autogenerated by the current Gateway."),
+    )
+    subscribe_with_struct_timestamp: bool = Field(
+        default=False,
+        description=("If False, replaces the timestamp field on the GatewayStruct with a timestamp autogenerated by the current Gateway."),
+    )
+
+
+T = TypeVar("T")
+K = TypeVar("K")
+
+
+class ReadFileDrop(GatewayModule):
+    """The module to read data from files dropped in specific directories and send them as structs to specific channels"""
+
+    directory_configs: Dict[str, List[ReadFileDropConfiguration]] = Field(
+        description="Mapping of directories to a list of ReadFileDropConfiguration, for that directory in the filedrop module"
+    )
+
+    @csp.node
+    def handle_list_basket(self, data: csp.ts[List[T]], list_size: int) -> csp.OutputBasket(List[csp.ts[T]], shape="list_size"):
+        if csp.ticked(data):
+            return data
+
+    @csp.node
+    def handle_dict_basket(self, data: csp.ts[Dict[K, T]], dict_keys: List[K]) -> csp.OutputBasket(Dict[K, csp.ts[T]], shape="dict_keys"):
+        if csp.ticked(data):
+            return data
+
+    def connect(self, channels: GatewayChannels):
+        channel_data = {}
+        channel_basket_types = {}
+        channel_base_types = {}
+        for dir, configs in self.directory_configs.items():
+            # ensure that directories exist
+            os.makedirs(dir, exist_ok=True)
+            for config in configs:
+                channel_data[config.channel_name] = []
+        for dir, configs in self.directory_configs.items():
+            for config in configs:
+                context = {
+                    "force_new_id": not config.subscribe_with_struct_id,
+                    "force_new_timestamp": not config.subscribe_with_struct_timestamp,
+                }
+                adapter_config = FileDropAdapterConfiguration(
+                    dir_path=dir,
+                    filedrop_type=config.filedrop_type,
+                    field_map=config.field_map,
+                    extensions=config.extensions,
+                    type_adapter_args=context,
+                )
+                channel_type = channels.get_outer_type(config.channel_name)
+                if isTsType(channel_type):
+                    non_ts_type = channel_type.typ
+                    channel_basket_types[config.channel_name] = ""
+                else:
+                    normalized_type = ContainerTypeNormalizer.normalize_type(channel_type)
+                    if get_origin(normalized_type) is list:
+                        inner_type = get_args(normalized_type)[0]
+                        if not isTsType(inner_type):
+                            raise ValueError(f"Channel type for {config.channel_name} should be of the form List[TsType], got: {channel_type}")
+                        non_ts_type = List[inner_type.typ]
+                        channel_basket_types[config.channel_name] = "list"
+                    elif get_origin(normalized_type) is dict:
+                        key_type, inner_type = get_args(normalized_type)
+                        if not isTsType(inner_type):
+                            raise ValueError(
+                                f"Channel type for {config.channel_name} should be of the form Dict[KeyType, TsType], got: {channel_type}"
+                            )
+                        non_ts_type = Dict[key_type, inner_type.typ]
+                        channel_basket_types[config.channel_name] = "dict"
+                    else:
+                        raise Exception(f"Channel type cannot be handled: {channel_type}")
+                channel_base_types[config.channel_name] = non_ts_type
+                # NOTE: We have to pass in the type as [type] because of the type normalization issue in csp: https://github.com/Point72/csp/issues/569
+                data = filedrop_adapter_def(config=adapter_config, ts_typ=[non_ts_type], deserializer=None)
+                channel_data[config.channel_name].append(data)
+        for channel_name, data_list in channel_data.items():
+            data = csp.flatten(data_list)
+            typed_data = csp.apply(data, lambda x: x, channel_base_types[channel_name])
+            if channel_basket_types[channel_name] == "list":
+                channels.set_channel(channel_name, self.handle_list_basket(typed_data, list_size=channels.dynamic_keys()[channel_name]))
+            elif channel_basket_types[channel_name] == "dict":
+                channels.set_channel(channel_name, self.handle_dict_basket(typed_data, dict_keys=channels.dynamic_keys()[channel_name]))
+            else:
+                channels.set_channel(channel_name, typed_data)