Add depend strategy for reflections

Author: bcmeireles

dbt/adapters/dremio/api/rest/client.py

@@ -207,7 +207,7 @@ def delete_tags(self, dataset_id: str, version: str):
 
 
     def get_reflections(self, dataset_id):
-        url = UrlBuilder.get_reflection_url(self._parameters, dataset_id)
+        url = UrlBuilder.get_reflections_from_dataset_url(self._parameters, dataset_id)
         return _get(
             url,
             self._parameters.authentication.get_headers(),
@@ -231,4 +231,11 @@ def update_reflection(self, reflection_id, payload):
             json=payload,
             ssl_verify=self._parameters.authentication.verify_ssl,
         )
-    
+
+    def get_reflection(self, reflection_id):
+        url = UrlBuilder.get_reflection_url(self._parameters, reflection_id)
+        return _get(
+            url,
+            self._parameters.authentication.get_headers(),
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )

dbt/adapters/dremio/api/rest/url_builder.py

@@ -185,7 +185,7 @@ def update_reflection_url(cls, parameters: Parameters, dataset_id):
         return url_path + endpoint
 
     @classmethod
-    def get_reflection_url(cls, parameters: Parameters, dataset_id):
+    def get_reflections_from_dataset_url(cls, parameters: Parameters, dataset_id):
         url_path = parameters.base_url
         if type(parameters) is CloudParameters:
             url_path += UrlBuilder.CLOUD_DATASET_ENDPOINT.format(
@@ -196,3 +196,16 @@ def get_reflection_url(cls, parameters: Parameters, dataset_id):
 
         endpoint = "/{}/reflection".format(dataset_id)
         return url_path + endpoint
+
+    @classmethod
+    def get_reflection_url(cls, parameters: Parameters, reflection_id):
+        url_path = parameters.base_url
+        if type(parameters) is CloudParameters:
+            url_path += UrlBuilder.CLOUD_REFLECTIONS_ENDPOINT.format(
+                parameters.cloud_project_id
+            )
+        else:
+            url_path += UrlBuilder.SOFTWARE_REFLECTIONS_ENDPOINT
+
+        endpoint = "/{}".format(reflection_id)
+        return url_path + endpoint

dbt/adapters/dremio/connections.py

@@ -326,7 +326,7 @@ def create_reflection(self, name: str, reflection_type: str, anchor: DremioRelat
                           date_dimensions: List[str], measures: List[str],
                           computations: List[str], partition_by: List[str], partition_transform: List[str],
                           partition_method: str, distribute_by: List[str], localsort_by: List[str],
-                          arrow_cache: bool) -> None:
+                          arrow_cache: bool, reflection_strategy: str, max_wait_time: int, check_interval: int) -> None:
         thread_connection = self.get_thread_connection()
         connection = self.open(thread_connection)
         rest_client = connection.handle.get_client()
@@ -357,13 +357,41 @@ def create_reflection(self, name: str, reflection_type: str, anchor: DremioRelat
             if reflection.get("name") == name:
                 logger.debug(f"Reflection {name} already exists. Updating it")
                 payload["tag"] = reflection.get("tag")
-                rest_client.update_reflection(reflection.get("id"), payload)
+                created_reflection = rest_client.update_reflection(reflection.get("id"), payload)
                 updated = True
                 break
 
         if not updated:
             logger.debug(f"Reflection {name} does not exist. Creating it")
-            rest_client.create_reflection(payload)
+            created_reflection = rest_client.create_reflection(payload)
+
+        if reflection_strategy == "wait":
+            reflection_id = created_reflection["id"]
+            start_time = time.time()
+            end_time = start_time + max_wait_time
+
+            while time.time() < end_time:
+                reflection_info = rest_client.get_reflection(reflection_id)
+                status = reflection_info["status"]["availability"]
+
+                if status == "AVAILABLE":
+                    return
+
+                time.sleep(check_interval)
+
+            logger.info(f"Reflection {name} did not become available within {max_wait_time} seconds, skipping wait")
+
+        elif reflection_strategy == "depend":
+            reflection_id = created_reflection["id"]
+
+            while True:
+                reflection_info = rest_client.get_reflection(reflection_id)
+                status = reflection_info["status"]["availability"]
+
+                if status == "AVAILABLE":
+                    return
+
+                time.sleep(check_interval)
 
     def _make_new_space_json(self, name) -> json:
         python_dict = {"entityType": "space", "name": name}

dbt/adapters/dremio/impl.py

@@ -207,10 +207,10 @@ def process_tags(self, relation: DremioRelation, tags: list[str]) -> None:
     def create_reflection(self, name: str, type: str, anchor: DremioRelation, display: List[str], dimensions: List[str],
                           date_dimensions: List[str], measures: List[str], computations: List[str],
                           partition_by: List[str], partition_transform: List[str], partition_method: str,
-                          distribute_by: List[str], localsort_by: List[str], arrow_cache: bool) -> None:
+                          distribute_by: List[str], localsort_by: List[str], arrow_cache: bool, reflection_strategy: str, max_wait_time: int, check_interval: int) -> None:
         self.connections.create_reflection(name, type, anchor, display, dimensions, date_dimensions, measures,
                                            computations, partition_by, partition_transform, partition_method,
-                                           distribute_by, localsort_by, arrow_cache)
+                                           distribute_by, localsort_by, arrow_cache, reflection_strategy, max_wait_time, check_interval)
 
 
 COLUMNS_EQUAL_SQL = """

dbt/include/dremio/macros/materializations/reflection/create_reflection.sql

@@ -38,7 +38,7 @@ USING target
 #}
 
 {%- macro create_reflection(reflection_name, reflection_type, anchor,
-  display=none, dimensions=none, date_dimensions=none, measures=none, computations=none, partition_by=none, partition_transform=none, partition_method=none, distribute_by=none, localsort_by=none, arrow_cache=false) %}
+  display=none, dimensions=none, date_dimensions=none, measures=none, computations=none, partition_by=none, partition_transform=none, partition_method=none, distribute_by=none, localsort_by=none, arrow_cache=false, reflection_strategy=reflection_strategy, max_wait_time=max_wait_time, check_interval=check_interval) %}
 
   {%- if reflection_type == 'raw' %}
     {% set reflection_type = 'RAW' %}
@@ -48,7 +48,7 @@ USING target
     {% do exceptions.CompilationError("invalid reflection type") %}
   {%- endif -%}
 
-  {% do adapter.create_reflection(reflection_name, reflection_type, anchor, display, dimensions, date_dimensions, measures, computations, partition_by, partition_transform, partition_method, distribute_by, localsort_by, arrow_cache) %}
+  {% do adapter.create_reflection(reflection_name, reflection_type, anchor, display, dimensions, date_dimensions, measures, computations, partition_by, partition_transform, partition_method, distribute_by, localsort_by, arrow_cache, reflection_strategy, max_wait_time, check_interval) %}
 
   SELECT 1
 {% endmacro -%}

dbt/include/dremio/macros/materializations/reflection/reflection.sql

@@ -13,6 +13,11 @@ See the License for the specific language governing permissions and
 limitations under the License.*/
 
 {% materialization reflection, adapter='dremio' %}
+  {%- set reflection_strategy = config.get('reflection_strategy', validator=validation.any[basestring]) or 'trigger' -%}
+
+  {% if reflection_strategy not in ['depend', 'trigger'] %}
+    {% do exceptions.CompilationError("Invalid reflection strategy. Valid strategies are 'trigger', 'wait' or 'depend'") %}
+  {%- endif -%}
 
   {% set reflection_name = config.get('name', validator=validation.any[basetring]) or config.get('alias', validator=validation.any[basetring]) or model.name %}
   {% set raw_reflection_type = config.get('reflection_type', validator=validation.any[basestring]) or 'raw' %}
@@ -30,6 +35,8 @@ limitations under the License.*/
   {%- set distribute_by = config.get('distribute_by', validator=validation.any[basestring])-%}
   {%- set localsort_by = config.get('localsort_by', validator=validation.any[basestring]) -%}
   {%- set arrow_cache = config.get('arrow_cache') -%}
+  {%- set max_wait_time = config.get('max_wait_time', validator=validation.any[int]) or 30 -%}
+  {%- set check_interval = config.get('check_interval', validator=validation.any[int]) or 5 -%}
 
   {% set relation = this %}
 
@@ -115,7 +122,7 @@ limitations under the License.*/
   -- build model
   {% call statement('main') -%}
     {{ create_reflection(reflection_name, reflection_type, anchor,
-      display=display, dimensions=dimensions, date_dimensions=date_dimensions, measures=measures, computations=computations, partition_by=partition_by, partition_transform=partition_transform, partition_method=partition_method, distribute_by=distribute_by, localsort_by=localsort_by, arrow_cache=arrow_cache) }}
+      display=display, dimensions=dimensions, date_dimensions=date_dimensions, measures=measures, computations=computations, partition_by=partition_by, partition_transform=partition_transform, partition_method=partition_method, distribute_by=distribute_by, localsort_by=localsort_by, arrow_cache=arrow_cache, reflection_strategy=reflection_strategy, max_wait_time=max_wait_time, check_interval=check_interval) }}
   {%- endcall %}
 
   {{ run_hooks(post_hooks) }}

tests/functional/adapter/dremio_specific/test_reflections.py

@@ -3,9 +3,11 @@
 from dbt.adapters.dremio import DremioCredentials
 from dbt.adapters.dremio.api.parameters import ParametersBuilder
 from dbt.adapters.dremio.api.authentication import DremioAuthentication
-from dbt.tests.util import run_dbt
+from dbt.tests.util import run_dbt, run_dbt_and_capture
 from pydantic.experimental.pipeline import transform
 
+from tests.utils.util import BUCKET
+
 view1_model = """
 SELECT IncidntNum, Category, Descript, DayOfWeek, TO_DATE("SF_incidents2016.json"."Date", 'YYYY-MM-DD', 1) AS "Date", "SF_incidents2016.json"."Time" AS "Time", PdDistrict, Resolution, Address, X, Y, Location, PdId
 FROM Samples."samples.dremio.com"."SF_incidents2016.json" AS "SF_incidents2016.json";
@@ -141,6 +143,25 @@
 -- depends_on: {{ ref('view1') }}
 """
 
+depend_strategy_timeout_reflection = """
+{{ config(alias='This will mock a timeout when using depend',
+            materialized='reflection',
+            reflection_strategy='depend',
+            max_wait_time=1,
+            display=['Date', 'DayOfWeek', 'PdDistrict', 'Category'],
+            reflection_type='raw')}}
+-- depends_on: {{ ref('view1') }}
+"""
+
+trigger_strategy_timeout_reflection = """
+{{ config(alias='This will mock a timeout when using trigger',
+            materialized='reflection',
+            reflection_strategy='trigger',
+            max_wait_time=1,
+            display=['Date', 'DayOfWeek', 'PdDistrict', 'Category'],
+            reflection_type='raw')}}
+-- depends_on: {{ ref('view1') }}
+"""
 
 class TestReflectionsDremio:
     @pytest.fixture(scope="class")
@@ -170,6 +191,8 @@ def models(self):
             "default_displays_model.sql": default_displays_model,
             "name_reflection_from_alias.sql": name_reflection_from_alias_model,
             "name_reflection_from_filename.sql": name_reflection_from_filename_model,
+            "depend_strategy_timeout_reflection.sql": depend_strategy_timeout_reflection,
+            "trigger_strategy_timeout_reflection.sql": trigger_strategy_timeout_reflection,
         }
 
     def _create_path_list(self, database, schema):
@@ -407,3 +430,11 @@ def testNameReflectionFromFilename(self, project, client):
         assert "partitionFields" not in reflection
         assert "sortFields" not in reflection
         assert reflection["partitionDistributionStrategy"] == "STRIPED"
+
+    def testDependStrategyTimeoutReflection(self, project):
+        (results, log_output) = run_dbt_and_capture(["run", "--select", "view1", "depend_strategy_timeout_reflection"])
+        assert "did not become available within 1 seconds, skipping wait" in log_output
+
+    def testTriggerStrategyTimeoutReflection(self,    project):
+        (results, log_output) = run_dbt_and_capture(["run", "--select", "view1", "trigger_strategy_timeout_reflection"])
+        assert "did not become available within 1 seconds, skipping wait" not in log_output