Chore: Refactor new "Writers" interface for destinations and caches

airbyte/_future_cdk/catalog_providers.py

@@ -8,13 +8,15 @@
 
 from __future__ import annotations
 
+import copy
 from typing import TYPE_CHECKING, Any, final
 
 from airbyte_protocol.models import (
     ConfiguredAirbyteCatalog,
 )
 
 from airbyte import exceptions as exc
+from airbyte.strategies import WriteMethod, WriteStrategy
 
 
 if TYPE_CHECKING:
@@ -135,3 +137,70 @@ def from_read_result(
                 ]
             )
         )
+
+    def get_primary_keys(
+        self,
+        stream_name: str,
+    ) -> list[str]:
+        pks = self.get_configured_stream_info(stream_name).primary_key
+        if not pks:
+            return []
+
+        joined_pks = [".".join(pk) for pk in pks]
+        for pk in joined_pks:
+            if "." in pk:
+                msg = f"Nested primary keys are not yet supported. Found: {pk}"
+                raise NotImplementedError(msg)
+
+        return joined_pks
+
+    def get_cursor_key(
+        self,
+        stream_name: str,
+    ) -> str | None:
+        return self.get_configured_stream_info(stream_name).cursor_field
+
+    def resolve_write_method(
+        self,
+        stream_name: str,
+        write_strategy: WriteStrategy,
+    ) -> WriteMethod:
+        """Return the write method for the given stream."""
+        has_pks: bool = bool(self.get_primary_keys(stream_name))
+        has_incremental_key: bool = bool(self.get_cursor_key(stream_name))
+        if write_strategy == WriteStrategy.MERGE and not has_pks:
+            raise exc.PyAirbyteInputError(
+                message="Cannot use merge strategy on a stream with no primary keys.",
+                context={
+                    "stream_name": stream_name,
+                },
+            )
+
+        if write_strategy != WriteStrategy.AUTO:
+            return WriteMethod(write_strategy)
+
+        if has_pks:
+            return WriteMethod.MERGE
+
+        if has_incremental_key:
+            return WriteMethod.APPEND
+
+        return WriteMethod.REPLACE
+
+    def with_write_strategy(
+        self,
+        write_strategy: WriteStrategy,
+    ) -> CatalogProvider:
+        """Return a new catalog provider with the specified write strategy applied.
+
+        The original catalog provider is not modified.
+        """
+        new_catalog: ConfiguredAirbyteCatalog = copy.deepcopy(self.configured_catalog)
+        for stream in new_catalog.streams:
+            write_method = self.resolve_write_method(
+                stream_name=stream.stream.name,
+                write_strategy=write_strategy,
+            )
+            stream.destination_sync_mode = write_method.destination_sync_mode
+
+        return CatalogProvider(new_catalog)