[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

@@ -4,6 +4,7 @@
 from tempfile import TemporaryDirectory
 from typing import Any
 
+from etl_entities.hwm_store import BaseHWMStore
 from horizon.client.auth import LoginPassword
 from horizon_hwm_store import HorizonHWMStore
 from onetl.strategy import IncrementalStrategy
@@ -198,6 +199,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:
@@ -235,29 +239,43 @@ def get_handler(
             temp_dir=temp_dir,
         )
 
-    def _perform_incremental_transfer(self) -> None:
-        with HorizonHWMStore(
+    def _get_hwm_store(self) -> BaseHWMStore:
+        return 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:
+        ).force_create_namespace()
 
+    def _perform_incremental_transfer(self) -> None:
+        with self._get_hwm_store() 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 self._get_hwm_store() as hwm_store:
+            hwm_name = self._get_transfer_hwm_name()
+            hwm = hwm_store.get_hwm(hwm_name)
+            if hwm and hwm.value:
+                hwm.reset()
+                hwm_store.set_hwm(hwm)

syncmaster/worker/controller.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from typing import Literal
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 
 
 class HWMStoreSettings(BaseModel):
@@ -19,18 +19,37 @@ class HWMStoreSettings(BaseModel):
         SYNCMASTER__HWM_STORE__URL=http://horizon:8000
     """
 
-    type: Literal["horizon"] = Field(
-        description=("HWM Store type"),
+    enabled: bool = Field(
+        default=False,
+        description="Enable or disable HWM Store",
     )
-    url: str = Field(
-        description=("HWM Store URL"),
+    type: Literal["horizon"] | None = Field(
+        default=None,
+        description="HWM Store type",
     )
-    user: str = Field(
-        description=("HWM Store user"),
+    url: str | None = Field(
+        default=None,
+        description="HWM Store URL",
     )
-    password: str = Field(
-        description=("HWM Store password"),
+    user: str | None = Field(
+        default=None,
+        description="HWM Store user",
     )
-    namespace: str = Field(
-        description=("HWM Store namespace"),
+    password: str | None = Field(
+        default=None,
+        description="HWM Store password",
     )
+    namespace: str | None = Field(
+        default=None,
+        description="HWM Store namespace",
+    )
+
+    @model_validator(mode="after")
+    def check_required_fields_if_enabled(self):
+        if self.enabled:
+            missing_fields = [
+                field for field in ["type", "url", "user", "password", "namespace"] if getattr(self, field) is None
+            ]
+            if missing_fields:
+                raise ValueError(f"All fields must be set with enabled HWMStore. Missing {', '.join(missing_fields)}")
+        return self

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,56 @@ 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,
+):
+    """
+    Run three transfers:
+
+    1. Incremental transfer
+    2. Full transfer with HWM reset
+    3. Incremental transfer to verify that results match the snapshot.
+    """
+    _, 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 +289,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 +365,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