fix: added fixes for dremio cloud apis

datahub-project · Oct 28, 2024 · 7433b8b · 7433b8b
1 parent 82988fd
commit 7433b8b
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 13 deletions.
diff --git a/datahub-web-react/src/app/ingest/source/builder/sources.json b/datahub-web-react/src/app/ingest/source/builder/sources.json
@@ -309,6 +309,6 @@
         "displayName": "Dremio",
         "description": "Import Spaces, Sources, Tables and statistics from Dremio.",
         "docsUrl": "https://datahubproject.io/docs/metadata-ingestion/",
-        "recipe": "source:\n    type: dremio\n    config:\n        # Coordinates\n        hostname: null\n        port: null\n        #true if https, otherwise false\n        tls: true\n\n        #Credentials with personal access token\n        authentication_method: PAT\n        password: pass\n\n        #Or Credentials with basic auth\n        #authentication_method: password\n        #username: null\n        #password: null\n\n        stateful_ingestion:\n            enabled: true"
+        "recipe": "source:\n    type: dremio\n    config:\n        # Coordinates\n        hostname: null\n        port: null\n        #true if https, otherwise false\n        tls: true\n\n        #For cloud instance\n        #is_dremio_cloud: True\n        #dremio_cloud_project_id: <project_id>\n\n        #Credentials with personal access token\n        authentication_method: PAT\n        password: pass\n\n        #Or Credentials with basic auth\n        #authentication_method: password\n        #username: null\n        #password: null\n\n        stateful_ingestion:\n            enabled: true"
     }
 ]
diff --git a/metadata-ingestion/docs/sources/dremio/dremio_recipe.yml b/metadata-ingestion/docs/sources/dremio/dremio_recipe.yml
@@ -14,6 +14,10 @@ source:
     # username: user
     # password: pass
 
+    #For cloud instance
+    #is_dremio_cloud: True
+    #dremio_cloud_project_id: <project_id>
+
     include_query_lineage: True
 
     #Optional

diff --git a/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_api.py b/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_api.py
@@ -266,15 +266,11 @@ def execute_query(self, query: str, timeout: int = 300) -> List[Dict[str, Any]]:
                     return future.result(timeout=timeout)
                 except concurrent.futures.TimeoutError:
                     self.cancel_query(job_id)
-                    self.report.failure(
-                        f"Query execution timed out after {timeout} seconds"
-                    )
-                    raise TimeoutError(
+                    raise DremioAPIException(
                         f"Query execution timed out after {timeout} seconds"
                     )
                 except RuntimeError as e:
-                    self.report.failure("Query Execution failed", exc=e)
-                    raise DremioAPIException("Query Execution failed: {str(e)}")
+                    raise DremioAPIException(f"{str(e)}")
 
         except requests.RequestException as e:
             raise DremioAPIException(f"Error executing query: {str(e)}")
@@ -490,16 +486,15 @@ def get_all_tables_and_columns(self, containers: Deque) -> List[Dict]:
                     deny_schema_pattern=deny_schema_condition,
                     container_name=schema.container_name.lower(),
                 )
-
                 all_tables_and_columns.extend(
                     self.execute_query(
                         query=formatted_query,
                     )
                 )
-            except Exception as exc:
-                logger.debug(f"Failed with {exc}")
+            except DremioAPIException as e:
                 self.report.warning(
-                    f"{schema.subclass} {schema.container_name} had no tables or views"
+                    title=f"{schema.subclass} {schema.container_name} had no tables or views",
+                    message=str(e),
                 )
 
         tables = []

diff --git a/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_sql_queries.py b/metadata-ingestion/src/datahub/ingestion/source/dremio/dremio_sql_queries.py
@@ -135,15 +135,14 @@ class DremioSQLQueries:
         """
 
     QUERY_DATASETS_CLOUD = """
-        SELECT* FROM
+        SELECT * FROM
         (
         SELECT
             RESOURCE_ID,
             V.TABLE_NAME,
             OWNER,
             PATH AS TABLE_SCHEMA,
             CONCAT(REPLACE(REPLACE(
-            CONCAT(REPLACE(REPLACE(
             REPLACE(V.PATH, ', ', '.'),
              '[', ''), ']', ''
              )) AS FULL_TABLE_PATH,