diff --git a/elt-common/src/elt_common/sources/sqldatabase/__init__.py b/elt-common/src/elt_common/sources/sqldatabase/__init__.py index 9e42615d..23837346 100644 --- a/elt-common/src/elt_common/sources/sqldatabase/__init__.py +++ b/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, + ) diff --git a/elt-common/src/elt_common/sources/sqldatabase/schema.py b/elt-common/src/elt_common/sources/sqldatabase/schema.py index a49ffcc5..ddf5547a 100644 --- a/elt-common/src/elt_common/sources/sqldatabase/schema.py +++ b/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 diff --git a/elt-common/tests/unit_tests/sources/test_sqldatabase.py b/elt-common/tests/unit_tests/sources/test_sqldatabase.py index e00a5bca..c8d59e83 100644 --- a/elt-common/tests/unit_tests/sources/test_sqldatabase.py +++ b/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"}'] diff --git a/elt-common/tests/unit_tests/sources/test_sqldatabase_schema.py b/elt-common/tests/unit_tests/sources/test_sqldatabase_schema.py index df19aa2b..f97d9a4f 100644 --- a/elt-common/tests/unit_tests/sources/test_sqldatabase_schema.py +++ b/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): diff --git a/elt-pipelines/fase/ingest/fase/proposal/proposal.py b/elt-pipelines/fase/ingest/fase/proposal/proposal.py new file mode 100644 index 00000000..1e699e98 --- /dev/null +++ b/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 + } diff --git a/elt-pipelines/pyproject.toml b/elt-pipelines/pyproject.toml index a0a767ea..ff112369 100644 --- a/elt-pipelines/pyproject.toml +++ b/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", ] diff --git a/infra/local/warehouses/lakekeeper/fase_landing.json b/infra/local/warehouses/lakekeeper/fase_landing.json new file mode 100644 index 00000000..daa217e8 --- /dev/null +++ b/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" + ] + } +}