[DOP-22139] Reset HWM when changing strategy

Author: Ilyas Gasanov

.env.docker

@@ -39,6 +39,7 @@ SYNCMASTER__AUTH__ACCESS_TOKEN__SECRET_KEY=generate_another_random_string
 SYNCMASTER__SCHEDULER__TRANSFER_FETCHING_TIMEOUT_SECONDS=200
 
 # HWM Store
+SYNCMASTER__HWM_STORE__ENABLED=true
 SYNCMASTER__HWM_STORE__TYPE=horizon
 SYNCMASTER__HWM_STORE__URL=http://horizon:8000
 SYNCMASTER__HWM_STORE__NAMESPACE=syncmaster_namespace

.env.local

@@ -39,6 +39,7 @@ export SYNCMASTER__AUTH__ACCESS_TOKEN__SECRET_KEY=generate_another_random_string
 export SYNCMASTER__SCHEDULER__TRANSFER_FETCHING_TIMEOUT_SECONDS=200
 
 # HWM Store
+export SYNCMASTER__HWM_STORE__ENABLED=true
 export SYNCMASTER__HWM_STORE__TYPE=horizon
 export SYNCMASTER__HWM_STORE__URL=http://localhost:8020
 export SYNCMASTER__HWM_STORE__NAMESPACE=syncmaster_namespace

docs/changelog/next_release/217.improvement.rst

@@ -0,0 +1 @@
+Reset HWM when changing strategy from `incremental` to `full`

poetry.lock

@@ -60,7 +60,7 @@ pyjwt = { version = "^2.10.1", optional = true }
 jinja2 = { version = "^3.1.4", optional = true }
 python-multipart = { version = ">=0.0.9,<0.0.21", optional = true }
 celery = { version = "^5.4.0", optional = true }
-onetl = { version = "^0.13.1", extras = ["all"], optional = true }
+onetl = { version = "^0.13.3", extras = ["all"], optional = true }
 pyyaml = { version = "*", optional = true }
 # due to not supporting MacOS 14.x https://www.psycopg.org/psycopg3/docs/news.html#psycopg-3-1-20
 psycopg = { version = ">=3.1.0,<3.2.6", extras = ["binary"], optional = true }
@@ -72,7 +72,7 @@ apscheduler = { version = "^3.10.4", optional = true }
 starlette-exporter = {version = "^0.23.0", optional = true}
 itsdangerous = {version = "*", optional = true}
 python-keycloak = {version = ">=4.7,<6.0", optional = true}
-horizon-hwm-store = {version = ">=1.1.1", optional = true }
+horizon-hwm-store = {version = ">=1.1.2", optional = true }
 
 [tool.poetry.extras]
 server = [

pyproject.toml

@@ -198,6 +198,9 @@ def perform_transfer(self) -> None:
                 if self.source_handler.transfer_dto.strategy.type == "incremental":
                     return self._perform_incremental_transfer()
 
+                if self.source_handler.transfer_dto.strategy.type == "full" and self.settings.hwm_store.enabled:
+                    self._reset_transfer_hwm()
+
                 df = self.source_handler.read()
                 self.target_handler.write(df)
         finally:
@@ -243,21 +246,38 @@ def _perform_incremental_transfer(self) -> None:
         ).force_create_namespace() as hwm_store:
 
             with IncrementalStrategy():
-                if self.source_handler.connection_dto.type in FILE_CONNECTION_TYPES:
-                    hwm_name_suffix = self.source_handler.transfer_dto.directory_path
-                else:
-                    hwm_name_suffix = self.source_handler.transfer_dto.table_name
-                hwm_name = "_".join(
-                    [
-                        str(self.source_handler.transfer_dto.id),
-                        self.source_handler.connection_dto.type,
-                        hwm_name_suffix,
-                    ],
-                )
+                hwm_name = self._get_transfer_hwm_name()
                 hwm = hwm_store.get_hwm(hwm_name)
 
                 self.source_handler.hwm = hwm
                 self.target_handler.hwm = hwm
 
                 df = self.source_handler.read()
                 self.target_handler.write(df)
+
+    def _get_transfer_hwm_name(self) -> str:
+        if self.source_handler.connection_dto.type in FILE_CONNECTION_TYPES:
+            hwm_name_suffix = self.source_handler.transfer_dto.directory_path
+        else:
+            hwm_name_suffix = self.source_handler.transfer_dto.table_name
+        hwm_name = "_".join(
+            [
+                str(self.source_handler.transfer_dto.id),
+                self.source_handler.connection_dto.type,
+                hwm_name_suffix,
+            ],
+        )
+        return hwm_name
+
+    def _reset_transfer_hwm(self) -> None:
+        with HorizonHWMStore(
+            api_url=self.settings.hwm_store.url,
+            auth=LoginPassword(login=self.settings.hwm_store.user, password=self.settings.hwm_store.password),
+            namespace=self.settings.hwm_store.namespace,
+        ).force_create_namespace() as hwm_store:
+
+            hwm_name = self._get_transfer_hwm_name()
+            hwm = hwm_store.get_hwm(hwm_name)
+            if hwm:
+                hwm.reset()
+                hwm_store.set_hwm(hwm)

syncmaster/worker/controller.py

@@ -19,18 +19,21 @@ class HWMStoreSettings(BaseModel):
         SYNCMASTER__HWM_STORE__URL=http://horizon:8000
     """
 
+    enabled: bool = Field(
+        description="Enable or disable HWM Store",
+    )
     type: Literal["horizon"] = Field(
-        description=("HWM Store type"),
+        description="HWM Store type",
     )
     url: str = Field(
-        description=("HWM Store URL"),
+        description="HWM Store URL",
     )
     user: str = Field(
-        description=("HWM Store user"),
+        description="HWM Store user",
     )
     password: str = Field(
-        description=("HWM Store password"),
+        description="HWM Store password",
     )
     namespace: str = Field(
-        description=("HWM Store namespace"),
+        description="HWM Store namespace",
     )

syncmaster/worker/settings/hwm_store.py

@@ -124,6 +124,7 @@
     incremental_strategy_by_file_modified_since,
     incremental_strategy_by_file_name,
     incremental_strategy_by_number_column,
+    update_transfer_strategy,
 )
 from tests.test_integration.test_run_transfer.connection_fixtures.webdav_fixtures import (
     prepare_webdav,

tests/test_integration/test_run_transfer/connection_fixtures/__init__.py

@@ -1,4 +1,8 @@
 import pytest
+import pytest_asyncio
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from syncmaster.db.models import Transfer
 
 
 @pytest.fixture
@@ -30,3 +34,14 @@ def incremental_strategy_by_number_column():
         "type": "incremental",
         "increment_by": "NUMBER",
     }
+
+
+@pytest_asyncio.fixture
+async def update_transfer_strategy(session: AsyncSession, request: pytest.FixtureRequest):
+    async def _update_transfer_strategy(transfer: Transfer, strategy_fixture_name: str) -> None:
+        strategy = request.getfixturevalue(strategy_fixture_name)
+        transfer.strategy_params = strategy
+        await session.commit()
+        return transfer
+
+    return _update_transfer_strategy

tests/test_integration/test_run_transfer/connection_fixtures/strategy_fixtures.py

@@ -135,6 +135,49 @@ async def test_run_transfer_postgres_to_clickhouse_with_full_strategy(
     assert df.sort("ID").collect() == init_df.sort("ID").collect()
 
 
+@pytest.mark.parametrize(
+    "strategy, transformations",
+    [
+        (
+            lf("incremental_strategy_by_number_column"),
+            [],
+        ),
+    ],
+)
+async def test_run_transfer_postgres_to_clickhouse_with_different_strategies(
+    client: AsyncClient,
+    group_owner: MockUser,
+    prepare_postgres,
+    prepare_clickhouse,
+    init_df: DataFrame,
+    postgres_to_clickhouse: Transfer,
+    update_transfer_strategy,
+    strategy,
+    transformations,
+):
+    _, fill_with_data = prepare_postgres
+    clickhouse, _ = prepare_clickhouse
+
+    first_transfer_df, second_transfer_df = split_df(df=init_df, ratio=0.6, keep_sorted_by="number")
+    fill_with_data(first_transfer_df)
+    await run_transfer_and_verify(client, group_owner, postgres_to_clickhouse.id)
+
+    fill_with_data(second_transfer_df)
+    await update_transfer_strategy(postgres_to_clickhouse, "full_strategy")
+    await run_transfer_and_verify(client, group_owner, postgres_to_clickhouse.id)
+
+    await update_transfer_strategy(postgres_to_clickhouse, "incremental_strategy_by_number_column")
+    await run_transfer_and_verify(client, group_owner, postgres_to_clickhouse.id)
+
+    reader = DBReader(
+        connection=clickhouse,
+        table=f"{clickhouse.user}.target_table",
+    )
+    df = reader.run()
+    df, init_df = cast_dataframe_types(df, init_df)
+    assert df.sort("ID").collect() == init_df.sort("ID").collect()
+
+
 @pytest.mark.parametrize(
     "strategy, transformations",
     [
@@ -239,11 +282,11 @@ async def test_run_transfer_clickhouse_to_postgres_with_full_strategy(
     prepare_clickhouse,
     prepare_postgres,
     init_df: DataFrame,
+    clickhouse_to_postgres: Transfer,
     source_type,
     strategy,
     transformations,
     expected_filter,
-    clickhouse_to_postgres: Transfer,
 ):
     _, fill_with_data = prepare_clickhouse
     fill_with_data(init_df)
@@ -315,9 +358,9 @@ async def test_run_transfer_clickhouse_to_postgres_with_incremental_strategy(
     prepare_clickhouse,
     prepare_postgres,
     init_df: DataFrame,
+    clickhouse_to_postgres: Transfer,
     strategy,
     transformations,
-    clickhouse_to_postgres: Transfer,
 ):
     _, fill_with_data = prepare_clickhouse
     postgres, _ = prepare_postgres