Move Reflections to Rest API Calls and Extend Implementation (dremio#256)

### Summary

Reflections were being created through the SQL Runner and that was the
cause for some issues we were facing.

### Description

The current reflections implementation requires sys.reflections access,
which is only available to superusers (admins). When creating the
reflection with a regular user, it will error. To fix this issue, the
creation of the reflections has been moved from queries being ran inside
the SQL Runner to Dremio's Rest API calls.
The current logic of writing reflections through a custom
materialization remains to make sure reflections created under the
previous approach still work.

On top of this, the behaviour where all reflections would be dropped and
then recreated when running a dbt project has been changed to instead
first check if a reflection with the same name is already present in the
same dataset and, if that's the case, it will update the already
existing reflection instead.

### Test Results

A total of 11 tests have been added, approaching different situations
that may occur when dealing with reflections

### Changelog

- Reflections are now handled through the Rest API, allowing non-admin
users to also create them
- It is now possible to set a custom name for reflections
- If a reflection already exists in the dataset with the same name
defined in the model, it will be updated instead of creating a new one
- New `date_dimensions` parameter was added to the reflection
materialization, to set fields that have a `DATE` granularity
- Added distribution fields under `distribute_by`
- Added partition transformations under `partition_transform`
  - Defaults to Original/Identity if not defined
  - year/month/day/hour/bucket(n), truncate(n)
- Computations default to `SUM, COUNT` if mapped measure is numeric,
`COUNT` if not
- `reflections_enabled` adapter option has been renamed to
`reflections_metadata_enabled` (requires user privileges to run in
dremio)
- CI now creates and formats a `Samples` source as it is needed for the
reflection tests

Author: bcmeireles

.github/scripts/create_dremio_s3_source.sh renamed to .github/scripts/create_and_format_sources.sh

@@ -50,4 +50,28 @@ curl -s -X PUT "http://localhost:9047/apiv2/source/dbt_test_source" \
   -H "Authorization: _dremio$AUTH_TOKEN" \
   --data "{\"name\":\"dbt_test_source\",\"config\":{\"credentialType\":\"ACCESS_KEY\",\"accessKey\":\"$MINIO_ROOT_USER\",\"accessSecret\":\"$MINIO_ROOT_PASSWORD\",\"secure\":false,\"externalBucketList\":[],\"enableAsync\":true,\"enableFileStatusCheck\":true,\"rootPath\":\"/\",\"defaultCtasFormat\":\"ICEBERG\",\"propertyList\":[{\"name\":\"fs.s3a.path.style.access\",\"value\":\"true\"},{\"name\":\"fs.s3a.endpoint\",\"value\":\"minio:9000\"},{\"name\":\"dremio.s3.compat\",\"value\":\"true\"}],\"whitelistedBuckets\":[],\"isCachingEnabled\":false,\"maxCacheSpacePct\":100},\"type\":\"S3\",\"metadataPolicy\":{\"deleteUnavailableDatasets\":true,\"autoPromoteDatasets\":false,\"namesRefreshMillis\":3600000,\"datasetDefinitionRefreshAfterMillis\":3600000,\"datasetDefinitionExpireAfterMillis\":10800000,\"authTTLMillis\":86400000,\"updateMode\":\"PREFETCH_QUERIED\"}}"
 
-echo "S3 Source created in Dremio."
+echo "S3 Source created in Dremio."
+
+echo "Creating the Samples source in Dremio..."
+curl -s -X PUT  "http://localhost:9047/apiv2/source/Samples" \
+  -H "Content-Type: application/json" \
+  -H "Authorization: _dremio$AUTH_TOKEN" \
+  --data-raw "{\"name\":\"Samples\",\"config\":{\"externalBucketList\":[\"samples.dremio.com\"],\"credentialType\":\"NONE\",\"secure\":false,\"propertyList\":[]},\"name\":\"Samples\",\"accelerationRefreshPeriod\":3600000,\"accelerationGracePeriod\":10800000,\"accelerationNeverRefresh\":true,\"accelerationNeverExpire\":true,\"accelerationActivePolicyType\":\"PERIOD\",\"accelerationRefreshSchedule\":\"0 0 8 * * *\",\"accelerationRefreshOnDataChanges\":false,\"type\":\"S3\"}"
+
+echo "Samples source created in Dremio."
+
+echo "Formatting SF_incidents2016..."
+curl -s -X PUT "http://localhost:9047/apiv2/source/Samples/file_format/samples.dremio.com/SF_incidents2016.json" \
+  -H "Content-Type: application/json" \
+  -H "Authorization: _dremio$AUTH_TOKEN" \
+  --data-raw "{\"type\":\"JSON\"}"
+
+echo "SF_incidents2016 formatted in Dremio."
+
+echo "Formatting NYC-taxi-trips..."
+curl -s -X PUT "http://localhost:9047/apiv2/source/Samples/folder_format/samples.dremio.com/NYC-taxi-trips" \
+  -H "Content-Type: application/json" \
+  -H "Authorization: _dremio$AUTH_TOKEN" \
+  --data-raw "{\"ignoreOtherFileFormats\":false,\"type\":\"Parquet\"}"
+
+echo "NYC-taxi-trips formatted in Dremio."

.github/workflows/ci.yml

@@ -48,8 +48,8 @@ jobs:
       - name: Create MinIO bucket
         run: bash .github/scripts/create_minio_bucket.sh
 
-      - name: Create Dremio S3 Source
-        run: bash .github/scripts/create_dremio_s3_source.sh
+      - name: Create and Format Sources
+        run: bash .github/scripts/create_and_format_sources.sh
 
       - name: Install Dependencies
         uses: actions/setup-python@v4

CHANGELOG.md

@@ -3,8 +3,18 @@
 ## Changes
 
 -   Added [DremioRestClient](dbt/adapters/dremio/api/rest/client.py) to isolate all Dremio API calls inside one class
+- [#256](https://github.com/dremio/dbt-dremio/pull/256) Reflections are now handled through the Rest API
+  - Non-admin users are now able to use reflections
+  - It is now possible to set a custom name for reflections
+  - If a reflection already exists in the dataset with the same name defined in the model, it will be updated instead of creating a new one
+  - New `date_dimensions` parameter was added to the reflection materialization, to set fields that have a `DATE` granularity
+  - Added Distribution Fields under `distribute_by`
+  - Added partition transformations under `partition_transform`
+    - Defaults to Original/Identity if not defined
+    - `year/month/day/hour/bucket(n)/truncate(n)`
+  - Computations default to `SUM, COUNT` if mapped measure is numeric, `COUNT` if not
+  - `reflections_enabled` adapter option has been renamed to `reflections_metadata_enabled` (requires user privileges to run in dremio)
 -   Removing duplicated macros array_append, array_concat as Dremio already has SQL functions analogues.
-
 ## Dependency
 
 -   [#222](https://github.com/dremio/dbt-dremio/issues/222) Upgrade dbt-core to 1.8.8 and dbt-tests-adapter to 1.8.0

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

@@ -15,12 +15,11 @@
 # limitations under the License.
 
 
-
 import requests
 
 from dbt.adapters.dremio.api.authentication import DremioPatAuthentication
 from dbt.adapters.dremio.api.parameters import Parameters
-from dbt.adapters.dremio.api.rest.utils import _post, _get, _delete
+from dbt.adapters.dremio.api.rest.utils import _post, _get, _put, _delete
 from dbt.adapters.dremio.api.rest.url_builder import UrlBuilder
 
 from dbt.adapters.events.logging import AdapterLogger
@@ -132,4 +131,30 @@ def delete_catalog(self, cid):
             url,
             self._parameters.authentication.get_headers(),
             ssl_verify=self._parameters.authentication.verify_ssl,
-        )
+        )
+
+    def get_reflections(self, dataset_id):
+        url = UrlBuilder.get_reflection_url(self._parameters, dataset_id)
+        return _get(
+            url,
+            self._parameters.authentication.get_headers(),
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )
+
+    def create_reflection(self, payload):
+        url = UrlBuilder.create_reflection_url(self._parameters)
+        return _post(
+            url,
+            self._parameters.authentication.get_headers(),
+            json=payload,
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )
+
+    def update_reflection(self, reflection_id, payload):
+        url = UrlBuilder.update_reflection_url(self._parameters, reflection_id)
+        return _put(
+            url,
+            self._parameters.authentication.get_headers(),
+            json=payload,
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )

dbt/adapters/dremio/api/rest/entities/__init__.py

@@ -0,0 +1,141 @@
+from enum import Enum
+
+
+class TransformType(Enum):
+    YEAR = "YEAR"
+    MONTH = "MONTH"
+    DAY = "DAY"
+    HOUR = "HOUR"
+    IDENTITY = "IDENTITY"
+    BUCKET = "BUCKET"
+    TRUNCATE = "TRUNCATE"
+
+    @classmethod
+    def from_string(cls, transform_str):
+        transform_str = transform_str.upper()
+
+        if transform_str.startswith("BUCKET("):
+            return cls.BUCKET
+        elif transform_str.startswith("TRUNCATE("):
+            return cls.TRUNCATE
+
+        try:
+            return cls(transform_str)
+        except ValueError:
+            return cls.IDENTITY
+
+    def to_transform(self, raw_str):
+        if self in (
+                TransformType.YEAR,
+                TransformType.MONTH,
+                TransformType.DAY,
+                TransformType.HOUR,
+                TransformType.IDENTITY
+        ):
+            return {"type": self.value}
+
+        if self == TransformType.BUCKET:
+            bucket_count = int(raw_str.split("(")[1].split(")")[0])
+            return {
+                "type": "BUCKET",
+                "bucketTransform": {"bucketCount": bucket_count},
+            }
+
+        if self == TransformType.TRUNCATE:
+            truncate_length = int(raw_str.split("(")[1].split(")")[0])
+            return {
+                "type": "TRUNCATE",
+                "truncateTransform": {"truncateLength": truncate_length},
+            }
+
+        return {"type": TransformType.IDENTITY.value}
+
+
+# https://docs.dremio.com/current/reference/api/reflections/
+class ReflectionEntity:
+    def __init__(self, name, reflection_type, dataset_id, display_fields, dimensions, date_dimensions, measures,
+                 computations, partition_by, partition_transform, partition_method, distribute_by, localsort_by,
+                 arrow_cache):
+        self.__name = name
+        self.__type = reflection_type
+        self.__dataset_id = dataset_id
+        self.__partition_method = partition_method.upper()
+        self.__display_fields = display_fields
+        self.__dimensions_fields = dimensions
+        self.__date_dimensions_fields = date_dimensions
+        self.__measure_fields = measures
+        self.__computation_fields = computations
+        self.__partition_by_fields = partition_by
+        self.__partition_transformations = partition_transform
+        self.__partition_method = partition_method
+        self.__distribute_by_fields = distribute_by
+        self.__local_sort_fields = localsort_by
+        self.__arrow_cache = arrow_cache
+
+    def buildDisplayFields(self):
+        return [{"name": field} for field in self.__display_fields]
+
+    def buildDimensionFields(self):
+        return [{"name": field} for field in self.__dimensions_fields]
+
+    def buildDateFields(self):
+        return [{"name": date_dimension, "granularity": "DATE"} for date_dimension in self.__date_dimensions_fields]
+
+    def buildMeasureFields(self):
+        return [{"name": measure, "measureTypeList": computation.split(',')} for
+                measure, computation in zip(self.__measure_fields, self.__computation_fields)]
+
+    def buildPartitionFields(self):
+        if not self.__partition_transformations:
+            self.__partition_transformations = ["IDENTITY"] * len(self.__partition_by_fields)
+
+        partition_fields = []
+
+        for partition, transform in zip(self.__partition_by_fields, self.__partition_transformations):
+            transform_type = TransformType.from_string(transform)
+            partition_fields.append({
+                "name": partition,
+                "transform": transform_type.to_transform(transform)
+            })
+
+        return partition_fields
+
+    def buildDistributionFields(self):
+        return [{"name": distribute} for distribute in self.__distribute_by_fields]
+
+    def buildSortFields(self):
+        return [{"name": sort} for sort in self.__local_sort_fields]
+
+    def build_payload(self):
+        payload = {
+            "type": self.__type,
+            "name": self.__name,
+            "datasetId": self.__dataset_id,
+            "enabled": True,
+            "arrowCachingEnabled": self.__arrow_cache,
+            "partitionDistributionStrategy": self.__partition_method.upper(),
+            "entityType": "reflection"
+        }
+
+        if self.__display_fields:
+            payload["displayFields"] = self.buildDisplayFields()
+
+        if self.__dimensions_fields:
+            payload["dimensionFields"] = self.buildDimensionFields()
+
+        if self.__date_dimensions_fields:
+            payload["dateFields"] = self.buildDateFields()
+
+        if self.__measure_fields and self.__computation_fields:
+            payload["measureFields"] = self.buildMeasureFields()
+
+        if self.__partition_by_fields:
+            payload["partitionFields"] = self.buildPartitionFields()
+
+        if self.__distribute_by_fields:
+            payload["distributionFields"] = self.buildDistributionFields()
+
+        if self.__local_sort_fields:
+            payload["sortFields"] = self.buildSortFields()
+
+        return payload

dbt/adapters/dremio/api/rest/entities/reflection.py

@@ -34,6 +34,12 @@ class UrlBuilder:
     SOFTWARE_CATALOG_ENDPOINT = "/api/v3/catalog"
     CLOUD_CATALOG_ENDPOINT = CLOUD_PROJECT_ENDPOINT + "/{}/catalog"
 
+    SOFTWARE_REFLECTIONS_ENDPOINT = "/api/v3/reflection"
+    CLOUD_REFLECTIONS_ENDPOINT = CLOUD_PROJECT_ENDPOINT + "/{}/reflection"
+
+    SOFTWARE_DATASET_ENDPOIT = "/api/v3/dataset"
+    CLOUD_DATASET_ENDPOINT = CLOUD_PROJECT_ENDPOINT + "/{}/dataset"
+
     # https://docs.dremio.com/software/rest-api/jobs/get-job/
     OFFSET_DEFAULT = 0
     LIMIT_DEFAULT = 100
@@ -56,10 +62,10 @@ def sql_url(cls, parameters: Parameters):
     def job_status_url(cls, parameters: Parameters, job_id):
         if type(parameters) is CloudParameters:
             return (
-                parameters.base_url
-                + UrlBuilder.CLOUD_JOB_ENDPOINT.format(parameters.cloud_project_id)
-                + "/"
-                + job_id
+                    parameters.base_url
+                    + UrlBuilder.CLOUD_JOB_ENDPOINT.format(parameters.cloud_project_id)
+                    + "/"
+                    + job_id
             )
         return parameters.base_url + UrlBuilder.SOFTWARE_JOB_ENDPOINT + "/" + job_id
 
@@ -75,11 +81,11 @@ def job_cancel_url(cls, parameters: Parameters, job_id):
 
     @classmethod
     def job_results_url(
-        cls,
-        parameters: Parameters,
-        job_id,
-        offset=OFFSET_DEFAULT,
-        limit=LIMIT_DEFAULT,
+            cls,
+            parameters: Parameters,
+            job_id,
+            offset=OFFSET_DEFAULT,
+            limit=LIMIT_DEFAULT,
     ):
         url_path = parameters.base_url
         if type(parameters) is CloudParameters:
@@ -139,3 +145,41 @@ def catalog_item_by_path_url(cls, parameters: Parameters, path_list):
         joined_path_str = "/".join(quoted_path_list).replace('"', "")
         endpoint = f"/by-path/{joined_path_str}"
         return url_path + endpoint
+
+    @classmethod
+    def create_reflection_url(cls, parameters: Parameters):
+        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
+
+        return url_path
+
+    @classmethod
+    def update_reflection_url(cls, parameters: Parameters, dataset_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(dataset_id)
+        return url_path + endpoint
+
+    @classmethod
+    def get_reflection_url(cls, parameters: Parameters, dataset_id):
+        url_path = parameters.base_url
+        if type(parameters) is CloudParameters:
+            url_path += UrlBuilder.CLOUD_DATASET_ENDPOINT.format(
+                parameters.cloud_project_id
+            )
+        else:
+            url_path += UrlBuilder.SOFTWARE_DATASET_ENDPOIT
+
+        endpoint = "/{}/reflection".format(dataset_id)
+        return url_path + endpoint

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

@@ -31,7 +31,6 @@
 import json as jsonlib
 from requests.exceptions import HTTPError
 
-
 from dbt.adapters.events.logging import AdapterLogger
 
 logger = AdapterLogger("dremio")
@@ -45,12 +44,12 @@ def _get(url, request_headers, details="", ssl_verify=True):
 
 
 def _post(
-    url,
-    request_headers=None,
-    json=None,
-    details="",
-    ssl_verify=True,
-    timeout=None,
+        url,
+        request_headers=None,
+        json=None,
+        details="",
+        ssl_verify=True,
+        timeout=None,
 ):
     if isinstance(json, str):
         json = jsonlib.loads(json)
@@ -64,6 +63,13 @@ def _post(
     return _check_error(response, details)
 
 
+def _put(url, request_headers, json=None, details="", ssl_verify=True):
+    response = session.put(
+        url, headers=request_headers, verify=ssl_verify, json=json
+    )
+    return _check_error(response, details)
+
+
 def _delete(url, request_headers, details="", ssl_verify=True):
     response = session.delete(url, headers=request_headers, verify=ssl_verify)
     return _check_error(response, details)
@@ -148,5 +154,3 @@ def _check_error(response, details=""):
             "Gateway Timeout:" + details, error, response
         )
     raise DremioException("Unknown error", error)
-
-