Implementation of wikis and tags feature (#250)

### Summary Wikis and tags feature implementation. ### Description One of dbt's core features is source documentation. Right now, this documentation does not propagate into Dremio’s wiki, leading to two separate sources of truth. The goal is to display dbt documentation - more specifically, wikis and tags - inside of Dremio. ### Test Results All added tests are running - Added tests to check wikis and tags creation for table models - Added tests to check wikis and tags creation, update and deletion for view models ### Changelog - Possibility to integrate wikis and tags by enabling `relation` option from `persist_docs` configuration - New macro `dremio__persist_docs` created - Views also perform `persist_docs` macro - Integration via REST API ### Related Issue #86 #196
dremio · Jan 13, 2025 · 50fe20a · 50fe20a
1 parent 8ad45ec
commit 50fe20a
Show file tree

Hide file tree

Showing 8 changed files with 561 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -15,6 +15,11 @@
   - 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.
+- [#250](https://github.com/dremio/dbt-dremio/pull/250) Possibility to integrate wikis and tags by enabling `relation` option from `persist_docs` configuration
+  - New macro `dremio__persist_docs` created
+  - Views also perform `persist_docs` macro
+  - Integration via REST API
+
 ## 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
@@ -23,6 +28,7 @@
 
 -   [#223](https://github.com/dremio/dbt-dremio/issues/224) Implement merge strategy for incremental materializations
 -   [#229](https://github.com/dremio/dbt-dremio/issues/229) Add max operator to get_relation_last_modified macro 
+-   [#250](https://github.com/dremio/dbt-dremio/pull/250) Implementation of wikis and tags feature
 
 # dbt-dremio v1.7.0
 

diff --git a/dbt/adapters/dremio/api/rest/client.py b/dbt/adapters/dremio/api/rest/client.py
@@ -132,6 +132,79 @@ def delete_catalog(self, cid):
             self._parameters.authentication.get_headers(),
             ssl_verify=self._parameters.authentication.verify_ssl,
         )
+
+    # dbt docs integration within Dremio wikis and tags
+    def create_wiki(self, object_id: str, text: str):
+        url = UrlBuilder.wikis_management_url(self._parameters, object_id)
+        return _post(
+            url,
+            self._parameters.authentication.get_headers(),
+            json={"text": text},
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )
+
+    def retrieve_wiki(self, object_id: str):
+        url = UrlBuilder.wikis_management_url(self._parameters, object_id)
+        return _get(
+            url,
+            self._parameters.authentication.get_headers(),
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )
+
+    def update_wiki(self, object_id: str, text: str, version: int):
+        url = UrlBuilder.wikis_management_url(self._parameters, object_id)
+        return _post(
+            url,
+            self._parameters.authentication.get_headers(),
+            json={"text": text, "version": version},
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )
+
+    def delete_wiki(self, object_id: str, version: int):
+        url = UrlBuilder.wikis_management_url(self._parameters, object_id)
+        return _post(
+            url,
+            self._parameters.authentication.get_headers(),
+            json={"text": "", "version": version},
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )
+
+
+    def create_tags(self, dataset_id: str, tags: list[str]):
+        url = UrlBuilder.tags_management_url(self._parameters, dataset_id)
+        return _post(
+            url,
+            self._parameters.authentication.get_headers(),
+            json={"tags": tags},
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )
+
+    def retrieve_tags(self, dataset_id: str):
+        url = UrlBuilder.tags_management_url(self._parameters, dataset_id)
+        return _get(
+            url,
+            self._parameters.authentication.get_headers(),
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )
+
+    def update_tags(self, dataset_id: str, tags: list[str], version: str):
+        url = UrlBuilder.tags_management_url(self._parameters, dataset_id)
+        return _post(
+            url,
+            self._parameters.authentication.get_headers(),
+            json={"tags": tags, "version": version},
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )
+
+    def delete_tags(self, dataset_id: str, version: str):
+        url = UrlBuilder.tags_management_url(self._parameters, dataset_id)
+        return _post(
+            url,
+            self._parameters.authentication.get_headers(),
+            json={"tags": [], "version": version},
+            ssl_verify=self._parameters.authentication.verify_ssl,
+        )
+
 
     def get_reflections(self, dataset_id):
         url = UrlBuilder.get_reflection_url(self._parameters, dataset_id)
@@ -158,3 +231,4 @@ def update_reflection(self, reflection_id, payload):
             json=payload,
             ssl_verify=self._parameters.authentication.verify_ssl,
         )
+
diff --git a/dbt/adapters/dremio/api/rest/url_builder.py b/dbt/adapters/dremio/api/rest/url_builder.py
@@ -34,6 +34,10 @@ class UrlBuilder:
     SOFTWARE_CATALOG_ENDPOINT = "/api/v3/catalog"
     CLOUD_CATALOG_ENDPOINT = CLOUD_PROJECT_ENDPOINT + "/{}/catalog"
 
+    DREMIO_WIKIS_ENDPOINT = "/collaboration/wiki"
+
+    DREMIO_TAGS_ENDPOINT = "/collaboration/tag"
+
     SOFTWARE_REFLECTIONS_ENDPOINT = "/api/v3/reflection"
     CLOUD_REFLECTIONS_ENDPOINT = CLOUD_PROJECT_ENDPOINT + "/{}/reflection"
 
@@ -145,6 +149,15 @@ 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
+
+    # dbt docs integration within Dremio wikis and tags
+    @classmethod
+    def wikis_management_url(cls, parameters: Parameters, object_id: str) -> str:
+        return cls.catalog_url(parameters) + f"/{object_id}{UrlBuilder.DREMIO_WIKIS_ENDPOINT}"
+
+    @classmethod
+    def tags_management_url(cls, parameters: Parameters, dataset_id: str) -> str:
+        return cls.catalog_url(parameters) + f"/{dataset_id}{UrlBuilder.DREMIO_TAGS_ENDPOINT}"
 
     @classmethod
     def create_reflection_url(cls, parameters: Parameters):

diff --git a/dbt/adapters/dremio/connections.py b/dbt/adapters/dremio/connections.py
@@ -48,7 +48,6 @@
 
 logger = AdapterLogger("dremio")
 
-
 class DremioConnectionManager(SQLConnectionManager):
     TYPE = "dremio"
     DEFAULT_CONNECTION_RETRIES = 5
@@ -132,8 +131,8 @@ def add_commit_query(self):
 
     # Auto_begin may not be relevant with the rest_api
     def add_query(
-            self, sql, auto_begin=True, bindings=None, abridge_sql_log=False,
-            fetch=False
+        self, sql, auto_begin=True, bindings=None, abridge_sql_log=False,
+        fetch=False
     ):
         connection = self.get_thread_connection()
         if auto_begin and connection.transaction_open is False:
@@ -231,6 +230,92 @@ def create_catalog(self, relation):
             logger.debug(f"Creating folder(s): {database}.{schema}")
             self._create_folders(database, schema, rest_client)
         return
+
+    # dbt docs integration with Dremio wikis and tags
+    def process_wikis(self, relation, text: str):
+        logger.debug("Integrating wikis")
+        thread_connection = self.get_thread_connection()
+        connection = self.open(thread_connection)
+        rest_client = connection.handle.get_client()
+        database = relation.database
+        schema = relation.schema
+
+        path = self._create_path_list(database,schema)
+        identifier = relation.identifier
+        path.append(identifier)
+        try:
+            catalog_info = rest_client.get_catalog_item(
+                catalog_id=None,
+                catalog_path=path,
+            )
+        except DremioNotFoundException:
+            logger.debug("Catalog not found. Returning")
+            return
+
+        object_id = catalog_info.get("id")
+        stored_wiki = rest_client.retrieve_wiki(object_id)
+        wiki_content = stored_wiki.get("text")
+        wiki_version = stored_wiki.get("version", None)
+
+        if wiki_version is None:
+            logger.debug(f"Creating wiki for {'.'.join(path)}")
+            result = rest_client.create_wiki(object_id, text)
+            logger.debug(result)
+            return
+
+        if wiki_content != text:
+            if text == "": # text is empty, delete wiki
+                logger.debug(f"Deleting wiki for {'.'.join(path)}")
+                result = rest_client.delete_wiki(object_id, wiki_version)
+                logger.debug(result)
+                return
+
+            logger.debug(f"Updating wiki for {'.'.join(path)}")
+            result = rest_client.update_wiki(object_id, text, wiki_version)
+            logger.debug(result)
+
+    def process_tags(self, relation, tags: list[str]):
+        logger.debug("Integrating tags")
+        thread_connection = self.get_thread_connection()
+        connection = self.open(thread_connection)
+        rest_client = connection.handle.get_client()
+        database = relation.database
+        schema = relation.schema
+
+        path = self._create_path_list(database,schema)
+        identifier = relation.identifier
+        path.append(identifier)
+        try:
+            catalog_info = rest_client.get_catalog_item(
+                catalog_id=None,
+                catalog_path=path,
+            )
+        except DremioNotFoundException:
+            logger.debug("Catalog not found. Returning")
+            return
+
+        object_id = catalog_info.get("id")
+        stored_tags = rest_client.retrieve_tags(object_id)
+        tags_list = stored_tags.get("tags")
+        tags_version = stored_tags.get("version", None)
+
+        if tags_version is None:
+            logger.debug(f"Creating tags for {'.'.join(path)}")
+            result = rest_client.create_tags(object_id, tags)
+            logger.debug(result)
+            return
+
+        if tags_list != tags:
+            if tags == []:  # tags is empty, delete tags
+                logger.debug(f"Deleting tags for {'.'.join(path)}")
+                result = rest_client.delete_tags(object_id, tags_version)
+                logger.debug(result)
+                return
+
+            logger.debug(f"Updating tags for {'.'.join(path)}")
+            result = rest_client.update_tags(object_id, tags, tags_version)
+            logger.debug(result)
+
 
     def create_reflection(self, name: str, reflection_type: str, anchor: DremioRelation, display: List[str],
                           dimensions: List[str],

diff --git a/dbt/adapters/dremio/impl.py b/dbt/adapters/dremio/impl.py
@@ -178,6 +178,15 @@ def run_sql_for_tests(self, sql, fetch, conn):
         finally:
             conn.transaction_open = False
 
+    # dbt docs integration with Dremio wikis and tags
+    @available
+    def process_wikis(self, relation: DremioRelation, text: str) -> None:
+        self.connections.process_wikis(relation, text)
+
+    @available
+    def process_tags(self, relation: DremioRelation, tags: list[str]) -> None:
+        self.connections.process_tags(relation, tags)
+
     @available
     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],

diff --git a/dbt/include/dremio/macros/adapters/persist_docs.sql b/dbt/include/dremio/macros/adapters/persist_docs.sql
@@ -0,0 +1,6 @@
+{% macro dremio__persist_docs(relation, model, for_relation, for_columns) -%}
+  {% if for_relation and config.persist_relation_docs() %}
+    {% do adapter.process_wikis(relation, model.description) %}
+    {% do adapter.process_tags(relation, model.tags) %}
+  {% endif %}
+{% endmacro %}
diff --git a/dbt/include/dremio/macros/materializations/view/view.sql b/dbt/include/dremio/macros/materializations/view/view.sql
@@ -41,6 +41,8 @@ limitations under the License.*/
 
   {{ enable_default_reflection() }}
 
+  {% do persist_docs(target_relation, model) %}
+
   {% do apply_grants(target_relation, grant_config, should_revoke=should_revoke) %}
 
   {{ run_hooks(post_hooks) }}