Fix RabbitMQ connector: TLS support, rstream API update, AMQP message framing

- Add TLS parameters to DataStorage for RabbitMQ connections
- Update rstream Producer API (Producer() + start() instead of Producer.create())
- Send AMQP 1.0 framed messages from producer for compatibility with Rust consumer
- Add PYTHONUNBUFFERED to producer Dockerfile
- Remap postgres port to avoid conflicts

Author: danghoangnhan

CHANGELOG.md

@@ -6,9 +6,13 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 ## [Unreleased]
 
 ### Added
+- `pw.io.rabbitmq` connector for reading from and writing to RabbitMQ Streams. Supports JSON, plaintext, and raw formats with TLS/mutual TLS configuration. Includes `pw.io.rabbitmq.simple_read` for quick-start usage.
 - `pw.io.postgres.read` connector, which reads data from a PostgreSQL table directly by parsing the Write-Ahead Log (WAL).
 - `pw.io.postgres.write` and `pw.io.postgres.read` now support serialization/deserialization of `np.ndarray` (`int`/`float` elements), homogeneous `tuple` and `list` (via Postgres `ARRAY`; multidimensional rectangular arrays supported).
 
+### Removed
+- `pw.io.mssql` connector has been removed.
+
 ### Changed
 - **BREAKING**: The dependencies for `pw.io.pyfilesystem.read` are no longer included in the default package installation. To install them, please use `pip install pathway[pyfilesystem]`.
 - Asynchronous callback for `pw.io.python.write` is now available as `pw.io.OnChangeCallbackAsync`.

examples/projects/rabbitmq-ETL/docker-compose.yml

@@ -27,7 +27,7 @@ services:
       POSTGRES_PASSWORD: pathway
       POSTGRES_DB: etl_db
     ports:
-      - 5432:5432
+      - 5433:5432
     volumes:
       - ./sql/init-db.sql:/docker-entrypoint-initdb.d/init-db.sql
     healthcheck:

examples/projects/rabbitmq-ETL/pathway-src/etl.py

@@ -17,13 +17,13 @@
 # To use advanced features with Pathway Scale, get your free license key from
 # https://pathway.com/features and paste it below.
 # To use Pathway Community, comment out the line below.
-pw.set_license_key("demo-license-key-with-telemetry")
+# pw.set_license_key("demo-license-key-with-telemetry")
 
-RABBITMQ_URI = "rabbitmq-stream://guest:guest@rabbitmq:5552"
+RABBITMQ_URI = "rabbitmq-stream://guest:guest@0.0.0.0:5552"
 
 postgres_settings = {
-    "host": "postgres",
-    "port": "5432",
+    "host": "0.0.0.0",
+    "port": "5433",
     "dbname": "etl_db",
     "user": "pathway",
     "password": "pathway",

examples/projects/rabbitmq-ETL/producer-src/Dockerfile

@@ -1,5 +1,6 @@
 FROM python:3.10
 
+ENV PYTHONUNBUFFERED=1
 RUN pip install rstream
 COPY ./producer-src/create-streams.py create-streams.py
 

examples/projects/rabbitmq-ETL/producer-src/create-streams.py

@@ -170,15 +170,16 @@
 
 
 async def main():
-    from rstream import Producer
+    from rstream import AMQPMessage, Producer
 
     print("Connecting to RabbitMQ Streams...")
-    producer = await Producer.create(
+    producer = Producer(
         host=RABBITMQ_HOST,
         port=RABBITMQ_PORT,
         username=RABBITMQ_USER,
         password=RABBITMQ_PASSWORD,
     )
+    await producer.start()
 
     # Create all streams
     for stream in STREAMS:
@@ -191,21 +192,21 @@ async def main():
     # Publish initial employees
     print("Publishing employee records...")
     for emp in EMPLOYEES:
-        await producer.send("employees_raw", json.dumps(emp).encode())
+        await producer.send("employees_raw", AMQPMessage(body=json.dumps(emp).encode()))
         await asyncio.sleep(0.3)
     print(f"Published {len(EMPLOYEES)} employees")
 
     # Publish initial orders
     print("Publishing order records...")
     for order in ORDERS:
-        await producer.send("orders_raw", json.dumps(order).encode())
+        await producer.send("orders_raw", AMQPMessage(body=json.dumps(order).encode()))
         await asyncio.sleep(0.2)
     print(f"Published {len(ORDERS)} orders")
 
     # Publish initial listings
     print("Publishing listing records...")
     for listing in LISTINGS:
-        await producer.send("listings_raw", json.dumps(listing).encode())
+        await producer.send("listings_raw", AMQPMessage(body=json.dumps(listing).encode()))
         await asyncio.sleep(0.2)
     print(f"Published {len(LISTINGS)} listings")
 
@@ -229,7 +230,7 @@ async def main():
             "product": random.choice(PRODUCTS),
             "quantity": random.randint(1, 10),
         }
-        await producer.send("orders_raw", json.dumps(new_order).encode())
+        await producer.send("orders_raw", AMQPMessage(body=json.dumps(new_order).encode()))
         print(f"Published new order: {order_id}")
 
         await asyncio.sleep(3)
@@ -245,7 +246,7 @@ async def main():
             "posted_date": "2024-12-10",
             "agent_employee_id": random.choice(agent_ids),
         }
-        await producer.send("listings_raw", json.dumps(new_listing).encode())
+        await producer.send("listings_raw", AMQPMessage(body=json.dumps(new_listing).encode()))
         print(f"Published new listing: {listing_id}")
 
 

integration_tests/db_connectors/conftest.py

@@ -3,7 +3,6 @@
     DebeziumContext,
     DynamoDBContext,
     MongoDBContext,
-    MssqlContext,
     MySQLContext,
     PgvectorContext,
     PostgresContext,
@@ -50,8 +49,3 @@ def dynamodb():
 @pytest.fixture
 def mysql():
     return MySQLContext()
-
-
-@pytest.fixture
-def mssql():
-    return MssqlContext()

integration_tests/db_connectors/utils.py

@@ -80,19 +80,6 @@
     + f"@{MYSQL_DB_HOST}:{MYSQL_DB_PORT}/{MYSQL_DB_NAME}"
 )
 
-MSSQL_DB_HOST = "mssql"
-MSSQL_DB_PORT = 1433
-MSSQL_DB_NAME = "testdb"
-MSSQL_DB_USER = "sa"
-MSSQL_DB_PASSWORD = "YourStrong!Passw0rd"
-MSSQL_CONNECTION_STRING = (
-    f"Server=tcp:{MSSQL_DB_HOST},{MSSQL_DB_PORT};"
-    f"Database={MSSQL_DB_NAME};"
-    f"User Id={MSSQL_DB_USER};"
-    f"Password={MSSQL_DB_PASSWORD};"
-    "TrustServerCertificate=true"
-)
-
 
 def is_mysql_reachable():
     try:
@@ -110,22 +97,6 @@ def is_mysql_reachable():
     return True
 
 
-def is_mssql_reachable():
-    try:
-        import pymssql
-
-        pymssql.connect(
-            server=MSSQL_DB_HOST,
-            port=MSSQL_DB_PORT,
-            user=MSSQL_DB_USER,
-            password=MSSQL_DB_PASSWORD,
-            database=MSSQL_DB_NAME,
-        )
-    except Exception:
-        return False
-
-    return True
-
 
 @dataclass(frozen=True)
 class ColumnProperties:
@@ -550,109 +521,6 @@ def random_table_name(self) -> str:
         return f"mysql_{uuid.uuid4().hex}"
 
 
-class MssqlContext:
-    def __init__(self):
-        import pymssql
-
-        self.connection = pymssql.connect(
-            server=MSSQL_DB_HOST,
-            port=MSSQL_DB_PORT,
-            user=MSSQL_DB_USER,
-            password=MSSQL_DB_PASSWORD,
-            database=MSSQL_DB_NAME,
-            autocommit=True,
-        )
-        self.cursor = self.connection.cursor()
-
-    def random_table_name(self) -> str:
-        return f"mssql_{uuid.uuid4().hex}"
-
-    def execute_sql(self, query: str):
-        self.cursor.execute(query)
-
-    def insert_row(
-        self, table_name: str, values: dict[str, Union[int, bool, str, float]]
-    ) -> None:
-        field_names = list(values.keys())
-        placeholders = ", ".join(["%s"] * len(values))
-        query = f"INSERT INTO {table_name} ({','.join(field_names)}) VALUES ({placeholders})"
-        print(f"Inserting a row: {query}")
-        self.cursor.execute(query, tuple(values.values()))
-
-    def create_table(self, schema: type[pw.Schema], *, add_special_fields: bool) -> str:
-        table_name = self.random_table_name()
-
-        primary_key_found = False
-        fields = []
-        for field_name, field_schema in schema.columns().items():
-            parts = [f"[{field_name}]"]
-            field_type = field_schema.dtype
-            if field_type == dtype.STR:
-                parts.append("NVARCHAR(MAX)")
-            elif field_type == dtype.INT:
-                parts.append("BIGINT")
-            elif field_type == dtype.FLOAT:
-                parts.append("FLOAT")
-            elif field_type == dtype.BOOL:
-                parts.append("BIT")
-            else:
-                raise RuntimeError(f"Unsupported field type {field_type}")
-            if field_schema.primary_key:
-                if primary_key_found:
-                    raise AssertionError("Only single primary key supported")
-                primary_key_found = True
-                parts.append("PRIMARY KEY NOT NULL")
-            fields.append(" ".join(parts))
-
-        if add_special_fields:
-            fields.append("[time] BIGINT NOT NULL")
-            fields.append("[diff] BIGINT NOT NULL")
-
-        create_sql = (
-            f"IF OBJECT_ID(N'{table_name}', N'U') IS NULL "
-            f"CREATE TABLE {table_name} ({','.join(fields)})"
-        )
-        self.cursor.execute(create_sql)
-        return table_name
-
-    def get_table_contents(
-        self,
-        table_name: str,
-        column_names: list[str],
-        sort_by: Union[str, tuple, None] = None,
-    ) -> list[dict[str, Union[str, int, bool, float]]]:
-        select_query = f"SELECT {','.join(column_names)} FROM {table_name};"
-        self.cursor.execute(select_query)
-        rows = self.cursor.fetchall()
-        result = []
-        for row in rows:
-            row_map = dict(zip(column_names, row))
-            result.append(row_map)
-        if sort_by is not None:
-            if isinstance(sort_by, tuple):
-                result.sort(key=lambda item: tuple(item[key] for key in sort_by))
-            else:
-                result.sort(key=lambda item: item[sort_by])
-        return result
-
-    def get_table_schema(self, table_name: str) -> dict[str, ColumnProperties]:
-        query = """
-            SELECT column_name, data_type, is_nullable
-            FROM information_schema.columns
-            WHERE table_name = %s AND table_schema = 'dbo'
-            ORDER BY ordinal_position;
-        """
-        self.cursor.execute(query, (table_name,))
-        rows = self.cursor.fetchall()
-
-        schema_props = {}
-        for column_name, type_name, is_nullable in rows:
-            schema_props[column_name] = ColumnProperties(
-                type_name.lower(), is_nullable.upper() == "YES"
-            )
-        return schema_props
-
-
 class EntryCountChecker:
 
     def __init__(

python/pathway/engine.pyi

@@ -930,6 +930,10 @@ class DataStorage:
         ssl_cert_path: str | None = None,
         psql_replication: PsqlReplicationSettings | None = None,
         schema_name: str | None = None,
+        rabbitmq_tls_root_certificates: str | None = None,
+        rabbitmq_tls_client_cert: str | None = None,
+        rabbitmq_tls_client_key: str | None = None,
+        rabbitmq_tls_trust_certificates: bool = False,
     ) -> None: ...
     def delta_s3_storage_options(self, *args, **kwargs): ...