ergon-framework-python 0.1.0__py3-none-any.whl

ergon/connector/pipefy/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

ergon/connector/postgres/__init__.py

@@ -0,0 +1,11 @@
+from .async_connector import AsyncPostgresConnector
+from .async_service import AsyncPostgresService
+from .models import PostgresClient, PostgresConsumerConfig, PostgresProducerConfig
+
+__all__ = [
+    "AsyncPostgresConnector",
+    "AsyncPostgresService",
+    "PostgresClient",
+    "PostgresConsumerConfig",
+    "PostgresProducerConfig",
+]

ergon/connector/postgres/async_connector.py

@@ -0,0 +1,119 @@
+import asyncio
+import logging
+from typing import Any, List, Optional
+
+from ..connector import AsyncConnector
+from ..transaction import Transaction
+from .async_service import AsyncPostgresService
+from .models import PostgresClient, PostgresConsumerConfig, PostgresProducerConfig
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncPostgresConnector(AsyncConnector):
+    service: AsyncPostgresService
+
+    def __init__(
+        self,
+        client: PostgresClient,
+        consumer_config: Optional[PostgresConsumerConfig] = None,
+        producer_config: Optional[PostgresProducerConfig] = None,
+    ) -> None:
+        self.service = AsyncPostgresService(client)
+        self._consumer_config = consumer_config
+        self._producer_config = producer_config
+        self._notify_event: Optional[asyncio.Event] = None
+
+    async def init_async(self) -> None:
+        """
+        Called by the task runner on startup.
+
+        Sets up PG LISTEN if a listen_channel is configured, using an
+        asyncio.Event to allow the consume loop to react to notifications.
+        """
+        if self._consumer_config and self._consumer_config.listen_channel:
+            self._notify_event = asyncio.Event()
+
+            def _on_notify(connection, pid, channel, payload):
+                if self._notify_event is not None:
+                    self._notify_event.set()
+
+            await self.service.listen(self._consumer_config.listen_channel, _on_notify)
+
+    async def fetch_transactions_async(
+        self,
+        batch_size: Optional[int] = None,
+        fetch_query: Optional[str] = None,
+        fetch_params: Optional[List[Any]] = None,
+        *args,
+        **kwargs,
+    ) -> List[Transaction]:
+        if self._consumer_config is None:
+            raise ValueError("AsyncPostgresConnector requires a consumer_config to fetch transactions")
+
+        config = self._consumer_config
+        query = fetch_query or config.fetch_query
+        params = fetch_params or config.fetch_params
+        limit = batch_size or config.batch_size
+
+        rows = await self.service.fetch(query, params, limit=limit)
+
+        if not rows:
+            return []
+
+        if self._notify_event is not None:
+            self._notify_event.clear()
+
+        return [
+            Transaction(
+                id=str(row.get(config.id_column, "")),
+                payload=row,
+                metadata={},
+            )
+            for row in rows
+        ]
+
+    async def dispatch_transactions_async(
+        self,
+        transactions: List[Transaction],
+        dispatch_query: Optional[str] = None,
+        *args,
+        **kwargs,
+    ) -> None:
+        if self._producer_config is None:
+            raise ValueError("AsyncPostgresConnector requires a producer_config to dispatch transactions")
+
+        query = dispatch_query or self._producer_config.dispatch_query
+
+        for txn in transactions:
+            params = self._build_dispatch_params(txn)
+            await self.service.execute(query, params)
+
+    @staticmethod
+    def _build_dispatch_params(transaction: Transaction) -> List[Any]:
+        """
+        Build parameter list for the dispatch query from the transaction.
+
+        Passes transaction.id as the first param. Subclasses or callers
+        can override this for more complex parameter derivation.
+        """
+        return [transaction.id]
+
+    async def wait_for_notify(self, timeout: Optional[float] = None) -> bool:
+        """
+        Wait for a PG NOTIFY event. Returns True if notified, False on timeout.
+
+        Useful for the platform layer to build outbox-relay-style tasks
+        that sleep until the DB signals new rows.
+        """
+        if self._notify_event is None:
+            return False
+
+        try:
+            await asyncio.wait_for(self._notify_event.wait(), timeout=timeout)
+            return True
+        except asyncio.TimeoutError:
+            return False
+
+    async def close(self) -> None:
+        await self.service.close()

ergon/connector/postgres/async_service.py

@@ -0,0 +1,116 @@
+import logging
+from typing import Any, Callable, Dict, List, Optional
+
+import asyncpg
+
+from .models import PostgresClient
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncPostgresService:
+    def __init__(self, client: PostgresClient) -> None:
+        self.client = client
+
+        self._pool: Optional[asyncpg.Pool] = None
+        self._listener_conn: Optional[asyncpg.Connection] = None
+        self._listen_channel: Optional[str] = None
+
+    # ---------- Connection Pool ----------
+
+    async def _get_pool(self) -> asyncpg.Pool:
+        if self._pool is None or self._pool._closed:  # type: ignore[attr-defined]
+            dsn = self.client.get_dsn()
+            ssl_mode = "require" if self.client.ssl else None
+
+            self._pool = await asyncpg.create_pool(
+                dsn,
+                min_size=self.client.min_pool_size,
+                max_size=self.client.max_pool_size,
+                ssl=ssl_mode,
+            )
+            logger.info(
+                "PostgreSQL connection pool created (min=%d, max=%d)",
+                self.client.min_pool_size,
+                self.client.max_pool_size,
+            )
+
+        return self._pool
+
+    # ---------- Query ----------
+
+    async def fetch(
+        self,
+        query: str,
+        params: Optional[List[Any]] = None,
+        limit: Optional[int] = None,
+    ) -> List[Dict[str, Any]]:
+        """Execute a SELECT query and return rows as list of dicts."""
+        pool = await self._get_pool()
+
+        effective_query = query
+        if limit is not None and "LIMIT" not in query.upper():
+            effective_query = f"{query} LIMIT {limit}"
+
+        async with pool.acquire() as conn:
+            rows = await conn.fetch(effective_query, *(params or []))
+
+        return [dict(row) for row in rows]
+
+    async def execute(self, query: str, params: Optional[List[Any]] = None) -> str:
+        """Execute a single statement (INSERT, UPDATE, DELETE). Returns status string."""
+        pool = await self._get_pool()
+        async with pool.acquire() as conn:
+            result = await conn.execute(query, *(params or []))
+        return result
+
+    async def execute_many(self, query: str, params_list: List[List[Any]]) -> None:
+        """Execute a statement for each set of params in the list."""
+        pool = await self._get_pool()
+        async with pool.acquire() as conn:
+            await conn.executemany(query, params_list)
+
+    # ---------- LISTEN / NOTIFY ----------
+
+    async def listen(self, channel: str, callback: Callable[..., Any]) -> None:
+        """
+        Open a dedicated connection and subscribe to a PG NOTIFY channel.
+
+        The callback receives (connection, pid, channel, payload).
+        """
+        if self._listener_conn is not None:
+            logger.warning("Listener already active on channel=%s, closing before re-listen", self._listen_channel)
+            await self.unlisten()
+
+        dsn = self.client.get_dsn()
+        ssl_mode = "require" if self.client.ssl else None
+        conn = await asyncpg.connect(dsn, ssl=ssl_mode)
+        self._listener_conn = conn
+        self._listen_channel = channel
+
+        await conn.add_listener(channel, callback)
+        logger.info("Listening on PG channel=%s", channel)
+
+    async def unlisten(self) -> None:
+        """Remove listener and close the dedicated listener connection."""
+        if self._listener_conn is not None:
+            if self._listen_channel:
+                try:
+                    await self._listener_conn.remove_listener(self._listen_channel, lambda *a: None)
+                except Exception:
+                    pass
+            await self._listener_conn.close()
+            self._listener_conn = None
+            self._listen_channel = None
+            logger.info("PostgreSQL listener connection closed")
+
+    # ---------- Lifecycle ----------
+
+    async def close(self) -> None:
+        """Close pool and listener connection."""
+        await self.unlisten()
+
+        if self._pool is not None:
+            await self._pool.close()
+            self._pool = None
+            logger.info("PostgreSQL connection pool closed")

ergon/connector/postgres/models.py

@@ -0,0 +1,34 @@
+from typing import Any, List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class PostgresClient(BaseModel):
+    dsn: Optional[str] = Field(default=None, description="Full DSN (takes precedence over individual params)")
+    host: str = Field(default="localhost", description="PostgreSQL host")
+    port: int = Field(default=5432, description="PostgreSQL port")
+    user: str = Field(default="postgres", description="PostgreSQL user")
+    password: str = Field(default="postgres", description="PostgreSQL password")
+    database: str = Field(default="postgres", description="PostgreSQL database name")
+    min_pool_size: int = Field(default=1, ge=1, description="Minimum connection pool size")
+    max_pool_size: int = Field(default=10, ge=1, description="Maximum connection pool size")
+    ssl: bool = Field(default=False, description="Enable SSL/TLS connections")
+
+    def get_dsn(self) -> str:
+        if self.dsn:
+            return self.dsn
+        return f"postgresql://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}"
+
+
+class PostgresConsumerConfig(BaseModel):
+    fetch_query: str = Field(description="SQL query to fetch rows as transactions")
+    fetch_params: List[Any] = Field(default_factory=list, description="Bind parameters for the fetch query")
+    batch_size: int = Field(default=100, ge=1, description="Max rows to fetch per call")
+    listen_channel: Optional[str] = Field(
+        default=None, description="Optional PG LISTEN channel for push-based notifications"
+    )
+    id_column: str = Field(default="id", description="Column name to use as Transaction.id")
+
+
+class PostgresProducerConfig(BaseModel):
+    dispatch_query: str = Field(description="SQL statement to execute per dispatched transaction")

ergon/connector/rabbitmq/__init__.py

@@ -0,0 +1,25 @@
+from .async_connector import AsyncRabbitMQConnector
+from .async_service import AsyncRabbitMQService
+from .connector import RabbitMQConnector
+from .models import (
+    AsyncRabbitmqClient,
+    AsyncRabbitmqConsumerConfig,
+    AsyncRabbitmqProducerConfig,
+    RabbitmqClient,
+    RabbitmqConsumerMessage,
+    RabbitmqProducerMessage,
+)
+from .service import RabbitMQService
+
+__all__ = [
+    "AsyncRabbitMQConnector",
+    "AsyncRabbitMQService",
+    "AsyncRabbitmqClient",
+    "AsyncRabbitmqConsumerConfig",
+    "AsyncRabbitmqProducerConfig",
+    "RabbitMQConnector",
+    "RabbitMQService",
+    "RabbitmqClient",
+    "RabbitmqConsumerMessage",
+    "RabbitmqProducerMessage",
+]

ergon/connector/rabbitmq/async_connector.py

@@ -0,0 +1,120 @@
+import json
+import logging
+from typing import Any, Dict, List, Optional
+
+from ..connector import AsyncConnector
+from ..transaction import Transaction
+from .async_service import AsyncRabbitMQService
+from .models import AsyncRabbitmqClient, AsyncRabbitmqConsumerConfig, AsyncRabbitmqProducerConfig
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncRabbitMQConnector(AsyncConnector):
+    service: AsyncRabbitMQService
+
+    def __init__(
+        self,
+        client: AsyncRabbitmqClient,
+        consumer_config: Optional[AsyncRabbitmqConsumerConfig] = None,
+        producer_config: Optional[AsyncRabbitmqProducerConfig] = None,
+    ) -> None:
+        self.service = AsyncRabbitMQService(client)
+        self._consumer_config = consumer_config
+        self._producer_config = producer_config or AsyncRabbitmqProducerConfig()
+
+    async def fetch_transactions_async(
+        self,
+        batch_size: int = 1,
+        queue_name: Optional[str] = None,
+        exchange_name: Optional[str] = None,
+        binding_keys: Optional[List[str]] = None,
+        *args,
+        **kwargs,
+    ) -> List[Transaction]:
+        if self._consumer_config is None:
+            raise ValueError("AsyncRabbitMQConnector requires a consumer_config to fetch transactions")
+
+        config = self._consumer_config
+        if queue_name or exchange_name or binding_keys:
+            overrides: Dict[str, Any] = {}
+            if queue_name is not None:
+                overrides["queue_name"] = queue_name
+            if exchange_name is not None:
+                overrides["exchange_name"] = exchange_name
+            if binding_keys is not None:
+                overrides["binding_keys"] = binding_keys
+            config = config.model_copy(update=overrides)
+
+        messages = await self.service.consume(config, batch_size)
+
+        if not messages:
+            return []
+
+        transactions = []
+        for msg in messages:
+            delivery_tag = msg.get("delivery_tag")
+
+            transactions.append(
+                Transaction(
+                    id=str(delivery_tag),
+                    payload=msg.get("body"),
+                    metadata={
+                        "delivery_tag": delivery_tag,
+                        "routing_key": msg.get("routing_key", ""),
+                        "headers": msg.get("headers", {}),
+                        "content_type": msg.get("content_type"),
+                        "message_id": msg.get("message_id"),
+                        "correlation_id": msg.get("correlation_id"),
+                        "queue_name": config.queue_name,
+                        "_message": msg.get("_message"),
+                    },
+                )
+            )
+
+        return transactions
+
+    async def dispatch_transactions_async(
+        self,
+        transactions: List[Transaction],
+        routing_key: Optional[str] = None,
+        exchange_name: Optional[str] = None,
+        headers: Optional[Dict[str, Any]] = None,
+        *args,
+        **kwargs,
+    ) -> None:
+        config = self._producer_config
+        if exchange_name:
+            config = config.model_copy(update={"exchange_name": exchange_name})
+
+        for txn in transactions:
+            body = txn.payload
+            if not isinstance(body, (str, bytes)):
+                body = json.dumps(body, default=str)
+            if isinstance(body, str):
+                body = body.encode("utf-8")
+
+            rk = routing_key or (txn.metadata.get("routing_key") if txn.metadata else None) or config.routing_key
+            msg_headers = headers or (txn.metadata.get("headers") if txn.metadata else None)
+
+            await self.service.publish(
+                config=config,
+                body=body,
+                routing_key=rk,
+                headers=msg_headers,
+            )
+
+    async def ack_transaction(self, transaction: Transaction) -> None:
+        raw_message = transaction.metadata.get("_message") if transaction.metadata else None
+        if raw_message is None:
+            raise ValueError(f"Cannot ack transaction {transaction.id}: no raw message in metadata")
+        await self.service.ack(raw_message)
+
+    async def nack_transaction(self, transaction: Transaction, requeue: bool = True) -> None:
+        raw_message = transaction.metadata.get("_message") if transaction.metadata else None
+        if raw_message is None:
+            raise ValueError(f"Cannot nack transaction {transaction.id}: no raw message in metadata")
+        await self.service.nack(raw_message, requeue=requeue)
+
+    async def close(self) -> None:
+        await self.service.close()