Issue 33: Migrate stream clear to the data plane (#37)

* issue-33: Add data plane client to Decodable adapter

* issue-33: Migrate stream clear to use the data plane

* issue-33: Formatting

* issue-33: Pyright

Author: nicoweidner

dbt/adapters/decodable/impl.py

@@ -19,10 +19,11 @@
 
 from agate.table import Table as AgateTable
 from dbt.adapters.base import BaseAdapter, BaseRelation, Column
+from dbt.adapters.base.meta import available
+from dbt.adapters.protocol import AdapterConfig
 from dbt.contracts.connection import Connection
 from dbt.contracts.graph.manifest import Manifest
 from dbt.contracts.graph.parsed import ParsedNode
-from dbt.adapters.base.meta import available
 from dbt.contracts.relation import RelationType
 from dbt.events import AdapterLogger
 from dbt.exceptions import (
@@ -38,8 +39,11 @@
 )
 from dbt.adapters.decodable.handler import DecodableHandler
 from dbt.adapters.decodable.relation import DecodableRelation
-from dbt.adapters.protocol import AdapterConfig
-from decodable.client.client import DecodableControlPlaneApiClient, SchemaField
+from decodable.client.client import (
+    DecodableControlPlaneApiClient,
+    SchemaField,
+    DecodableDataPlaneApiClient,
+)
 from decodable.client.types import (
     FieldType,
     PrimaryKey,
@@ -189,7 +193,7 @@ def drop_relation(self, relation: BaseRelation) -> None:
         if relation.identifier is None:
             return
 
-        client = self._client()
+        client = self._control_plane_client()
 
         self.logger.debug(f"Dropping pipeline '{relation}'...")
 
@@ -243,23 +247,25 @@ def truncate_relation(self, relation: BaseRelation) -> None:
         if not relation.identifier:
             raise_compiler_error("Cannot truncate an unnamed relation")
 
-        client = self._client()
-        stream_id = client.get_stream_id(relation.render())
+        control_plane_client = self._control_plane_client()
+        data_plane_client = self._data_plane_client()
+        stream_id = control_plane_client.get_stream_id(relation.render())
 
         if not stream_id:
             raise_database_error(
                 f"Error clearing stream `{relation.render()}`: stream doesn't exist"
             )
 
-        client.clear_stream(stream_id)
+        clear_token_response = control_plane_client.get_clear_stream_token(stream_id)
+        data_plane_client.clear_stream(stream_id, clear_token_response.token)
 
     @available.parse_none
     def rename_relation(self, from_relation: BaseRelation, to_relation: BaseRelation) -> None:
         """Rename the relation from from_relation to to_relation.
 
         Implementors must call self.cache.rename() to preserve cache state.
         """
-        client = self._client()
+        client = self._control_plane_client()
         self.cache_renamed(from_relation, to_relation)
 
         if not from_relation.identifier:
@@ -345,7 +351,7 @@ def expand_column_types(self, goal: BaseRelation, current: BaseRelation) -> None
     def list_relations_without_caching(self, schema_relation: BaseRelation) -> List[BaseRelation]:
         relations: List[BaseRelation] = []
 
-        stream_list: List[Dict[str, Any]] = self._client().list_streams().items
+        stream_list: List[Dict[str, Any]] = self._control_plane_client().list_streams().items
         for stream in stream_list:
             relations.append(
                 self.Relation.create(
@@ -367,7 +373,9 @@ def get_columns_in_relation(
         if not relation.identifier:
             return []
 
-        stream_info = self._client().get_stream_information(stream_id=relation.render())
+        stream_info = self._control_plane_client().get_stream_information(
+            stream_id=relation.render()
+        )
 
         for schema_column in stream_info["schema"]:
             columns.append(
@@ -384,7 +392,7 @@ def has_changed(
         watermark: Optional[str] = None,
         primary_key: Optional[str] = None,
     ) -> bool:
-        client = self._client()
+        client = self._control_plane_client()
 
         new_pipe_sql = self._wrap_as_pipeline(relation.render(), sql)
         fields: List[Dict[str, str]] = client.get_stream_from_sql(new_pipe_sql)["schema"]
@@ -452,7 +460,7 @@ def create_table(
         if not model:
             self.logger.debug(f"Model {relation.render()} not found in dbt graph")
 
-        client = self._client()
+        client = self._control_plane_client()
 
         fields: List[Dict[str, str]] = client.get_stream_from_sql(
             self._wrap_as_pipeline(relation.render(), sql)
@@ -537,7 +545,7 @@ def create_seed_table(
 
             schema.append(SchemaField(name=col_name, type=field_type))
 
-        client = self._client()
+        client = self._control_plane_client()
 
         self.logger.debug(f"Creating connection and stream for seed `{table_name}`...")
         response = client.create_connection(name=table_name, schema=schema)
@@ -550,7 +558,7 @@ def create_seed_table(
     @available
     def send_seed_as_events(self, seed_name: str, data: AgateTable):
         self.logger.debug(f"Sending data to connection `{seed_name}`")
-        client = self._client()
+        client = self._control_plane_client()
 
         conn_id = client.get_connection_id(seed_name)
         if not conn_id:
@@ -576,7 +584,7 @@ def send_seed_as_events(self, seed_name: str, data: AgateTable):
 
     @available
     def reactivate_connection(self, connection: Relation):
-        client = self._client()
+        client = self._control_plane_client()
 
         conn_id = client.get_connection_id(connection.render())
         if not conn_id:
@@ -586,7 +594,7 @@ def reactivate_connection(self, connection: Relation):
 
     @available
     def stop_pipeline(self, pipe: Relation):
-        client = self._client()
+        client = self._control_plane_client()
 
         pipe_id = client.get_pipeline_id(pipe.render())
         if not pipe_id:
@@ -596,7 +604,7 @@ def stop_pipeline(self, pipe: Relation):
 
     @available
     def delete_pipeline(self, pipe: Relation):
-        client = self._client()
+        client = self._control_plane_client()
 
         pipeline_id = client.get_pipeline_id(pipe.render())
         if pipeline_id:
@@ -607,7 +615,7 @@ def delete_pipeline(self, pipe: Relation):
 
     @available
     def delete_stream(self, stream: Relation, skip_errors: bool = False):
-        client = self._client()
+        client = self._control_plane_client()
 
         stream_id = client.get_stream_id(stream.render())
         if not stream_id:
@@ -623,7 +631,7 @@ def delete_stream(self, stream: Relation, skip_errors: bool = False):
 
     @available
     def delete_connection(self, conn: Relation):
-        client = self._client()
+        client = self._control_plane_client()
 
         conn_id = client.get_connection_id(conn.render())
         if not conn_id:
@@ -645,12 +653,18 @@ def should_materialize_tests(self) -> bool:
             return False
         return credentials.materialize_tests
 
-    def _client(self) -> DecodableControlPlaneApiClient:
+    def _control_plane_client(self) -> DecodableControlPlaneApiClient:
         handle: DecodableHandler = (
             self.get_thread_connection().handle
         )  # pyright: ignore [reportGeneralTypeIssues]
         return handle.control_plane_client
 
+    def _data_plane_client(self) -> DecodableDataPlaneApiClient:
+        handle: DecodableHandler = (
+            self.get_thread_connection().handle
+        )  # pyright: ignore [reportGeneralTypeIssues]
+        return handle.data_plane_client
+
     @classmethod
     def _get_model_schema_hints(cls, model: ParsedNode) -> Set[SchemaField]:
         schema: Set[SchemaField] = set()

decodable/client/client.py

@@ -97,6 +97,22 @@ def from_dict(cls, response: Dict[str, Any]) -> PreviewTokensResponse:
         )
 
 
+@dataclass
+class DataPlaneTokenResponse:
+    data_plane_request: Optional[str]
+    token: str
+
+    @classmethod
+    def from_dict(cls, response: Dict[str, Any]) -> DataPlaneTokenResponse:
+        return cls(
+            # Some requests don't include a body, so check whether it's present
+            data_plane_request=response["data_plane_request"]
+            if "data_plane_request" in response
+            else None,
+            token=response["token"],
+        )
+
+
 class DecodableAPIException(Exception):
     @classmethod
     def category(cls) -> str:
@@ -169,6 +185,11 @@ def get_preview(self, auth_token: str, next_token: str) -> PreviewResponse:
         )
         return PreviewResponse.from_dict(response.json())
 
+    def clear_stream(self, stream_id: str, token: str) -> None:
+        self._post_api_request(
+            bearer_token=token, endpoint_url=f"{self.config.api_url}/streams/{stream_id}/clear"
+        )
+
     def _get_api_request(
         self,
         bearer_token: str,
@@ -291,10 +312,12 @@ def delete_stream(self, stream_id: str) -> None:
             endpoint_url=f"{self.config.decodable_api_url()}/streams/{stream_id}"
         )
 
-    def clear_stream(self, stream_id: str) -> None:
-        self._post_api_request(
-            payload={}, endpoint_url=f"{self.config.decodable_api_url()}/streams/{stream_id}/clear"
+    def get_clear_stream_token(self, stream_id: str) -> DataPlaneTokenResponse:
+        response = self._post_api_request(
+            payload={},
+            endpoint_url=f"{self.config.decodable_api_url()}/streams/{stream_id}/clear/token",
         )
+        return DataPlaneTokenResponse.from_dict(response.json())
 
     def list_pipelines(self) -> ApiResponse:
         response = self._get_api_request(