diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b435f17..0e94aa4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,9 +2,9 @@ name: CI on: push: - branches: [ main, master ] + branches: [ main] pull_request: - branches: [ main, master ] + branches: [ main] schedule: # Run tests weekly on Monday at 00:00 UTC to catch dependency issues - cron: '0 0 * * 1' @@ -17,12 +17,14 @@ jobs: - name: Set up Python uses: actions/setup-python@v5 with: - python-version: "3.10" + python-version: "3.11" cache: "pip" - - name: Install hatch - run: pip install hatch - - name: Lint - run: hatch run lint + - name: Install Ruff + run: | + python -m pip install --upgrade pip + python -m pip install ruff + - name: Lint (Ruff, config from pyproject.toml) + run: ruff check --output-format=github . build: needs: lint @@ -42,5 +44,17 @@ jobs: - name: Type check run: hatch run types:check - name: Test - run: hatch run test -- --cov=iceberg_loader --cov=tests --cov-report=xml + run: hatch run test -- --cov=iceberg_loader --cov=tests --cov-report=xml --junitxml=pytest-report.xml + - name: Upload coverage report + uses: actions/upload-artifact@v4 + with: + name: coverage-${{ matrix.python-version }} + path: coverage.xml + if-no-files-found: error + - name: Upload pytest report + uses: actions/upload-artifact@v4 + with: + name: pytest-report-${{ matrix.python-version }} + path: pytest-report.xml + if-no-files-found: error diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index ff54ead..622f1dc 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -2,7 +2,7 @@ name: Documentation on: push: - branches: [ main, master ] + branches: [ main] workflow_dispatch: permissions: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 32403c1..8841c4e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -3,10 +3,10 @@ repos: rev: v0.6.9 hooks: - id: ruff - args: [--fix, "--target-version=py310"] + args: [--fix, "--config=pyproject.toml"] exclude: ^examples/ - id: ruff-format - args: ["--target-version=py310"] + args: ["--config=pyproject.toml"] exclude: ^examples/ - repo: https://github.com/pre-commit/mirrors-mypy rev: v1.13.0 diff --git a/examples/README.md b/examples/README.md index 3a69481..5594b7d 100644 --- a/examples/README.md +++ b/examples/README.md @@ -7,9 +7,9 @@ This directory contains runnable examples demonstrating various features of `ice You need a running Iceberg catalog (e.g., REST catalog) and MinIO/S3. A `docker-compose.yml` is provided to spin up a local environment. -```bash -docker-compose up -d -``` + ```bash + docker-compose up -d + ``` ## Running Examples diff --git a/examples/advanced_scenarios.py b/examples/advanced_scenarios.py index 94044f7..dc165ca 100644 --- a/examples/advanced_scenarios.py +++ b/examples/advanced_scenarios.py @@ -1,8 +1,14 @@ import logging +import sys +from pathlib import Path + +# Ensure parent directory (examples/) is on path +sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from catalog import get_catalog -from iceberg_loader import load_data_to_iceberg +from iceberg_loader import LoaderConfig, load_data_to_iceberg from iceberg_loader.arrow_utils import create_arrow_table_from_data logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') @@ -27,14 +33,8 @@ def scenario_initial_append(catalog): ] table_arrow = create_arrow_table_from_data(data_day_1) - load_data_to_iceberg( - table_data=table_arrow, - table_identifier=table_id, - catalog=catalog, - write_mode='append', - partition_col='ts', - schema_evolution=True, - ) + config = LoaderConfig(write_mode='append', partition_col='ts', schema_evolution=True) + load_data_to_iceberg(table_data=table_arrow, table_identifier=table_id, catalog=catalog, config=config) verify_table(catalog, table_id, expected_rows=2) @@ -47,13 +47,12 @@ def scenario_append_new_partition(catalog): {'id': 1, 'category': 'A', 'ts': '2023-01-01', 'value': 100}, {'id': 2, 'category': 'B', 'ts': '2023-01-01', 'value': 200}, ] + config = LoaderConfig(write_mode='append', partition_col='month(ts)', schema_evolution=True) load_data_to_iceberg( table_data=create_arrow_table_from_data(data_day_1), table_identifier=table_id, catalog=catalog, - write_mode='append', - partition_col='month(ts)', - schema_evolution=True, + config=config, ) # Append day 2 data_day_2 = [ @@ -63,7 +62,7 @@ def scenario_append_new_partition(catalog): table_data=create_arrow_table_from_data(data_day_2), table_identifier=table_id, catalog=catalog, - write_mode='append', + config=LoaderConfig(write_mode='append'), ) verify_table(catalog, table_id, expected_rows=3) @@ -78,13 +77,12 @@ def scenario_idempotent_replace_partition(catalog): {'id': 2, 'category': 'B', 'ts': '2023-01-01', 'value': 200}, {'id': 3, 'category': 'A', 'ts': '2023-01-02', 'value': 150}, ] + config_base = LoaderConfig(write_mode='append', partition_col='ts', schema_evolution=True) load_data_to_iceberg( table_data=create_arrow_table_from_data(base_data), table_identifier=table_id, catalog=catalog, - write_mode='append', - partition_col='ts', - schema_evolution=True, + config=config_base, ) # Corrected day1 @@ -92,12 +90,12 @@ def scenario_idempotent_replace_partition(catalog): {'id': 1, 'category': 'A', 'ts': '2023-01-01', 'value': 999}, {'id': 2, 'category': 'B', 'ts': '2023-01-01', 'value': 200}, ] + config_replace = LoaderConfig(write_mode='append', replace_filter="ts == '2023-01-01'") load_data_to_iceberg( table_data=create_arrow_table_from_data(corrected_day1), table_identifier=table_id, catalog=catalog, - write_mode='append', - replace_filter="ts == '2023-01-01'", + config=config_replace, ) verify_table(catalog, table_id, expected_rows=3) @@ -109,13 +107,12 @@ def scenario_schema_evolution(catalog): base_data = [ {'id': 1, 'category': 'A', 'ts': '2023-01-01', 'value': 100}, ] + config_base = LoaderConfig(write_mode='append', partition_col='ts', schema_evolution=True) load_data_to_iceberg( table_data=create_arrow_table_from_data(base_data), table_identifier=table_id, catalog=catalog, - write_mode='append', - partition_col='ts', - schema_evolution=True, + config=config_base, ) evolved = [ @@ -125,8 +122,7 @@ def scenario_schema_evolution(catalog): table_data=create_arrow_table_from_data(evolved), table_identifier=table_id, catalog=catalog, - write_mode='append', - schema_evolution=True, + config=LoaderConfig(write_mode='append', schema_evolution=True), ) verify_table(catalog, table_id, expected_rows=2) @@ -145,13 +141,12 @@ def scenario_full_overwrite(catalog): {'id': 1, 'category': 'A', 'ts': '2023-01-01', 'value': 100}, {'id': 2, 'category': 'B', 'ts': '2023-01-02', 'value': 200}, ] + config_base = LoaderConfig(write_mode='append', partition_col='ts', schema_evolution=True) load_data_to_iceberg( table_data=create_arrow_table_from_data(initial), table_identifier=table_id, catalog=catalog, - write_mode='append', - partition_col='ts', - schema_evolution=True, + config=config_base, ) replace_all = [ @@ -161,8 +156,7 @@ def scenario_full_overwrite(catalog): table_data=create_arrow_table_from_data(replace_all), table_identifier=table_id, catalog=catalog, - write_mode='overwrite', - schema_evolution=True, + config=LoaderConfig(write_mode='overwrite', schema_evolution=True), ) verify_table(catalog, table_id, expected_rows=1) diff --git a/examples/compare_complex_json_fail.py b/examples/compare_complex_json_fail.py index 721b8e4..d2632ff 100644 --- a/examples/compare_complex_json_fail.py +++ b/examples/compare_complex_json_fail.py @@ -1,9 +1,16 @@ import logging +import sys +from pathlib import Path import pyarrow as pa + +# Ensure parent directory (examples/) is on path +sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + from catalog import get_catalog -from iceberg_loader import load_data_to_iceberg +from iceberg_loader import LoaderConfig, load_data_to_iceberg from iceberg_loader.arrow_utils import create_arrow_table_from_data logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') @@ -40,12 +47,12 @@ def run_comparison(): logger.info("Created Arrow table with schema:\n%s", arrow_table.schema) # 2. Load to Iceberg + config = LoaderConfig(write_mode='overwrite', schema_evolution=True) load_data_to_iceberg( table_data=arrow_table, table_identifier=table_id, catalog=catalog, - write_mode='overwrite', - schema_evolution=True, + config=config, ) logger.info("Successfully loaded data to Iceberg table '%s'", table_id) diff --git a/examples/load_complex_json.py b/examples/load_complex_json.py index bc85cac..8cdd979 100644 --- a/examples/load_complex_json.py +++ b/examples/load_complex_json.py @@ -1,9 +1,15 @@ import json import logging +import sys +from pathlib import Path + +# Ensure parent directory (examples/) is on path +sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from catalog import get_catalog -from iceberg_loader import load_data_to_iceberg +from iceberg_loader import LoaderConfig, load_data_to_iceberg from iceberg_loader.arrow_utils import create_arrow_table_from_data logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') @@ -43,9 +49,8 @@ def run_complex_load(): logger.info(arrow_table.to_pydict()) logger.info('\nLoading to Iceberg...') - load_data_to_iceberg( - table_data=arrow_table, table_identifier=table_id, catalog=catalog, write_mode='append', schema_evolution=True - ) + config = LoaderConfig(write_mode='append', schema_evolution=True) + load_data_to_iceberg(table_data=arrow_table, table_identifier=table_id, catalog=catalog, config=config) logger.info('\nVerifying data in Iceberg...') table = catalog.load_table(table_id) diff --git a/examples/load_example.py b/examples/load_example.py index ce5d0df..40bf627 100644 --- a/examples/load_example.py +++ b/examples/load_example.py @@ -2,7 +2,7 @@ from catalog import get_catalog -from iceberg_loader import load_data_to_iceberg +from iceberg_loader import LoaderConfig, load_data_to_iceberg logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) @@ -54,16 +54,17 @@ def run_example(): except Exception: pass + config = LoaderConfig( + write_mode='append', + partition_col='signup_date', + schema_evolution=True, + ) + result = load_data_to_iceberg( table_data=arrow_table, table_identifier=table_id, catalog=catalog, - write_mode='append', - # New flexible partitioning: - partition_col='signup_date', - # Optional: idempotency filter (if we were reloading data for a specific date) - # replace_filter="signup_date == '2023-01-01'", - schema_evolution=True, + config=config, ) logger.info('Load result: %s', result) diff --git a/examples/load_upsert.py b/examples/load_upsert.py index ec224d2..ec7ffaa 100644 --- a/examples/load_upsert.py +++ b/examples/load_upsert.py @@ -1,11 +1,18 @@ import logging +import sys import time from datetime import datetime +from pathlib import Path import pyarrow as pa + +# Ensure parent directory (examples/) is on path +sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + from catalog import get_catalog -from iceberg_loader import load_data_to_iceberg +from iceberg_loader import LoaderConfig, load_data_to_iceberg logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) @@ -32,12 +39,8 @@ def run_upsert_example(): } ) - load_data_to_iceberg( - initial_data, - table_identifier, - catalog, - write_mode="overwrite", - ) + config_overwrite = LoaderConfig(write_mode="overwrite") + load_data_to_iceberg(initial_data, table_identifier, catalog, config=config_overwrite) table = catalog.load_table(table_identifier) rows = table.scan().to_arrow() @@ -56,13 +59,8 @@ def run_upsert_example(): } ) - load_data_to_iceberg( - upsert_data, - table_identifier, - catalog, - write_mode="upsert", - join_cols=["id"], - ) + config_upsert = LoaderConfig(write_mode="upsert", join_cols=["id"]) + load_data_to_iceberg(upsert_data, table_identifier, catalog, config_upsert) rows_after = table.scan().to_arrow() logger.info("Rows after upsert: %d", len(rows_after)) diff --git a/examples/load_with_commits.py b/examples/load_with_commits.py index d571250..34bd244 100644 --- a/examples/load_with_commits.py +++ b/examples/load_with_commits.py @@ -1,10 +1,17 @@ import logging +import sys import time +from pathlib import Path import pyarrow as pa + +# Ensure parent directory (examples/) is on path +sys.path.insert(0, str(Path(__file__).resolve().parent)) +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + from catalog import get_catalog -from iceberg_loader import load_batches_to_iceberg +from iceberg_loader import LoaderConfig, load_batches_to_iceberg logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) @@ -38,12 +45,13 @@ def run_example(): # We will load 20 batches, committing every 5 batches. # This means we expect roughly 4 snapshots (transactions) to be created. + config = LoaderConfig(write_mode='append', commit_interval=5) + result = load_batches_to_iceberg( batch_iterator=generate_batches(num_batches=20, batch_size=100), table_identifier=table_id, catalog=catalog, - write_mode='append', - commit_interval=5, # Commit every 5 batches + config=config, ) logger.info('Load complete. Result: %s', result) diff --git a/src/iceberg_loader/__init__.py b/src/iceberg_loader/__init__.py index b56add3..2bffe41 100644 --- a/src/iceberg_loader/__init__.py +++ b/src/iceberg_loader/__init__.py @@ -5,6 +5,7 @@ from iceberg_loader.__about__ import __version__ from iceberg_loader.iceberg_loader import ( IcebergLoader, + LoaderConfig, load_batches_to_iceberg, load_data_to_iceberg, load_ipc_stream_to_iceberg, @@ -18,4 +19,5 @@ 'expire_snapshots', '__version__', 'IcebergLoader', + 'LoaderConfig', ] diff --git a/src/iceberg_loader/iceberg_loader.py b/src/iceberg_loader/iceberg_loader.py index ec346c1..2e7c586 100644 --- a/src/iceberg_loader/iceberg_loader.py +++ b/src/iceberg_loader/iceberg_loader.py @@ -1,5 +1,6 @@ import logging from collections.abc import Iterator +from dataclasses import dataclass from typing import Any, BinaryIO, Literal import pyarrow as pa @@ -13,29 +14,48 @@ logger = logging.getLogger(__name__) +@dataclass +class LoaderConfig: + write_mode: Literal['overwrite', 'append', 'upsert'] = 'overwrite' + partition_col: str | None = None + replace_filter: str | None = None + schema_evolution: bool = False + table_properties: dict[str, Any] | None = None + commit_interval: int = 0 + join_cols: list[str] | None = None + + class IcebergLoader: """ Facade for loading data into Iceberg tables. Orchestrates SchemaManager and WriteStrategy to handle complex ingestion scenarios. """ - def __init__(self, catalog: Catalog, table_properties: dict[str, Any] | None = None): + def __init__( + self, + catalog: Catalog, + table_properties: dict[str, Any] | None = None, + default_config: LoaderConfig | None = None, + ): self.catalog = catalog self.table_properties = TABLE_PROPERTIES.copy() if table_properties: self.table_properties.update(table_properties) self.schema_manager = SchemaManager(self.catalog, self.table_properties) + self.default_config = default_config or LoaderConfig() def load_data( self, table_data: pa.Table, table_identifier: tuple[str, str], - write_mode: Literal['overwrite', 'append', 'upsert'] = 'overwrite', + write_mode: Literal['overwrite', 'append', 'upsert'] | None = None, partition_col: str | None = None, replace_filter: str | None = None, - schema_evolution: bool = False, + schema_evolution: bool | None = None, + table_properties: dict[str, Any] | None = None, join_cols: list[str] | None = None, + config: LoaderConfig | None = None, ) -> dict[str, Any]: """ Load PyArrow Table into Iceberg table. @@ -49,19 +69,23 @@ def load_data( partition_col=partition_col, replace_filter=replace_filter, schema_evolution=schema_evolution, + table_properties=table_properties, join_cols=join_cols, + config=config, ) def load_ipc_stream( self, stream_source: str | BinaryIO | pa.NativeFile, table_identifier: tuple[str, str], - write_mode: Literal['overwrite', 'append', 'upsert'] = 'overwrite', + write_mode: Literal['overwrite', 'append', 'upsert'] | None = None, partition_col: str | None = None, replace_filter: str | None = None, - schema_evolution: bool = False, - commit_interval: int = 0, + schema_evolution: bool | None = None, + commit_interval: int | None = None, join_cols: list[str] | None = None, + table_properties: dict[str, Any] | None = None, + config: LoaderConfig | None = None, ) -> dict[str, Any]: """Loads data from an Apache Arrow IPC stream source.""" with pa.ipc.open_stream(stream_source) as reader: @@ -74,18 +98,22 @@ def load_ipc_stream( schema_evolution=schema_evolution, commit_interval=commit_interval, join_cols=join_cols, + table_properties=table_properties, + config=config, ) def load_data_batches( self, batch_iterator: Iterator[pa.RecordBatch] | pa.RecordBatchReader, table_identifier: tuple[str, str], - write_mode: Literal['overwrite', 'append', 'upsert'] = 'overwrite', + write_mode: Literal['overwrite', 'append', 'upsert'] | None = None, partition_col: str | None = None, replace_filter: str | None = None, - schema_evolution: bool = False, - commit_interval: int = 0, + schema_evolution: bool | None = None, + commit_interval: int | None = None, join_cols: list[str] | None = None, + table_properties: dict[str, Any] | None = None, + config: LoaderConfig | None = None, ) -> dict[str, Any]: """ Main orchestration method. @@ -97,8 +125,19 @@ def load_data_batches( # Buffer pending_batches: list[pa.RecordBatch] = [] - # Strategy selection - strategy = get_write_strategy(write_mode, replace_filter, join_cols) + resolved = self._resolve_config( + config=config, + write_mode=write_mode, + partition_col=partition_col, + replace_filter=replace_filter, + schema_evolution=schema_evolution, + commit_interval=commit_interval, + join_cols=join_cols, + table_properties=table_properties, + ) + self.table_properties = resolved['table_properties'] + strategy = get_write_strategy(resolved['write_mode'], resolved['replace_filter'], resolved['join_cols']) + effective_schema_evolution = resolved['schema_evolution'] # State tracking table = None @@ -113,7 +152,7 @@ def process_buffer(batches: list[pa.RecordBatch]) -> None: combined_table = None - if schema_evolution: + if effective_schema_evolution: # Try fast path try: combined_table = pa.Table.from_batches(batches) @@ -122,9 +161,8 @@ def process_buffer(batches: list[pa.RecordBatch]) -> None: logger.info('Mixed schemas in batch buffer. Normalizing...') if table is None: - # Use first batch to ensure table exists table = self.schema_manager.ensure_table_exists( - table_identifier, batches[0].schema, partition_col + table_identifier, batches[0].schema, resolved['partition_col'] ) if table.current_snapshot() is None: new_table_created = True @@ -149,19 +187,20 @@ def process_buffer(batches: list[pa.RecordBatch]) -> None: # 1. Ensure Table Exists if table is None: - table = self.schema_manager.ensure_table_exists(table_identifier, combined_table.schema, partition_col) + table = self.schema_manager.ensure_table_exists( + table_identifier, combined_table.schema, resolved['partition_col'] + ) if table.current_snapshot() is None: new_table_created = True # 2. Schema Evolution - if schema_evolution: + if effective_schema_evolution: self.schema_manager.evolve_schema_if_needed(table, combined_table.schema) # 3. Type Conversion target_schema = self.schema_manager.get_arrow_schema(table) combined_table = convert_table_types(combined_table, target_schema) - # 4. Write via Strategy # The strategy now handles transaction management internally strategy.write(table, combined_table, is_first_write) @@ -173,7 +212,7 @@ def process_buffer(batches: list[pa.RecordBatch]) -> None: pending_batches.append(batch) batches_processed += 1 - limit = 1 if commit_interval <= 1 else commit_interval + limit = 1 if resolved['commit_interval'] <= 1 else resolved['commit_interval'] if len(pending_batches) >= limit: process_buffer(pending_batches) @@ -185,14 +224,41 @@ def process_buffer(batches: list[pa.RecordBatch]) -> None: return { 'rows_loaded': total_rows, - 'write_mode': write_mode, - 'partition_col': partition_col if partition_col else 'none', + 'write_mode': resolved['write_mode'], + 'partition_col': resolved['partition_col'] if resolved['partition_col'] else 'none', 'table_location': table.location() if table else 'none', 'snapshot_id': table.current_snapshot().snapshot_id if table and table.current_snapshot() else 'none', 'batches_processed': batches_processed, 'new_table_created': new_table_created, } + def _resolve_config( + self, + config: LoaderConfig | None, + write_mode: Literal['overwrite', 'append', 'upsert'] | None, + partition_col: str | None, + replace_filter: str | None, + schema_evolution: bool | None, + commit_interval: int | None, + join_cols: list[str] | None, + table_properties: dict[str, Any] | None, + ) -> dict[str, Any]: + base = config or self.default_config + merged_table_props = self.table_properties.copy() + if base.table_properties: + merged_table_props.update(base.table_properties) + if table_properties: + merged_table_props.update(table_properties) + return { + 'write_mode': write_mode or base.write_mode, + 'partition_col': partition_col if partition_col is not None else base.partition_col, + 'replace_filter': replace_filter if replace_filter is not None else base.replace_filter, + 'schema_evolution': base.schema_evolution if schema_evolution is None else schema_evolution, + 'commit_interval': base.commit_interval if commit_interval is None else commit_interval, + 'join_cols': join_cols if join_cols is not None else base.join_cols, + 'table_properties': merged_table_props, + } + # Public API functions (thin wrappers) @@ -201,14 +267,15 @@ def load_data_to_iceberg( table_data: pa.Table, table_identifier: tuple[str, str], catalog: Catalog, - write_mode: Literal['overwrite', 'append', 'upsert'] = 'overwrite', + write_mode: Literal['overwrite', 'append', 'upsert'] | None = None, partition_col: str | None = None, replace_filter: str | None = None, - schema_evolution: bool = False, + schema_evolution: bool | None = None, table_properties: dict[str, Any] | None = None, join_cols: list[str] | None = None, + config: LoaderConfig | None = None, ) -> dict[str, Any]: - loader = IcebergLoader(catalog, table_properties) + loader = IcebergLoader(catalog, table_properties, default_config=config) return loader.load_data( table_data, table_identifier, @@ -216,7 +283,9 @@ def load_data_to_iceberg( partition_col, replace_filter, schema_evolution, + table_properties, join_cols, + config=config, ) @@ -224,15 +293,16 @@ def load_batches_to_iceberg( batch_iterator: Iterator[pa.RecordBatch] | pa.RecordBatchReader, table_identifier: tuple[str, str], catalog: Catalog, - write_mode: Literal['overwrite', 'append', 'upsert'] = 'overwrite', + write_mode: Literal['overwrite', 'append', 'upsert'] | None = None, partition_col: str | None = None, replace_filter: str | None = None, - schema_evolution: bool = False, + schema_evolution: bool | None = None, table_properties: dict[str, Any] | None = None, - commit_interval: int = 0, + commit_interval: int | None = None, join_cols: list[str] | None = None, + config: LoaderConfig | None = None, ) -> dict[str, Any]: - loader = IcebergLoader(catalog, table_properties) + loader = IcebergLoader(catalog, table_properties, default_config=config) return loader.load_data_batches( batch_iterator, table_identifier, @@ -242,6 +312,8 @@ def load_batches_to_iceberg( schema_evolution, commit_interval, join_cols, + table_properties, + config, ) @@ -249,15 +321,16 @@ def load_ipc_stream_to_iceberg( stream_source: str | BinaryIO | pa.NativeFile, table_identifier: tuple[str, str], catalog: Catalog, - write_mode: Literal['overwrite', 'append', 'upsert'] = 'overwrite', + write_mode: Literal['overwrite', 'append', 'upsert'] | None = None, partition_col: str | None = None, replace_filter: str | None = None, - schema_evolution: bool = False, + schema_evolution: bool | None = None, table_properties: dict[str, Any] | None = None, - commit_interval: int = 0, + commit_interval: int | None = None, join_cols: list[str] | None = None, + config: LoaderConfig | None = None, ) -> dict[str, Any]: - loader = IcebergLoader(catalog, table_properties) + loader = IcebergLoader(catalog, table_properties, default_config=config) return loader.load_ipc_stream( stream_source, table_identifier, @@ -267,4 +340,6 @@ def load_ipc_stream_to_iceberg( schema_evolution, commit_interval, join_cols, + table_properties, + config, ) diff --git a/tests/test_iceberg_loader.py b/tests/test_iceberg_loader.py index e069565..f70308c 100644 --- a/tests/test_iceberg_loader.py +++ b/tests/test_iceberg_loader.py @@ -105,7 +105,7 @@ def test_public_api_wrapper(self): load_data_to_iceberg(self.arrow_table, self.table_identifier, self.mock_catalog) - mock_loader_cls.assert_called_with(self.mock_catalog, None) + mock_loader_cls.assert_called_with(self.mock_catalog, None, default_config=None) mock_instance.load_data.assert_called_once() def test_field_ids_preserved_on_evolution(self):