openmetadata-ingestion 1.9.13.0__py3-none-any.whl → 1.10.1.0__py3-none-any.whl

metadata/ingestion/source/dashboard/redash/metadata.py

@@ -110,6 +110,8 @@ class RedashSource(DashboardServiceSource):
         )
 
     def get_dashboards_list(self) -> Optional[List[dict]]:
+        if not self.source_config.includeOwners:
+            logger.debug("Skipping owner information as includeOwners is False")
         return self.dashboard_list
 
     def get_dashboard_name(self, dashboard: dict) -> str:
@@ -123,6 +125,8 @@ class RedashSource(DashboardServiceSource):
         Get owner from email
         """
         try:
+            if not self.source_config.includeOwners:
+                return None
             if dashboard_details.get("user") and dashboard_details["user"].get("email"):
                 return self.metadata.get_reference_by_email(
                     dashboard_details["user"].get("email")

metadata/ingestion/source/dashboard/sigma/metadata.py

@@ -91,6 +91,8 @@ class SigmaSource(DashboardServiceSource):
         """
         get list of dashboard
         """
+        if not self.source_config.includeOwners:
+            logger.debug("Skipping owner information as includeOwners is False")
         return self.client.get_dashboards()
 
     def get_dashboard_name(self, dashboard: Workbook) -> Optional[str]:
@@ -367,6 +369,8 @@ class SigmaSource(DashboardServiceSource):
         Get owner from email
         """
         try:
+            if not self.source_config.includeOwners:
+                return None
             if dashboard_details.ownerId:
                 owner = self.client.get_owner_detail(dashboard_details.ownerId)
                 return self.metadata.get_reference_by_email(owner.email)

metadata/ingestion/source/dashboard/superset/api_source.py

@@ -77,6 +77,8 @@ class SupersetAPISource(SupersetSourceMixin):
         """
         Get List of all dashboards
         """
+        if not self.source_config.includeOwners:
+            logger.debug("Skipping owner information as includeOwners is False")
         current_page = 0
         page_size = 25
         total_dashboards = self.client.fetch_total_dashboards()

metadata/ingestion/source/dashboard/superset/db_source.py

@@ -111,6 +111,8 @@ class SupersetDBSource(SupersetSourceMixin):
         """
         Get List of all dashboards
         """
+        if not self.source_config.includeOwners:
+            logger.debug("Skipping owner information as includeOwners is False")
         query = (
             FETCH_DASHBOARDS
             if self.source_config.includeDraftDashboard

metadata/ingestion/source/dashboard/superset/mixin.py

@@ -117,6 +117,8 @@ class SupersetSourceMixin(DashboardServiceSource):
         self, dashboard_details: Union[DashboardResult, FetchDashboard]
     ) -> Optional[EntityReferenceList]:
         try:
+            if not self.source_config.includeOwners:
+                return None
             if hasattr(dashboard_details, "owners"):
                 for owner in dashboard_details.owners or []:
                     if owner.email:

metadata/ingestion/source/dashboard/tableau/client.py

@@ -106,8 +106,15 @@ class TableauClient:
     def site_id(self) -> str:
         return self.tableau_server.site_id
 
-    def get_tableau_owner(self, owner_id: str) -> Optional[TableauOwner]:
+    def get_tableau_owner(
+        self, owner_id: str, include_owners: bool = True
+    ) -> Optional[TableauOwner]:
+        """
+        Get tableau owner with optional include_owners flag
+        """
         try:
+            if not include_owners:
+                return None
             if owner_id in self.owner_cache:
                 return self.owner_cache[owner_id]
             owner = self.tableau_server.users.get_by_id(owner_id) if owner_id else None
@@ -122,7 +129,7 @@ class TableauClient:
         return None
 
     def get_workbook_charts_and_user_count(
-        self, views: List[ViewItem]
+        self, views: List[ViewItem], include_owners: bool = True
     ) -> Optional[Tuple[Optional[int], Optional[List[TableauChart]]]]:
         """
         Fetches workbook charts and dashboard user view count
@@ -136,7 +143,7 @@ class TableauClient:
                         id=str(view.id),
                         name=view.name,
                         tags=view.tags,
-                        owner=self.get_tableau_owner(view.owner_id),
+                        owner=self.get_tableau_owner(view.owner_id, include_owners),
                         contentUrl=view.content_url,
                         sheetType=view.sheet_type,
                     )
@@ -202,7 +209,7 @@ class TableauClient:
             logger.debug(f"Failed to get project parents by id: {str(e)}")
         return None
 
-    def get_workbooks(self) -> Iterable[TableauDashboard]:
+    def get_workbooks(self, include_owners: bool = True) -> Iterable[TableauDashboard]:
         """
         Fetch all tableau workbooks
         """
@@ -212,7 +219,7 @@ class TableauClient:
             try:
                 self.tableau_server.workbooks.populate_views(workbook, usage=True)
                 charts, user_views = self.get_workbook_charts_and_user_count(
-                    workbook.views
+                    workbook.views, include_owners
                 )
                 workbook = TableauDashboard(
                     id=str(workbook.id),
@@ -220,7 +227,7 @@ class TableauClient:
                     project=TableauBaseModel(
                         id=str(workbook.project_id), name=workbook.project_name
                     ),
-                    owner=self.get_tableau_owner(workbook.owner_id),
+                    owner=self.get_tableau_owner(workbook.owner_id, include_owners),
                     description=workbook.description,
                     tags=workbook.tags,
                     webpageUrl=workbook.webpage_url,
@@ -248,9 +255,11 @@ class TableauClient:
             "Please check if the user has permissions to access the Dashboards information"
         )
 
-    def test_get_workbook_views(self):
+    def test_get_workbook_views(self, include_owners: bool = True):
         workbook = self.test_get_workbooks()
-        charts, _ = self.get_workbook_charts_and_user_count(workbook.views)
+        charts, _ = self.get_workbook_charts_and_user_count(
+            workbook.views, include_owners
+        )
         if charts:
             return True
         raise TableauChartsException(
@@ -258,9 +267,11 @@ class TableauClient:
             "Please check if the user has permissions to access the Charts information"
         )
 
-    def test_get_owners(self) -> Optional[List[TableauOwner]]:
+    def test_get_owners(
+        self, include_owners: bool = True
+    ) -> Optional[List[TableauOwner]]:
         workbook = self.test_get_workbooks()
-        owners = self.get_tableau_owner(workbook.owner_id)
+        owners = self.get_tableau_owner(workbook.owner_id, include_owners)
         if owners is not None:
             return owners
         raise TableauOwnersNotFound(

metadata/ingestion/source/dashboard/tableau/metadata.py

@@ -132,7 +132,11 @@ class TableauSource(DashboardServiceSource):
         return cls(config, metadata)
 
     def get_dashboards_list(self) -> Iterable[TableauDashboard]:
-        yield from self.client.get_workbooks()
+        if not self.source_config.includeOwners:
+            logger.debug("Skipping owner information as includeOwners is False")
+        yield from self.client.get_workbooks(
+            include_owners=self.source_config.includeOwners
+        )
 
     def get_dashboard_name(self, dashboard: TableauDashboard) -> str:
         return dashboard.name
@@ -151,6 +155,8 @@ class TableauSource(DashboardServiceSource):
         Get dashboard owner from email
         """
         try:
+            if not self.source_config.includeOwners:
+                return None
             if dashboard_details.owner and dashboard_details.owner.email:
                 return self.metadata.get_reference_by_email(
                     dashboard_details.owner.email

metadata/ingestion/source/database/bigquery/connection.py

@@ -153,6 +153,10 @@ def test_connection(
             return policy_tags
 
     def test_tags():
+        if not service_connection.includePolicyTags:
+            logger.info("'includePolicyTags' is set to false, so skipping this test.")
+            return None
+
         taxonomy_project_ids = []
         if engine.url.host:
             taxonomy_project_ids.append(engine.url.host)

metadata/ingestion/source/database/bigquery/metadata.py

@@ -187,34 +187,12 @@ def get_columns(bq_schema):
             "max_length": field.max_length,
             "system_data_type": get_system_data_type(col_type),
             "is_complex": is_complex_type(str(col_type)),
-            "policy_tags": None,
+            "policy_tags": field.policy_tags,
         }
         if getattr(field, "fields", None):
             # Nested Columns available
             col_obj["children"] = get_columns(field.fields)
-        try:
-            if field.policy_tags:
-                policy_tag_name = field.policy_tags.names[0]
-                taxonomy_name = (
-                    policy_tag_name.split("/policyTags/")[0] if policy_tag_name else ""
-                )
-                if not taxonomy_name:
-                    raise NotImplementedError(
-                        f"Taxonomy Name not present for {field.name}"
-                    )
-                col_obj["taxonomy"] = (
-                    PolicyTagManagerClient()
-                    .get_taxonomy(name=taxonomy_name)
-                    .display_name
-                )
-                col_obj["policy_tags"] = (
-                    PolicyTagManagerClient()
-                    .get_policy_tag(name=policy_tag_name)
-                    .display_name
-                )
-        except Exception as exc:
-            logger.debug(traceback.format_exc())
-            logger.warning(f"Skipping Policy Tag: {exc}")
+
         col_list.append(col_obj)
     return col_list
 
@@ -513,6 +491,13 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
                         metadata=self.metadata,
                         system_tags=True,
                     )
+
+            if not self.service_connection.includePolicyTags:
+                logger.info(
+                    "'includePolicyTags' is set to false so skipping policy tag ingestion"
+                )
+                return
+
             # Fetching policy tags on the column level
             list_project_ids = [self.context.get().database]
             if not self.service_connection.taxonomyProjectID:
@@ -712,6 +697,32 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
                     table_tag_labels.append(tag_label)
         return table_tag_labels
 
+    def get_policy_tags_for_column(self, column: dict) -> dict:
+        try:
+            if column.get("policy_tags"):
+                policy_tag_name = column["policy_tags"].names[0]
+                taxonomy_name = (
+                    policy_tag_name.split("/policyTags/")[0] if policy_tag_name else ""
+                )
+                if not taxonomy_name:
+                    raise NotImplementedError(
+                        f"Taxonomy Name not present for {column['name']}"
+                    )
+                column["taxonomy"] = (
+                    PolicyTagManagerClient()
+                    .get_taxonomy(name=taxonomy_name)
+                    .display_name
+                )
+                column["policy_tags"] = (
+                    PolicyTagManagerClient()
+                    .get_policy_tag(name=policy_tag_name)
+                    .display_name
+                )
+                return column
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(f"Skipping Policy Tag: {exc}")
+
     def get_column_tag_labels(
         self, table_name: str, column: dict
     ) -> Optional[List[TagLabel]]:
@@ -719,12 +730,14 @@ class BigquerySource(LifeCycleQueryMixin, CommonDbSourceService, MultiDBSource):
         This will only get executed if the tags context
         is properly informed
         """
-        if column.get("policy_tags"):
+        if self.service_connection.includePolicyTags and column.get("policy_tags"):
+            self.get_policy_tags_for_column(column)
             return fetch_tag_labels_om(
                 metadata=self.metadata,
                 tags=[column["policy_tags"]],
                 classification_name=column["taxonomy"],
-                include_tags=self.source_config.includeTags,
+                include_tags=self.source_config.includeTags
+                and self.service_connection.includePolicyTags,
             )
         return None
 

metadata/ingestion/source/database/column_type_parser.py

@@ -293,7 +293,7 @@ class ColumnTypeParser:
         "STRG": "STRING",
         "RSTR": "STRING",
         # azuresql
-        "HIERARCHYID": "UNKNOWN",
+        "HIERARCHYID": "HIERARCHYID",
     }
 
     _COMPLEX_TYPE = re.compile("^(struct|map|array|uniontype)")

metadata/ingestion/source/database/common_pg_mappings.py

@@ -45,5 +45,12 @@ ischema_names.update(
         "txid_snapshot": create_sqlalchemy_type("TXID_SNAPSHOT"),
         "xid": SqlAlchemyString,
         "xml": create_sqlalchemy_type("XML"),
+        # PostgreSQL range types (used by TimescaleDB for chunk boundaries)
+        "int4range": create_sqlalchemy_type("INT4RANGE"),
+        "int8range": create_sqlalchemy_type("INT8RANGE"),
+        "numrange": create_sqlalchemy_type("NUMRANGE"),
+        "tsrange": create_sqlalchemy_type("TSRANGE"),
+        "tstzrange": create_sqlalchemy_type("TSTZRANGE"),
+        "daterange": create_sqlalchemy_type("DATERANGE"),
     }
 )

metadata/ingestion/source/database/databricks/auth.py

@@ -0,0 +1,99 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+"""
+This module provides authentication utilities for Databricks and Unity Catalog connections.
+"""
+from typing import Union
+
+from databricks.sdk.core import Config, azure_service_principal, oauth_service_principal
+
+from metadata.generated.schema.entity.services.connections.database.databricks.azureAdSetup import (
+    AzureAdSetup,
+)
+from metadata.generated.schema.entity.services.connections.database.databricks.databricksOAuth import (
+    DatabricksOauth,
+)
+from metadata.generated.schema.entity.services.connections.database.databricks.personalAccessToken import (
+    PersonalAccessToken,
+)
+from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
+    DatabricksConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
+    UnityCatalogConnection,
+)
+
+
+def get_personal_access_token_auth(
+    connection: Union[DatabricksConnection, UnityCatalogConnection],
+) -> dict:
+    """
+    Configure Personal Access Token authentication
+    """
+    return {"access_token": connection.authType.token.get_secret_value()}
+
+
+def get_databricks_oauth_auth(
+    connection: Union[DatabricksConnection, UnityCatalogConnection],
+):
+    """
+    Create Databricks OAuth2 M2M credentials provider for Service Principal authentication
+    """
+
+    def credential_provider():
+        hostname = connection.hostPort.split(":")[0]
+        config = Config(
+            host=f"https://{hostname}",
+            client_id=connection.authType.clientId,
+            client_secret=connection.authType.clientSecret.get_secret_value(),
+        )
+        return oauth_service_principal(config)
+
+    return {"credentials_provider": credential_provider}
+
+
+def get_azure_ad_auth(connection: Union[DatabricksConnection, UnityCatalogConnection]):
+    """
+    Create Azure AD credentials provider for Azure Service Principal authentication
+    """
+
+    def credential_provider():
+        hostname = connection.hostPort.split(":")[0]
+        config = Config(
+            host=f"https://{hostname}",
+            azure_client_secret=connection.authType.azureClientSecret.get_secret_value(),
+            azure_client_id=connection.authType.azureClientId,
+            azure_tenant_id=connection.authType.azureTenantId,
+        )
+        return azure_service_principal(config)
+
+    return {"credentials_provider": credential_provider}
+
+
+def get_auth_config(
+    connection: Union[DatabricksConnection, UnityCatalogConnection],
+) -> dict:
+    """
+    Get authentication configuration for Databricks connection
+    """
+    auth_method = {
+        PersonalAccessToken: get_personal_access_token_auth,
+        DatabricksOauth: get_databricks_oauth_auth,
+        AzureAdSetup: get_azure_ad_auth,
+    }.get(type(connection.authType))
+
+    if not auth_method:
+        raise ValueError(
+            f"Unsupported authentication type: {type(connection.authType)}"
+        )
+
+    return auth_method(connection)

metadata/ingestion/source/database/databricks/client.py

@@ -61,14 +61,13 @@ class DatabricksClient:
     ):
         self.config = config
         base_url, *_ = self.config.hostPort.split(":")
-        auth_token = self.config.token.get_secret_value()
         self.base_url = f"https://{base_url}{API_VERSION}"
         self.base_query_url = f"{self.base_url}{QUERIES_PATH}"
         self.base_job_url = f"https://{base_url}{JOB_API_VERSION}/jobs"
         self.jobs_list_url = f"{self.base_job_url}/list"
         self.jobs_run_list_url = f"{self.base_job_url}/runs/list"
         self.headers = {
-            "Authorization": f"Bearer {auth_token}",
+            **self._get_auth_header(),
             "Content-Type": "application/json",
         }
         self.api_timeout = self.config.connectionTimeout or 120
@@ -81,6 +80,12 @@ class DatabricksClient:
         self.engine = engine
         self.client = requests
 
+    def _get_auth_header(self) -> dict[str, str]:
+        """
+        Method to get auth header
+        """
+        return {"Authorization": f"Bearer {self.config.token.get_secret_value()}"}
+
     def test_query_api_access(self) -> None:
         res = self.client.get(
             self.base_query_url, headers=self.headers, timeout=self.api_timeout

metadata/ingestion/source/database/databricks/connection.py

@@ -12,6 +12,7 @@
 """
 Source connection handler
 """
+from copy import deepcopy
 from functools import partial
 from typing import Optional
 
@@ -38,9 +39,17 @@ from metadata.ingestion.connections.test_connections import (
     test_connection_steps,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.databricks.auth import get_auth_config
 from metadata.ingestion.source.database.databricks.queries import (
     DATABRICKS_GET_CATALOGS,
     DATABRICKS_SQL_STATEMENT_TEST,
+    TEST_CATALOG_TAGS,
+    TEST_COLUMN_LINEAGE,
+    TEST_COLUMN_TAGS,
+    TEST_SCHEMA_TAGS,
+    TEST_TABLE_LINEAGE,
+    TEST_TABLE_TAGS,
+    TEST_VIEW_DEFINITIONS,
 )
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import ingestion_logger
@@ -122,8 +131,7 @@ class DatabricksEngineWrapper:
 
 
 def get_connection_url(connection: DatabricksConnection) -> str:
-    url = f"{connection.scheme.value}://token:{connection.token.get_secret_value()}@{connection.hostPort}"
-    return url
+    return f"{connection.scheme.value}://{connection.hostPort}"
 
 
 def get_connection(connection: DatabricksConnection) -> Engine:
@@ -131,17 +139,27 @@ def get_connection(connection: DatabricksConnection) -> Engine:
     Create connection
     """
 
+    if not connection.connectionArguments:
+        connection.connectionArguments = init_empty_connection_arguments()
+
     if connection.httpPath:
-        if not connection.connectionArguments:
-            connection.connectionArguments = init_empty_connection_arguments()
         connection.connectionArguments.root["http_path"] = connection.httpPath
 
-    return create_generic_db_connection(
+    auth_args = get_auth_config(connection)
+
+    original_connection_arguments = connection.connectionArguments
+    connection.connectionArguments = deepcopy(original_connection_arguments)
+    connection.connectionArguments.root.update(auth_args)
+
+    engine = create_generic_db_connection(
         connection=connection,
         get_connection_url_fn=get_connection_url,
         get_connection_args_fn=get_connection_args_common,
     )
 
+    connection.connectionArguments = original_connection_arguments
+    return engine
+
 
 def test_connection(
     metadata: OpenMetadata,
@@ -170,6 +188,11 @@ def test_connection(
     # Create wrapper to avoid multiple schema calls
     engine_wrapper = DatabricksEngineWrapper(connection)
 
+    # Helper function to get first catalog for tag queries
+    def get_first_catalog():
+        catalogs = engine_wrapper.get_catalogs(catalog_name=service_connection.catalog)
+        return catalogs[0] if catalogs else service_connection.catalog or "main"
+
     test_fn = {
         "CheckAccess": partial(test_connection_engine_step, connection),
         "GetSchemas": partial(
@@ -187,6 +210,41 @@ def test_connection(
                 query_history=service_connection.queryHistoryTable
             ),
         ),
+        "GetViewDefinitions": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_VIEW_DEFINITIONS,
+        ),
+        "GetCatalogTags": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_CATALOG_TAGS.format(database_name=get_first_catalog()),
+        ),
+        "GetSchemaTags": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_SCHEMA_TAGS.format(database_name=get_first_catalog()),
+        ),
+        "GetTableTags": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_TABLE_TAGS.format(database_name=get_first_catalog()),
+        ),
+        "GetColumnTags": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_COLUMN_TAGS.format(database_name=get_first_catalog()),
+        ),
+        "GetTableLineage": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_TABLE_LINEAGE,
+        ),
+        "GetColumnLineage": partial(
+            test_database_query,
+            engine=connection,
+            statement=TEST_COLUMN_LINEAGE,
+        ),
     }
 
     return test_connection_steps(

metadata/ingestion/source/database/databricks/queries.py

@@ -114,3 +114,64 @@ DATABRICKS_GET_COLUMN_LINEAGE_FOR_JOB = textwrap.dedent(
         AND event_time >= current_date() - INTERVAL 90 DAYS
     """
 )
+
+# Test connection queries
+TEST_VIEW_DEFINITIONS = textwrap.dedent(
+    """
+    SELECT
+        TABLE_NAME,
+        TABLE_SCHEMA,
+        VIEW_DEFINITION
+    FROM INFORMATION_SCHEMA.VIEWS
+    WHERE VIEW_DEFINITION IS NOT NULL
+    LIMIT 1
+    """
+)
+
+TEST_CATALOG_TAGS = textwrap.dedent(
+    """
+    SELECT COUNT(*) as count
+    FROM `{database_name}`.information_schema.catalog_tags
+    WHERE 1=0
+    """
+)
+
+TEST_SCHEMA_TAGS = textwrap.dedent(
+    """
+    SELECT COUNT(*) as count
+    FROM `{database_name}`.information_schema.schema_tags
+    WHERE 1=0
+    """
+)
+
+TEST_TABLE_TAGS = textwrap.dedent(
+    """
+    SELECT COUNT(*) as count
+    FROM `{database_name}`.information_schema.table_tags
+    WHERE 1=0
+    """
+)
+
+TEST_COLUMN_TAGS = textwrap.dedent(
+    """
+    SELECT COUNT(*) as count
+    FROM `{database_name}`.information_schema.column_tags
+    WHERE 1=0
+    """
+)
+
+TEST_TABLE_LINEAGE = textwrap.dedent(
+    """
+    SELECT COUNT(*) as count
+    FROM system.access.table_lineage
+    WHERE 1=0
+    """
+)
+
+TEST_COLUMN_LINEAGE = textwrap.dedent(
+    """
+    SELECT COUNT(*) as count
+    FROM system.access.column_lineage
+    WHERE 1=0
+    """
+)