Merge pull request #101 from 0xgouda/add-pyiceberg-sink

[+] add python based iceberg sink

Author: destrex271

.github/workflows/test.yml

@@ -47,10 +47,26 @@ jobs:
     - name: Setup Protobuf
       uses: ./.github/actions/setup-protobuf
 
-    - name: Test
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.12' 
+
+    - name: Install dependencies
+      run: |
+        python3 -m pip install --upgrade pip
+        pip install pytest
+        pip install -r cmd/pyiceberg_receiver/requirements.txt
+
+    - name: Golang Tests
       run: |
         go generate ./sinks/pb
-        go test -timeout 10m -failfast -v -coverprofile=profile.cov ./...
+        go test -timeout 20m -failfast -v -coverprofile=profile.cov ./...
+
+    - name: Python Tests
+      run: |
+        python3 -m grpc_tools.protoc -I sinks/pb --python_out=cmd/pyiceberg_receiver --grpc_python_out=cmd/pyiceberg_receiver sinks/pb/pgwatch.proto
+        pytest
 
     - name: Coveralls
       uses: coverallsapp/github-action@v2

.gitignore

@@ -1,3 +1,6 @@
 ./*/*/*.txt
 .vscode/
-*pb.go
+*pb.go
+*pb2*.py
+__pycache__/
+.pytest_cache/

cmd/pyiceberg_receiver/.pyiceberg.yaml

@@ -0,0 +1,7 @@
+catalog:
+  pgcatalog:
+    type: sql
+    uri: postgresql://username:password@localhost:5432/database
+    init_catalog_tables: true
+    echo: false
+    pool_pre_ping: false

cmd/pyiceberg_receiver/README.md

@@ -0,0 +1,47 @@
+# Iceberg Receiver
+
+A gRPC server that writes metrics received from pgwatch in Iceberg Table format.
+
+- The server assumes a PostgreSQL catalog is used and creates `pgwatch` namespace and `pgwatch.metrics` table within it if they don't exist.
+- The table is partitioned by `MetricName` and `DBName` (in order).
+- Metrics are written in the local file system as Apache Arrow records with the following schema:
+    ```python
+    Schema(
+        NestedField(field_id=1, name="DBName", field_type=StringType(), required=True),
+        NestedField(field_id=2, name="MetricName", field_type=StringType(), required=True),
+        NestedField(field_id=3, name="Data", field_type=StringType(), required=True),
+    )
+    ```
+- Catalog configurations should be provided in [.pyiceberg.yaml](./.pyiceberg.yaml) file under `pgcatalog` see [PyIceberg SQL Catalog](https://py.iceberg.apache.org/configuration/#sql-catalog) for details.
+
+## Flags
+
+```bash
+usage: pyiceberg_receiver [-h] -p PORT -d DIR
+
+options:
+  -h, --help            show this help message and exit
+  -p PORT, --port PORT  The port number to use for the gRPC server.
+  -d DIR, --iceberg-data-dir DIR
+                        Directory to store iceberg tables in.
+```
+
+## Usage example
+
+```bash
+# generate python gRPC code from protobuf
+python3 -m grpc_tools.protoc -I sinks/pb --python_out=cmd/pyiceberg_receiver --grpc_python_out=cmd/pyiceberg_receiver sinks/pb/pgwatch.proto
+# install dependencies
+pip install -r requirements.txt
+# tell PyIceberg about the dir to look for .pyiceberg.yaml in
+export PYICEBERG_HOME="cmd/pyiceberg_receiver"
+# run the server
+python3 cmd/pyiceberg_receiver -p <grpc-server-port> -d <dir-path>
+```
+
+## TODO
+
+- [ ] Use object storage instead of the local file system.
+- [ ] Support TLS over the gRPC connection.
+- [ ] Add authentication interceptor.
+- [ ] Cache measurements to minimize the number of Parquet files written.

cmd/pyiceberg_receiver/__main__.py

@@ -0,0 +1,41 @@
+import grpc
+import argparse
+from concurrent import futures
+from iceberg_receiver import IcebergReceiver
+from pgwatch_pb2_grpc import add_ReceiverServicer_to_server
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "-p", "--port", 
+    type=int,
+    dest="port", 
+    required=True, 
+    action="store", 
+    help="The port number to use for the gRPC server."
+)
+
+parser.add_argument(
+    "-d", "--iceberg-data-dir",
+    type=str, 
+    dest="icebergDataDir",
+    metavar="DIR",
+    required=True,
+    action="store",
+    help="Directory to store iceberg tables in."
+)
+args = parser.parse_args()
+
+def serve(port: int):
+    """Starts gRPC server listening on port"""
+
+    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
+    add_ReceiverServicer_to_server(
+        IcebergReceiver(args.icebergDataDir),
+        server,
+    )
+    server.add_insecure_port(f"0.0.0.0:{port}")
+    server.start()
+    print(f"gRPC server started, listening on port {port}")
+    server.wait_for_termination()
+
+serve(args.port)

cmd/pyiceberg_receiver/iceberg_receiver.py

@@ -0,0 +1,69 @@
+import pyarrow as pa
+import json
+from pgwatch_pb2_grpc import ReceiverServicer
+from pgwatch_pb2 import Reply
+from google.protobuf import json_format
+from pyiceberg.catalog import load_catalog
+from pyiceberg.schema import Schema
+from pyiceberg.partitioning import PartitionSpec, PartitionField
+from pyiceberg.transforms import IdentityTransform
+from pyiceberg.types import (
+    NestedField,
+    StringType
+)
+
+class IcebergReceiver(ReceiverServicer):
+    def __init__(self, icebergDataDir: str):
+        """
+        Creates pgwatch.metrics table in pgcatalog if it doesn't exist.
+
+        The table is partitioned by DBName and MetricName fields.
+
+        Args:
+            icebergDataDir (str): Local file system dir path to store data at.
+        """
+
+        catalog = load_catalog("pgcatalog")
+        catalog.create_namespace_if_not_exists("pgwatch")
+
+        schema = Schema(
+            NestedField(field_id=1, name="DBName", field_type=StringType(), required=True),
+            NestedField(field_id=2, name="MetricName", field_type=StringType(), required=True),
+            NestedField(field_id=3, name="Data", field_type=StringType(), required=True),
+        )
+
+        partition_spec = PartitionSpec(
+            PartitionField(
+                source_id=2, field_id=1000, transform=IdentityTransform(), name="MetricName"
+            ),
+            PartitionField(
+                source_id=1, field_id=1001, transform=IdentityTransform(), name="DBName"
+            ),
+        )
+
+        tbl = catalog.create_table_if_not_exists(
+            identifier=("pgwatch", "metrics"),
+            schema=schema,
+            location=icebergDataDir,
+            partition_spec=partition_spec
+        )        
+
+        self.catalog = catalog
+        self.tbl = tbl
+        self.arrow_schema = tbl.schema().as_arrow()
+
+
+    def UpdateMeasurements(self, request, context):
+        data = [json_format.MessageToDict(row) for row in request.Data]
+        dataJson = json.dumps(data)
+
+        measurement = [{
+            "DBName": request.DBName,
+            "MetricName": request.MetricName,
+            "Data": dataJson,
+        }]
+
+        df = pa.Table.from_pylist(measurement, schema=self.arrow_schema)
+        self.tbl.append(df)
+
+        return Reply(logmsg="Metrics Inserted in iceberg.")

cmd/pyiceberg_receiver/iceberg_receiver_test.py

@@ -0,0 +1,45 @@
+import yaml
+import os
+import pytest
+from testcontainers.postgres import PostgresContainer
+from pgwatch_pb2 import MeasurementEnvelope, Reply
+
+@pytest.fixture(scope="module", autouse=True)
+def setup_catalog(tmp_path_factory):
+    with PostgresContainer("postgres:16") as postgres:
+        tmp_path = tmp_path_factory.getbasetemp()
+        os.environ["PYICEBERG_HOME"] = tmp_path.as_posix()
+
+        test_iceberg_yaml = {
+            "catalog": {
+                "pgcatalog": {
+                    "uri": postgres.get_connection_url(),
+                    "type": "sql",
+                    "init_catalog_tables": True
+                }
+            }
+        }
+
+        test_pyiceberg_file = tmp_path / ".pyiceberg.yaml"
+        with open(test_pyiceberg_file.as_posix(), 'w') as file:
+            yaml.dump(test_iceberg_yaml, file, default_flow_style=False, allow_unicode=True)
+
+        yield 
+
+
+def test_IcebergReceiver(tmp_path):
+    # Late import to allow `PYICEERG_HOME` env to be
+    # set by `setup_catalog` fixture before 
+    # pyiceberg reads it on init
+    from iceberg_receiver import IcebergReceiver
+
+    recv = IcebergReceiver(tmp_path.as_posix()) 
+    msg = MeasurementEnvelope(DBName="test",MetricName="test")
+    reply = recv.UpdateMeasurements(msg, None)
+    assert reply == Reply(logmsg="Metrics Inserted in iceberg.")
+
+    paTable = recv.tbl.scan().to_arrow()
+    pyList = paTable.to_pylist()
+
+    assert pyList[0]["DBName"] == "test"
+    assert pyList[0]["MetricName"] == "test"

cmd/pyiceberg_receiver/requirements.txt

@@ -0,0 +1,42 @@
+annotated-types==0.7.0
+cachetools==5.5.2
+certifi==2025.8.3
+charset-normalizer==3.4.3
+click==8.2.1
+docker==7.1.0
+fsspec==2025.7.0
+greenlet==3.2.4
+grpcio==1.74.0
+grpcio-tools==1.74.0
+idna==3.10
+iniconfig==2.1.0
+markdown-it-py==4.0.0
+mdurl==0.1.2
+mmh3==5.2.0
+packaging==25.0
+pluggy==1.6.0
+protobuf==6.32.0
+psycopg2-binary==2.9.10
+pyarrow==21.0.0
+pydantic==2.11.7
+pydantic_core==2.33.2
+Pygments==2.19.2
+pyiceberg==0.9.1
+pyparsing==3.2.3
+pytest==8.4.1
+python-dateutil==2.9.0.post0
+python-dotenv==1.1.1
+PyYAML==6.0.2
+requests==2.32.5
+rich==13.9.4
+setuptools==80.9.0
+six==1.17.0
+sortedcontainers==2.4.0
+SQLAlchemy==2.0.43
+strictyaml==1.7.3
+tenacity==9.1.2
+testcontainers==4.12.0
+typing-inspection==0.4.1
+typing_extensions==4.14.1
+urllib3==2.5.0
+wrapt==1.17.3