diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index fc059bd..01555ac 100644 --- a/.github/workflows/test.yml +++ b/.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 diff --git a/.gitignore b/.gitignore index 2d90b05..3c10904 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ ./*/*/*.txt .vscode/ -*pb.go \ No newline at end of file +*pb.go +*pb2*.py +__pycache__/ +.pytest_cache/ \ No newline at end of file diff --git a/cmd/pyiceberg_receiver/.pyiceberg.yaml b/cmd/pyiceberg_receiver/.pyiceberg.yaml new file mode 100644 index 0000000..c1143be --- /dev/null +++ b/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 \ No newline at end of file diff --git a/cmd/pyiceberg_receiver/README.md b/cmd/pyiceberg_receiver/README.md new file mode 100644 index 0000000..dbe2fe0 --- /dev/null +++ b/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 -d +``` + +## 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. \ No newline at end of file diff --git a/cmd/pyiceberg_receiver/__main__.py b/cmd/pyiceberg_receiver/__main__.py new file mode 100644 index 0000000..3329270 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/cmd/pyiceberg_receiver/iceberg_receiver.py b/cmd/pyiceberg_receiver/iceberg_receiver.py new file mode 100644 index 0000000..6834074 --- /dev/null +++ b/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.") \ No newline at end of file diff --git a/cmd/pyiceberg_receiver/iceberg_receiver_test.py b/cmd/pyiceberg_receiver/iceberg_receiver_test.py new file mode 100644 index 0000000..3d9ef44 --- /dev/null +++ b/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" diff --git a/cmd/pyiceberg_receiver/requirements.txt b/cmd/pyiceberg_receiver/requirements.txt new file mode 100644 index 0000000..3e4b484 --- /dev/null +++ b/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