feat(elt-pipelines): Proposal postgresql sources pipeline

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

@@ -1,25 +1,32 @@
 """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 typing import Callable, Generator, Iterator, NamedTuple, Optional
 
 import pyarrow as pa
 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__)
 
+DEFAULT_PA_TYPE_MAPPING = {
+    "bigint": pa.int64(),
+    "bool": pa.bool_(),
+    "double": pa.float64(),
+    "timestamp": pa.timestamp("us"),
+    "date": pa.date32(),
+    "text": pa.string(),
+}
 
-class SqlDatabaseSourceConfig(BaseSettings):
-    """Configuration required to connect to a database"""
 
-    # connection
+class SqlDatabaseSourceConfig(BaseSettings):
     drivername: str
     database: str
     database_schema: Optional[str] = None
@@ -28,7 +35,6 @@ class SqlDatabaseSourceConfig(BaseSettings):
     username: Optional[str] = None
     password: Optional[SecretStr] = 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"""
 
@@ -100,6 +106,25 @@ def __init__(self, config: SqlDatabaseSourceConfig):
         self._engine = sa.create_engine(config.connection_url)
         self._metadata = sa.MetaData(schema=config.database_schema)
 
+    def get_table_schema(self, table_name: str) -> pa.Schema:
+        """Autoloads the SQL table and constructs a PyArrow schema using schema.py."""
+        table = sa.Table(
+            table_name,
+            self._metadata,
+            autoload_with=self._engine,
+        )
+        schema = to_pyarrow_schema(table)
+
+        # Apply column normalization (e.g. lowercase) if overridden by subclasses
+        normalized_fields = [
+            pa.field(self.normalize_column_name(field.name), field.type, field.nullable)
+            for field in schema
+        ]
+        return pa.schema(normalized_fields)
+
+    def normalize_column_name(self, name: str) -> str:
+        return name
+
     @abstractmethod
     def table_info(self) -> dict[str, Optional[TableInfo]]:
         """Define the tables to be extracted from the DB.
@@ -179,11 +204,66 @@ def _extract_table(
 
         query = query.limit(self._row_limit)
 
-        # If all the values in a column are null pyarrow won't know what type
-        # the column should be, so we need to explicitly create a schema from
-        # the table
-        pa_schema = to_pyarrow_schema(table)
         result = conn.execution_options(yield_per=self._chunk_size).execute(query)
-        for partition in result.mappings().partitions():
-            table = pa.Table.from_pylist(partition, schema=pa_schema)
-            yield table
+
+        target_schema = self.get_table_schema(name)
+
+        column_names = list(result.keys())
+
+        has_data = False
+        while True:
+            rows = result.fetchmany(self._chunk_size)
+
+            if not rows:
+                break
+
+            has_data = True
+
+            # Convert SQLAlchemy Row objects to column arrays
+            columns = {}
+
+            for idx, column_name in enumerate(column_names):
+                columns[column_name] = [row[idx] for row in rows]
+
+            arrow_arrays = []
+
+            for field in target_schema:
+                values = columns.get(
+                    field.name,
+                    [None] * len(rows),
+                )
+
+                # JSON / JSONB -> string for Iceberg
+                if pa.types.is_string(field.type):
+                    values = [
+                        json.dumps(v)
+                        if isinstance(v, (dict, list))
+                        else str(v)
+                        if v is not None and not isinstance(v, str)
+                        else v
+                        for v in values
+                    ]
+                elif pa.types.is_integer(field.type):
+                    values = [int(v) if v is not None else None for v in values]
+                elif pa.types.is_floating(field.type):
+                    values = [float(v) if v is not None else None for v in values]
+
+                array = pa.array(
+                    values,
+                    type=field.type,
+                )
+                arrow_arrays.append(array)
+
+            yield pa.Table.from_arrays(
+                arrow_arrays,
+                schema=target_schema,
+            )
+
+        # Return empty table with schema when no rows
+        if not has_data:
+            empty_arrays = [pa.array([], type=field.type) for field in target_schema]
+
+            yield pa.Table.from_arrays(
+                empty_arrays,
+                schema=target_schema,
+            )

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.string,
     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.String],
+    postgresql.JSONB: _SQL_ROOT_TYPES[sa.String],
 }
 
 _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,53 @@ 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"}', "jsonb_col": '{"key": "val2"}'},
+            ],
+        )
+
+    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"}']
+            assert data["jsonb_col"].to_pylist() == ['{"key": "val2"}']

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.string()),
+        (postgresql.JSON(), pa.string()),
+        (postgresql.JSONB(), pa.string()),
     ],
 )
 def test_supported_sqlalchemy_types(sql_type, expected_type):

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

@@ -0,0 +1,31 @@
+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"
+    table: str
+
+    @property
+    def target_tables(self) -> list[str]:
+        return [t.strip() for t in self.table.split(",") if t.strip()]
+
+
+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.target_tables
+        }

elt-pipelines/pyproject.toml

@@ -10,8 +10,12 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-statusdisplay = [
+proposal = [
     "pyarrow>=24.0.0",
+    "sqlalchemy>=2.0.0",
+    "psycopg[binary]>=3.1.0",
+]
+statusdisplay = [
     "requests>=2.34.2",
 ]
 

infra/local/warehouses/lakekeeper/fase_landing.json

@@ -0,0 +1,27 @@
+{
+  "warehouse-name": "fase_landing",
+  "storage-credential": {
+    "type": "s3",
+    "aws-access-key-id": "adpsuperuser",
+    "aws-secret-access-key": "adppassword",
+    "credential-type": "access-key"
+  },
+  "storage-profile": {
+    "type": "s3",
+    "bucket": "fase-landing",
+    "key-prefix": "iceberg",
+    "endpoint": "http://adp-router:59000",
+    "region": "local-01",
+    "path-style-access": true,
+    "sts-enabled": false,
+    "flavor": "s3-compat"
+  },
+  "delete-profile": {
+    "type": "hard"
+  },
+  "permissions": {
+    "service-account-trino": [
+      "select"
+    ]
+  }
+}