feat: Add support for new BATCH message type in taps and targets (#904

) * Add SDK Batch message type * Fix docstring * Restore _check_max_record_limit call * Set default batch size to 1000 * Implement Stream.get_batch_encoding * Implement Stream.get_batch_config * Implement batch for sinks * Remove custom SQLite insert * Fix type annotation * Refactor Sink.process_batch_file -> Sink.process_batch_files * Address mapper feedback * docs: Add page for BATCH * Make linter happy * Apply docs suggestions from code review * Update docs/batch.md Co-authored-by: Aaron ("AJ") Steers <[email protected]> * Get target batch storage config from absolute manifest URLs * Apply suggestions from code review Co-authored-by: Aaron ("AJ") Steers <[email protected]> Co-authored-by: Aaron ("AJ") Steers <[email protected]>
meltano · Sep 15, 2022 · 5bab2d8 · 5bab2d8
1 parent aab68a1
commit 5bab2d8
Show file tree

Hide file tree

Showing 20 changed files with 914 additions and 79 deletions.
diff --git a/docs/batch.md b/docs/batch.md
@@ -0,0 +1,95 @@
+# Batch Messages
+
+```{warning}
+The `BATCH` message functionality is currently in preview and is subject to change.
+You can [open an issue](https://github.com/meltano/sdk/issues) or [join the discussion](https://github.com/meltano/sdk/discussions/963) on GitHub to provide feedback during the preview period.
+```
+
+[The Singer message specification](https://github.com/singer-io/getting-started/blob/master/docs/SPEC.md#output) defines the three basic types of messages: `RECORD`, `STATE`, and `SCHEMA`. The `RECORD` message is used to send data from the tap to the target. The `STATE` message is used to send state data from the tap to the target. The `SCHEMA` message is used to send schema data from the tap to the target, and for example, create tables with the correct column types.
+
+However, the Singer specification can be extended to support additional types of messages. For example, the [`ACTIVATE_VERSION`](https://sdk.meltano.com/en/latest/capabilities.html#singer_sdk.helpers.capabilities.PluginCapabilities.ACTIVATE_VERSION) message is used to manage hard deletes in the target.
+
+This library's implementation of the `BATCH` message is used to send records in bulk from the tap to the target, using an intermediate filesystem to store _batch_ files. This is useful, for example
+
+- when the tap outputs records at a much higher rate than the target can consume them, creating backpressure
+- when the source system can directly export data in bulk (e.g. a database dump)
+
+Currently only a local filesystem is supported, but other filesystems like AWS S3, FTP, etc. could be supported in the future.
+
+## The `BATCH` Message
+
+```json
+{
+  "type": "BATCH",
+  "stream": "users",
+  "encoding": {
+    "format": "jsonl",
+    "compression": "gzip"
+  },
+  "manifest": [
+    "file://path/to/batch/file/1",
+    "file://path/to/batch/file/2"
+  ]
+}
+```
+
+### `encoding`
+
+The `encoding` field is used to specify the format and compression of the batch files. Currently only `jsonl` and `gzip` are supported, respectively.
+
+### `manifest`
+
+The `manifest` field is used to specify the paths to the batch files. The paths are relative to the `root` directory specified in the [`batch_config`](#batch-configuration) storage configuration.
+
+## Batch configuration
+
+When local storage is used, targets do no require special configuration to process `BATCH` messages.
+
+Taps may be configured to specify a root storage `root` directory, file path `prefix`, and `encoding` for batch files using a configuration like the below:
+
+
+In `config.json`:
+
+```js
+{
+  // ...
+  "batch_config": {
+    "encoding": {
+      "format": "jsonl",
+      "compression": "gzip",
+    },
+    "storage": {
+      "root": "file://tests/core/resources",
+      "prefix": "test-batch-",
+    }
+  }
+}
+```
+
+## Custom batch file creation and processing
+
+### Tap side
+
+Taps can optionally customize the batch file creation by implementing the [`get_batches`](singer_sdk.Stream.get_batches). This method should return a _tuple_ of an encoding and a list of batch files:
+
+```python
+class MyStream(Stream):
+    def get_batches(self, records):
+        return (
+            ParquetEncoding(compression="snappy"),
+            [
+                "s3://my-bucket/my-batch-file-1.parquet",
+                "s3://my-bucket/my-batch-file-2.parquet",
+            ]
+        )
+```
+
+### Target side
+
+Targets can optionally customize the batch file processing by implementing the [`process_batch_files`](singer_sdk.Sink.process_batch_files).
+
+```python
+class MySink(Sink):
+    def process_batch_files(self, encoding, storage, files):
+        # process the batch files
+```
diff --git a/docs/index.rst b/docs/index.rst
@@ -64,6 +64,7 @@ Advanced Topics
    parent_streams
    partitioning
    stream_maps
+   batch
    porting
    sinks
    CONTRIBUTING

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -40,6 +40,7 @@ pipelinewise-singer-python = "1.2.0"
 backoff = ">=1.8.0,<2.0"
 pendulum = "^2.1.0"
 click = "~=8.0"
+fs = "^2.4.16"
 PyJWT = "~=2.4"
 requests = "^2.25.1"
 cryptography = ">=3.4.6,<39.0.0"
@@ -148,6 +149,7 @@ exclude_lines = [
     "if __name__ == .__main__.:",
     '''class .*\bProtocol\):''',
     '''@(abc\.)?abstractmethod''',
+    "if TYPE_CHECKING:",
 ]
 fail_under = 82
 

diff --git a/samples/sample_tap_countries/countries_tap.py b/samples/sample_tap_countries/countries_tap.py
@@ -28,3 +28,7 @@ def discover_streams(self) -> List[Stream]:
             CountriesStream(tap=self),
             ContinentsStream(tap=self),
         ]
+
+
+if __name__ == "__main__":
+    SampleTapCountries.cli()
diff --git a/samples/sample_target_sqlite/__init__.py b/samples/sample_target_sqlite/__init__.py
@@ -1,8 +1,11 @@
 """A sample implementation for SQLite."""
 
-from typing import Any, Dict
+from __future__ import annotations
+
+from typing import Any
 
 import sqlalchemy
+from sqlalchemy.dialects.sqlite import insert
 
 from singer_sdk import SQLConnector, SQLSink, SQLTarget
 from singer_sdk import typing as th
@@ -20,7 +23,7 @@ class SQLiteConnector(SQLConnector):
     allow_column_alter = False
     allow_merge_upsert = True
 
-    def get_sqlalchemy_url(self, config: Dict[str, Any]) -> str:
+    def get_sqlalchemy_url(self, config: dict[str, Any]) -> str:
         """Generates a SQLAlchemy URL for SQLite."""
         return f"sqlite:///{config[DB_PATH_CONFIG]}"