diff --git a/src/snowflake/snowpark/dataframe_reader.py b/src/snowflake/snowpark/dataframe_reader.py index 2b855251ce..31ac944a79 100644 --- a/src/snowflake/snowpark/dataframe_reader.py +++ b/src/snowflake/snowpark/dataframe_reader.py @@ -325,31 +325,30 @@ def _extract_iceberg_changes_from_options(options: dict) -> dict: def _reader_options_conflict_with_incremental_read(options: dict) -> list[str]: """Return reader option keys that cannot coexist with incremental read.""" - incremental_keys = { - "start-snapshot-id", - "start_snapshot_id", - "end-snapshot-id", - "end_snapshot_id", - } - if not any( - k.upper().replace("_", "-") in {x.replace("_", "-") for x in incremental_keys} - for k in options - ): + + def normalize(key: str) -> str: + return key.upper().replace("_", "-") + + incremental_keys = {"START-SNAPSHOT-ID", "END-SNAPSHOT-ID"} + normalized_keys = {normalize(key) for key in options} + if not normalized_keys.intersection(incremental_keys): return [] + + time_travel_keys = { + *(normalize(key) for key in _TIME_TRAVEL_OPTIONS_PARAMS_MAP), + "SNAPSHOT-ID", + "AS-OF-TIMESTAMP", + "VERSION-TAG", + "VERSION-REF", + "BRANCH", + "TAG", + } blocked = [] for key in options: - upper = key.upper() - if upper in incremental_keys or upper.replace("_", "-") in { - x.replace("_", "-") for x in incremental_keys - }: + normalized_key = normalize(key) + if normalized_key in incremental_keys: continue - if upper in _TIME_TRAVEL_OPTIONS_PARAMS_MAP or upper in ( - "SNAPSHOT-ID", - "SNAPSHOT_ID", - "AS-OF-TIMESTAMP", - "VERSION_TAG", - "VERSION-TAG", - ): + if normalized_key in time_travel_keys: blocked.append(key) return blocked diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index beeb5d19a2..74591c9baa 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -1300,6 +1300,60 @@ def test_extract_time_travel_branch_and_version_ref_options(): ) +@pytest.mark.parametrize( + "incremental_key", + [ + "start-snapshot-id", + "START-SNAPSHOT-ID", + "start_snapshot_id", + "START_SNAPSHOT_ID", + "end-snapshot-id", + "END_SNAPSHOT_ID", + ], +) +@pytest.mark.parametrize( + "conflicting_key", + [ + "time_travel_mode", + "statement", + "offset", + "timestamp", + "timestamp_type", + "stream", + "snapshot-id", + "as-of-timestamp", + "version-tag", + "version_ref", + "branch", + "tag", + ], +) +def test_reader_options_conflict_with_incremental_read( + incremental_key, conflicting_key +): + from snowflake.snowpark.dataframe_reader import ( + _reader_options_conflict_with_incremental_read, + ) + + options = {incremental_key: 1, conflicting_key: "value"} + + assert _reader_options_conflict_with_incremental_read(options) == [conflicting_key] + + +def test_reader_options_conflict_with_incremental_read_ignores_unrelated_options(): + from snowflake.snowpark.dataframe_reader import ( + _reader_options_conflict_with_incremental_read, + ) + + options = { + "START_SNAPSHOT_ID": 1, + "END-SNAPSHOT-ID": 2, + "compression": "gzip", + } + + assert _reader_options_conflict_with_incremental_read(options) == [] + + def test_time_travel_version_ref_validation(): with pytest.raises(ValueError, match="'version_ref' must be a non-empty"): TimeTravelConfig.validate_and_normalize_params(