[DOP-21408] Allow using etl-entities==2.4.0

MobileTeleSystems · Nov 12, 2024 · eb3177f · eb3177f
1 parent 91b83b1
commit eb3177f
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 6 deletions.
diff --git a/docs/changelog/next_release/+.dependency.rst b/docs/changelog/next_release/+.dependency.rst
@@ -0,0 +1 @@
+Allow using ``etl-entities==2.4.0``.
diff --git a/onetl/connection/db_connection/kafka/connection.py b/onetl/connection/db_connection/kafka/connection.py
@@ -527,13 +527,14 @@ def get_min_max_values(
                 max_offsets[partition_id] = end_offset
 
         log.info("|%s| Received min and max offset values for each partition.", self.__class__.__name__)
-        for partition_id in sorted(min_offsets.keys()):
+        partitions = sorted(set(min_offsets.keys() | max_offsets.keys()))
+        for partition_id in partitions:
             log.debug(
                 "|%s| Partition %d: Min Offset = %d, Max Offset = %d",
                 self.__class__.__name__,
                 partition_id,
-                min_offsets[partition_id],
-                max_offsets[partition_id],
+                min_offsets.get(partition_id),
+                max_offsets.get(partition_id),
             )
 
         return min_offsets, max_offsets

diff --git a/requirements/core.txt b/requirements/core.txt
@@ -1,4 +1,4 @@
-etl-entities>=2.2,<2.4
+etl-entities>=2.2,<2.5
 evacuator>=1.0,<1.1
 frozendict
 humanize

diff --git a/...ategy_integration/tests_incremental_strategy_integration/test_strategy_increment_kafka.py b/...ategy_integration/tests_incremental_strategy_integration/test_strategy_increment_kafka.py
@@ -28,7 +28,6 @@ def test_kafka_strategy_incremental(
 ):
     from pyspark.sql.functions import max as spark_max
 
-    hwm_type = KeyValueIntHWM
     hwm_name = secrets.token_hex(5)
     store = HWMStoreStackManager.get_current()
 
@@ -77,7 +76,7 @@ def test_kafka_strategy_incremental(
 
     hwm = store.get_hwm(hwm_name)
     assert hwm is not None
-    assert isinstance(hwm, hwm_type)
+    assert isinstance(hwm, KeyValueIntHWM)
 
     # HWM contains mapping `partition: max offset + 1`
     partition_offsets_initial = dict.fromkeys(range(num_partitions or 1), 0)