Skip to content
This repository was archived by the owner on Sep 4, 2025. It is now read-only.
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 18 additions & 2 deletions .github/workflows/test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
5 changes: 4 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
./*/*/*.txt
.vscode/
*pb.go
*pb.go
*pb2*.py
__pycache__/
.pytest_cache/
7 changes: 7 additions & 0 deletions cmd/pyiceberg_receiver/.pyiceberg.yaml
Original file line number Diff line number Diff line change
@@ -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
47 changes: 47 additions & 0 deletions cmd/pyiceberg_receiver/README.md
Original file line number Diff line number Diff line change
@@ -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.
41 changes: 41 additions & 0 deletions cmd/pyiceberg_receiver/__main__.py
Original file line number Diff line number Diff line change
@@ -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)
69 changes: 69 additions & 0 deletions cmd/pyiceberg_receiver/iceberg_receiver.py
Original file line number Diff line number Diff line change
@@ -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.")
45 changes: 45 additions & 0 deletions cmd/pyiceberg_receiver/iceberg_receiver_test.py
Original file line number Diff line number Diff line change
@@ -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"
42 changes: 42 additions & 0 deletions cmd/pyiceberg_receiver/requirements.txt
Original file line number Diff line number Diff line change
@@ -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