feat(elt-pipelines): Proposal postgresql sources pipeline

elt-common/pyproject.toml

@@ -47,6 +47,13 @@ dev = [
   "ruff>=0.16.1",
 ]
 
+test = [
+    "authlib>=1.7.2",
+    "httpx>=0.28.1",
+    "tenacity>=9.1.2",
+    "sqlalchemy>=2.0.51",
+    "trino>=0.336.0",
+]
 
 [tool.ruff]
 line-length = 100

elt-common/src/elt_common/sources/sqldatabase/__init__.py

@@ -1,38 +1,71 @@
 """Support for ingesting data from an SQL database."""
 
+import json
 import logging
 from abc import abstractmethod
-from typing import Generator, Iterator, NamedTuple, Optional, Callable
+from collections.abc import Callable, Generator, Iterable, Iterator
+from typing import NamedTuple
 
 import pyarrow as pa
+import pyarrow.compute as pc
 import sqlalchemy as sa
-from pydantic import SecretStr, PositiveInt
+from pydantic import PositiveInt, SecretStr
 from pydantic_settings import BaseSettings
 from sqlalchemy import Select
 
-from elt_common.extract import ResourceProperties, ResourceWriteProperties, Watermark, BaseExtract
+from elt_common.extract import BaseExtract, ResourceProperties, ResourceWriteProperties, Watermark
 from elt_common.sources.sqldatabase.schema import to_pyarrow_schema
 
 LOGGER = logging.getLogger(__name__)
 
 
-class SqlDatabaseSourceConfig(BaseSettings):
-    """Configuration required to connect to a database"""
+def _serialize_json_values(row: dict) -> dict:
+    """Ensure dict/list objects in rows are serialized to JSON strings for PyArrow json_ types."""
+    return {k: json.dumps(v) if isinstance(v, (dict, list)) else v for k, v in row.items()}
+
+
+def json_to_str(table: pa.Table) -> pa.Table:
+    """Fast-path helper to cast any PyArrow JSON columns in a Table to String columns."""
+    # PyArrow JSON extension types are instances of pa.JsonType
+    json_field_indices = [
+        i for i, field in enumerate(table.schema) if isinstance(field.type, pa.JsonType)
+    ]
+
+    if not json_field_indices:
+        return table
+
+    new_schema = table.schema
+    new_columns = list(table.columns)
+
+    for i in json_field_indices:
+        field = table.schema.field(i)
+        new_field = field.with_type(pa.string())
+        new_schema = new_schema.set(i, new_field)
+        new_columns[i] = pc.cast(table.column(i), pa.string())
+
+    return pa.Table.from_arrays(new_columns, schema=new_schema)
+
 
-    # connection
+def _partition_to_pyarrow_table(partition: Iterable[dict], schema: pa.Schema) -> pa.Table:
+    """Converts a partition of mapping rows into a PyArrow Table, handling JSON serialization and schema casting."""
+    rows = [_serialize_json_values(row) for row in partition]
+    pa_table = pa.Table.from_pylist(rows, schema=schema)
+    return json_to_str(pa_table)
+
+
+class SqlDatabaseSourceConfig(BaseSettings):
     drivername: str
     database: str
-    database_schema: Optional[str] = None
-    port: Optional[int] = None
-    host: Optional[str] = None
-    username: Optional[str] = None
-    password: Optional[SecretStr] = None
+    database_schema: str | None = None
+    port: int | None = None
+    host: str | None = None
+    username: str | None = None
+    password: SecretStr | None = None
 
-    # loading behaviour
     chunk_size: int = 5000
     """If the query returns more than chunk_size rows, fetch them in multiple chunks of at most this size"""
 
-    row_limit: Optional[PositiveInt] = None
+    row_limit: PositiveInt | None = None
     """Maximum number of rows to return from each table, primarily for testing purposes. No limit if 'None'"""
 
     @property
@@ -61,9 +94,9 @@ class TableInfo(NamedTuple):
     be different to the name of the DB table
     """
 
-    write_properties: Optional[ResourceWriteProperties] = None
-    watermark_column: Optional[str] = None
-    destination_table_name: Optional[str] = None
+    write_properties: ResourceWriteProperties | None = None
+    watermark_column: str | None = None
+    destination_table_name: str | None = None
 
 
 class SqlDatabaseExtract(BaseExtract[SqlDatabaseSourceConfig]):
@@ -99,7 +132,7 @@ def __init__(self, config: SqlDatabaseSourceConfig):
         self._metadata = sa.MetaData(schema=config.database_schema)
 
     @abstractmethod
-    def table_info(self) -> dict[str, Optional[TableInfo]]:
+    def table_info(self) -> dict[str, TableInfo | None]:
         """Define the tables to be extracted from the DB.
 
         Each key in the returned dict is a table name. Their values can include
@@ -111,7 +144,6 @@ def table_info(self) -> dict[str, Optional[TableInfo]]:
         (e.g. filtering) extend :py:meth:`extract_resource_properties` with
         custom extractors.
         """
-        pass
 
     def extract_resource_properties(self):
         """Open a connection to the DB and return ingest properties for tables
@@ -142,6 +174,11 @@ def _make_table_properties(
                 if table_props and table_props.watermark_column
                 else None
             )
+            resource_name = (
+                table_props.destination_table_name
+                if table_props and table_props.destination_table_name
+                else name
+            )
 
             def extractor(watermark, *, _name=name):
                 return self._extract_table(_name, watermark=watermark, conn=conn)
@@ -152,11 +189,7 @@ def extractor(watermark, *, _name=name):
                 watermark_column=watermark_column,
             )
 
-            destination_table = name
-            if table_props is not None and table_props.destination_table_name is not None:
-                destination_table = table_props.destination_table_name
-
-            yield destination_table, properties
+            yield resource_name, properties
 
     def _extract_table(
         self,
@@ -188,6 +221,12 @@ def _extract_table(
         # the table
         pa_schema = to_pyarrow_schema(table)
         result = conn.execution_options(yield_per=self.config.chunk_size).execute(query)
+
+        has_data = False
         for partition in result.mappings().partitions():
-            table = pa.Table.from_pylist(partition, schema=pa_schema)
-            yield table
+            has_data = True
+            yield _partition_to_pyarrow_table(partition, schema=pa_schema)
+
+        if not has_data:
+            empty_table = pa.Table.from_batches([], schema=pa_schema)
+            yield empty_table

elt-common/src/elt_common/sources/sqldatabase/schema.py

@@ -5,6 +5,7 @@
 
 import pyarrow as pa
 import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
 
 
 def to_pyarrow_schema(table: sa.Table) -> pa.Schema:
@@ -25,6 +26,7 @@ def _to_pyarrow_field(column: sa.Column) -> pa.Field:
     sa.Float: pa.float64,
     sa.Integer: pa.int32,
     sa.Interval: lambda: pa.duration("us"),
+    sa.JSON: pa.json_,
     sa.LargeBinary: pa.binary,
     sa.SmallInteger: pa.int16,
     sa.String: pa.string,
@@ -59,6 +61,8 @@ def _to_pyarrow_field(column: sa.Column) -> pa.Field:
     sa.TIMESTAMP: _SQL_ROOT_TYPES[sa.DateTime],
     sa.UUID: _SQL_ROOT_TYPES[sa.Uuid],
     sa.VARCHAR: _SQL_ROOT_TYPES[sa.String],
+    postgresql.JSON: _SQL_ROOT_TYPES[sa.JSON],
+    postgresql.JSONB: _SQL_ROOT_TYPES[sa.JSON],
 }
 
 _SQL_TYPE_MAP = _SQL_ROOT_TYPES | _EXTENDED_SQL_TYPES

elt-common/tests/unit_tests/sources/test_sqldatabase.py

@@ -1,10 +1,12 @@
 from pathlib import Path
 from typing import Optional
+from unittest.mock import patch
 
 import pyarrow as pa
 import pyarrow.lib
 import pytest
 import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
 
 from elt_common.extract import ResourceWriteProperties, Watermark
 from elt_common.sources.sqldatabase import SqlDatabaseExtract, SqlDatabaseSourceConfig, TableInfo
@@ -224,3 +226,57 @@ def table_info(self) -> dict[str, Optional[TableInfo]]:
 
     for table_name, _ in e.extract_resource_properties():
         assert table_name == "a_different_name"
+
+
+def test_sql_database_json_jsonb_serialization(tmp_path: Path):
+    db_path = tmp_path / "test_json.db"
+    metadata = sa.MetaData()
+
+    db_table = sa.Table(
+        "json_table",
+        metadata,
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column("json_col", sa.Text),
+        sa.Column("jsonb_col", sa.Text),
+    )
+
+    pg_table = sa.Table(
+        "json_table",
+        sa.MetaData(),
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column("json_col", postgresql.JSON),
+        sa.Column("jsonb_col", postgresql.JSONB),
+    )
+
+    engine = sa.create_engine(f"sqlite:///{db_path}")
+    metadata.create_all(engine)
+
+    with engine.begin() as conn:
+        conn.execute(
+            db_table.insert(),
+            [
+                {
+                    "id": 1,
+                    "json_col": '{"key": "val1", "nested": {"a": 1}}',
+                    "jsonb_col": '[1, 2, "a"]',
+                },
+            ],
+        )
+
+    source_config = _create_config(db_path)
+
+    class Extract(SqlDatabaseExtract):
+        def table_info(self) -> dict[str, Optional[TableInfo]]:
+            return {"json_table": None}
+
+    e = Extract(source_config)
+
+    with patch("sqlalchemy.Table", return_value=pg_table):
+        for _, props in e.extract_resource_properties():
+            tables = list(props.extractor(None))
+            data = pyarrow.lib.concat_tables(tables)
+
+            assert data.schema.field("json_col").type == pa.string()
+            assert data.schema.field("jsonb_col").type == pa.string()
+            assert data["json_col"].to_pylist() == ['{"key": "val1", "nested": {"a": 1}}']
+            assert data["jsonb_col"].to_pylist() == ['[1, 2, "a"]']

elt-common/tests/unit_tests/sources/test_sqldatabase_schema.py

@@ -1,6 +1,7 @@
 import pyarrow as pa
 import pytest
 import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
 
 from elt_common.sources.sqldatabase.schema import to_pyarrow_schema
 
@@ -49,6 +50,9 @@ def test_builds_schema_from_multiple_columns():
         (sa.DECIMAL(), pa.float64()),
         (sa.DECIMAL(10, 2), pa.decimal128(10, 2)),
         (sa.DECIMAL(50, 2), pa.decimal256(50, 2)),
+        (sa.JSON(), pa.json_()),
+        (postgresql.JSON(), pa.json_()),
+        (postgresql.JSONB(), pa.json_()),
     ],
 )
 def test_supported_sqlalchemy_types(sql_type, expected_type):

elt-pipelines/fase/ingest/fase/proposal/proposal.py

@@ -0,0 +1,27 @@
+from elt_common.extract import (
+    ResourceWriteProperties,
+    Watermark,  # noqa: F401
+)
+from elt_common.sources.sqldatabase import (
+    SqlDatabaseExtract,
+    SqlDatabaseSourceConfig,
+    TableInfo,
+)
+
+
+class PipelinePostgresConfig(SqlDatabaseSourceConfig):
+    drivername: str = "postgresql+psycopg"
+    tables: list[str]
+
+
+class Extract(SqlDatabaseExtract):
+    config_cls = PipelinePostgresConfig
+
+    def table_info(self) -> dict[str, TableInfo]:
+        """Defines the target tables and their ingestion strategy."""
+        return {
+            table_name: TableInfo(
+                write_properties=ResourceWriteProperties(write_mode="replace")
+            )
+            for table_name in self.config.tables
+        }

elt-pipelines/pyproject.toml

@@ -10,8 +10,10 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+proposal = [
+    "sqlalchemy[postgresql-psycopgbinary]>=2.0.0",
+]
 statusdisplay = [
-    "pyarrow>=24.0.0",
     "requests>=2.34.2",
 ]