openmetadata-ingestion 1.3.0.1__py3-none-any.whl → 1.3.1.1__py3-none-any.whl

metadata/ingestion/models/patch_request.py

@@ -19,6 +19,7 @@ from pydantic import BaseModel
 
 from metadata.ingestion.api.models import Entity, T
 from metadata.ingestion.ometa.mixins.patch_mixin_utils import PatchOperation
+from metadata.ingestion.ometa.utils import model_str
 
 
 class PatchRequest(BaseModel):
@@ -138,12 +139,15 @@ ALLOWED_COMMON_PATCH_FIELDS = {
 
 RESTRICT_UPDATE_LIST = ["description", "tags", "owner"]
 
+ARRAY_ENTITY_FIELDS = ["columns", "tasks", "fields"]
+
 
 def build_patch(
     source: T,
     destination: T,
     allowed_fields: Optional[Dict] = None,
     restrict_update_fields: Optional[List] = None,
+    array_entity_fields: Optional[List] = None,
 ) -> Optional[jsonpatch.JsonPatch]:
     """
     Given an Entity type and Source entity and Destination entity,
@@ -163,6 +167,13 @@ def build_patch(
     source = _remove_change_description(source)
     destination = _remove_change_description(destination)
 
+    if array_entity_fields:
+        _sort_array_entity_fields(
+            source=source,
+            destination=destination,
+            array_entity_fields=array_entity_fields,
+        )
+
     # Get the difference between source and destination
     if allowed_fields:
         patch = jsonpatch.make_patch(
@@ -192,20 +203,61 @@ def build_patch(
     # for a user editable fields like descriptions, tags we only want to support "add" operation in patch
     # we will remove the other operations for replace, remove from here
     if restrict_update_fields:
-        patch.patch = [
-            patch_ops
-            for patch_ops in patch.patch
+        patch_ops_list = []
+        for patch_ops in patch.patch or []:
             if _determine_restricted_operation(
-                patch_ops=patch_ops,
-                restrict_update_fields=restrict_update_fields,
-            )
-        ]
-
+                patch_ops=patch_ops, restrict_update_fields=restrict_update_fields
+            ):
+                if (
+                    patch_ops.get("op") == PatchOperation.REPLACE.value
+                    and patch_ops.get("value") is None
+                ):
+                    patch_ops["op"] = PatchOperation.REMOVE.value
+                    del patch_ops["value"]
+                patch_ops_list.append(patch_ops)
+        patch.patch = patch_ops_list
     return patch
 
 
+def _sort_array_entity_fields(
+    source: T,
+    destination: T,
+    array_entity_fields: Optional[List] = None,
+):
+    """
+    Sort the array entity fields to make sure the order is consistent
+    """
+    for field in array_entity_fields or []:
+        if hasattr(destination, field) and hasattr(source, field):
+            destination_attributes = getattr(destination, field)
+            source_attributes = getattr(source, field)
+
+            # Create a dictionary of destination attributes for easy lookup
+            destination_dict = {
+                model_str(attr.name): attr for attr in destination_attributes
+            }
+
+            updated_attributes = []
+            for source_attr in source_attributes or []:
+                # Update the destination attribute with the source attribute
+                destination_attr = destination_dict.get(model_str(source_attr.name))
+                if destination_attr:
+                    updated_attributes.append(
+                        source_attr.copy(update=destination_attr.__dict__)
+                    )
+                    # Remove the updated attribute from the destination dictionary
+                    del destination_dict[model_str(source_attr.name)]
+                else:
+                    updated_attributes.append(None)
+
+            # Combine the updated attributes with the remaining destination attributes
+            final_attributes = updated_attributes + list(destination_dict.values())
+            setattr(destination, field, final_attributes)
+
+
 def _determine_restricted_operation(
-    patch_ops: Dict, restrict_update_fields: Optional[List] = None
+    patch_ops: Dict,
+    restrict_update_fields: Optional[List] = None,
 ) -> bool:
     """
     Only retain add operation for restrict_update_fields fields

metadata/ingestion/ometa/client.py

@@ -21,6 +21,7 @@ from requests.exceptions import HTTPError
 
 from metadata.config.common import ConfigModel
 from metadata.ingestion.ometa.credentials import URL, get_api_version
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -251,6 +252,7 @@ class REST:
 
         return None
 
+    @calculate_execution_time(context="GET")
     def get(self, path, data=None):
         """
         GET method
@@ -264,6 +266,7 @@ class REST:
         """
         return self._request("GET", path, data)
 
+    @calculate_execution_time(context="POST")
     def post(self, path, data=None):
         """
         POST method
@@ -277,6 +280,7 @@ class REST:
         """
         return self._request("POST", path, data)
 
+    @calculate_execution_time(context="PUT")
     def put(self, path, data=None):
         """
         PUT method
@@ -290,6 +294,7 @@ class REST:
         """
         return self._request("PUT", path, data)
 
+    @calculate_execution_time(context="PATCH")
     def patch(self, path, data=None):
         """
         PATCH method
@@ -308,6 +313,7 @@ class REST:
             headers={"Content-type": "application/json-patch+json"},
         )
 
+    @calculate_execution_time(context="DELETE")
     def delete(self, path, data=None):
         """
         DELETE method

metadata/ingestion/ometa/mixins/custom_property_mixin.py

@@ -15,7 +15,8 @@ To be used by OpenMetadata class
 """
 from typing import Dict
 
-from metadata.generated.schema.api.data.createCustomProperty import PropertyType
+from metadata.generated.schema.type.customProperty import PropertyType
+from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.models.custom_properties import (
     CustomPropertyDataTypes,
     CustomPropertyType,
@@ -54,16 +55,6 @@ class OMetaCustomPropertyMixin:
             f"/metadata/types/name/{entity_type}?category=field"
         )
 
-        # Get the data type of the custom property
-        if not ometa_custom_property.createCustomPropertyRequest.propertyType:
-            custom_property_type = self.get_custom_property_type(
-                data_type=ometa_custom_property.custom_property_type
-            )
-            property_type = PropertyType(id=custom_property_type.id, type="type")
-            ometa_custom_property.createCustomPropertyRequest.propertyType = (
-                property_type
-            )
-
         resp = self.client.put(
             f"/metadata/types/{entity_schema.get('id')}",
             data=ometa_custom_property.createCustomPropertyRequest.json(),
@@ -78,3 +69,12 @@ class OMetaCustomPropertyMixin:
         """
         resp = self.client.get(f"/metadata/types/name/{data_type.value}?category=field")
         return CustomPropertyType(**resp)
+
+    def get_property_type_ref(self, data_type: CustomPropertyDataTypes) -> PropertyType:
+        """
+        Get the PropertyType for custom properties
+        """
+        custom_property_type = self.get_custom_property_type(data_type=data_type)
+        return PropertyType(
+            __root__=EntityReference(id=custom_property_type.id, type="type")
+        )

metadata/ingestion/ometa/mixins/patch_mixin.py

@@ -119,6 +119,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         destination: T,
         allowed_fields: Optional[Dict] = None,
         restrict_update_fields: Optional[List] = None,
+        array_entity_fields: Optional[List] = None,
     ) -> Optional[T]:
         """
         Given an Entity type and Source entity and Destination entity,
@@ -140,6 +141,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
                 destination=destination,
                 allowed_fields=allowed_fields,
                 restrict_update_fields=restrict_update_fields,
+                array_entity_fields=array_entity_fields,
             )
 
             if not patch:

metadata/ingestion/ometa/ometa_api.py

@@ -418,7 +418,7 @@ class OpenMetadata(
         self,
         entity: Type[T],
         fields: Optional[List[str]] = None,
-        limit: int = 1000,
+        limit: int = 100,
         params: Optional[Dict[str, str]] = None,
         skip_on_failure: bool = False,
     ) -> Iterable[T]:

metadata/ingestion/sink/metadata_rest.py

@@ -62,6 +62,7 @@ from metadata.ingestion.models.ometa_classification import OMetaTagAndClassifica
 from metadata.ingestion.models.ometa_topic_data import OMetaTopicSampleData
 from metadata.ingestion.models.patch_request import (
     ALLOWED_COMMON_PATCH_FIELDS,
+    ARRAY_ENTITY_FIELDS,
     RESTRICT_UPDATE_LIST,
     PatchedEntity,
     PatchRequest,
@@ -82,7 +83,7 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.dashboard.dashboard_service import DashboardUsage
 from metadata.ingestion.source.database.database_service import DataModelLink
 from metadata.profiler.api.models import ProfilerResponse
-from metadata.utils.helpers import calculate_execution_time
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.logger import get_log_name, ingestion_logger
 
 logger = ingestion_logger()
@@ -129,7 +130,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         logger.debug(f"Processing Create request {type(record)}")
         return self.write_create_request(record)
 
-    @calculate_execution_time
+    @calculate_execution_time(store=False)
     def _run(self, record: Entity, *_, **__) -> Either[Any]:
         """
         Default implementation for the single dispatch
@@ -179,6 +180,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             destination=record.new_entity,
             allowed_fields=ALLOWED_COMMON_PATCH_FIELDS,
             restrict_update_fields=RESTRICT_UPDATE_LIST,
+            array_entity_fields=ARRAY_ENTITY_FIELDS,
         )
         patched_entity = PatchedEntity(new_entity=entity) if entity else None
         return Either(right=patched_entity)

metadata/ingestion/source/dashboard/looker/metadata.py

@@ -661,7 +661,7 @@ class LookerSource(DashboardServiceSource):
             ],
             # Dashboards are created from the UI directly. They are not linked to a project
             # like LookML assets, but rather just organised in folders.
-            project=self._get_dashboard_project(dashboard_details),
+            project=self.get_project_name(dashboard_details),
             sourceUrl=f"{clean_uri(self.service_connection.hostPort)}/dashboards/{dashboard_details.id}",
             service=self.context.dashboard_service,
             owner=self.get_owner_ref(dashboard_details=dashboard_details),
@@ -669,8 +669,7 @@ class LookerSource(DashboardServiceSource):
         yield Either(right=dashboard_request)
         self.register_record(dashboard_request=dashboard_request)
 
-    @staticmethod
-    def _get_dashboard_project(dashboard_details: LookerDashboard) -> Optional[str]:
+    def get_project_name(self, dashboard_details: LookerDashboard) -> Optional[str]:
         """
         Get dashboard project if the folder is informed
         """
@@ -680,7 +679,7 @@ class LookerSource(DashboardServiceSource):
             logger.debug(
                 f"Cannot get folder name from dashboard [{dashboard_details.title}] - [{exc}]"
             )
-            return None
+        return None
 
     @staticmethod
     def _clean_table_name(table_name: str) -> str:

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

@@ -124,6 +124,10 @@ class MetabaseClient:
         try:
             resp_dashboard = self.client.get(f"/dashboard/{dashboard_id}")
             if resp_dashboard:
+                # Small hack needed to support Metabase versions older than 0.48
+                # https://www.metabase.com/releases/metabase-48#fyi--breaking-changes
+                if "ordered_cards" in resp_dashboard:
+                    resp_dashboard["dashcards"] = resp_dashboard["ordered_cards"]
                 return MetabaseDashboardDetails(**resp_dashboard)
         except Exception:
             logger.debug(traceback.format_exc())

metadata/ingestion/source/dashboard/metabase/metadata.py

@@ -178,7 +178,7 @@ class MetabaseSource(DashboardServiceSource):
         Returns:
             Iterable[CreateChartRequest]
         """
-        charts = dashboard_details.ordered_cards
+        charts = dashboard_details.dashcards
         for chart in charts:
             try:
                 chart_details = chart.card
@@ -225,7 +225,7 @@ class MetabaseSource(DashboardServiceSource):
         if not db_service_name:
             return
         chart_list, dashboard_name = (
-            dashboard_details.ordered_cards,
+            dashboard_details.dashcards,
             str(dashboard_details.id),
         )
         for chart in chart_list:
@@ -324,8 +324,9 @@ class MetabaseSource(DashboardServiceSource):
         self, chart_details: MetabaseChart, db_service_name: str, dashboard_name: str
     ) -> Iterable[Either[AddLineageRequest]]:
         table = self.client.get_table(chart_details.table_id)
+        table_name = table.name or table.display_name
 
-        if table is None or table.display_name is None:
+        if table is None or table_name is None:
             return
 
         database_name = table.db.details.db if table.db and table.db.details else None
@@ -334,7 +335,7 @@ class MetabaseSource(DashboardServiceSource):
             database=database_name,
             service_name=db_service_name,
             database_schema=table.table_schema,
-            table=table.display_name,
+            table=table_name,
         )
 
         to_fqn = fqn.build(

metadata/ingestion/source/dashboard/metabase/models.py

@@ -67,7 +67,7 @@ class MetabaseChart(BaseModel):
     display: Optional[str]
 
 
-class OrderedCard(BaseModel):
+class DashCard(BaseModel):
     card: MetabaseChart
 
 
@@ -77,7 +77,7 @@ class MetabaseDashboardDetails(BaseModel):
     """
 
     description: Optional[str]
-    ordered_cards: List[OrderedCard]
+    dashcards: List[DashCard]
     name: Optional[str]
     id: int
     collection_id: Optional[str]

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

@@ -152,6 +152,23 @@ class TableauSource(DashboardServiceSource):
                 include_tags=self.source_config.includeTags,
             )
 
+    def _get_datamodel_sql_query(self, data_model: DataSource) -> Optional[str]:
+        """
+        Method to fetch the custom sql query from the tableau datamodels
+        """
+        try:
+            sql_queries = []
+            for table in data_model.upstreamTables or []:
+                for referenced_query in table.referencedByQueries or []:
+                    sql_queries.append(referenced_query.query)
+            return "\n\n".join(sql_queries) or None
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Error processing queries for datamodel [{data_model.id}]: {exc}"
+            )
+        return None
+
     def yield_datamodel(
         self, dashboard_details: TableauDashboard
     ) -> Iterable[Either[CreateDashboardDataModelRequest]]:
@@ -171,6 +188,7 @@ class TableauSource(DashboardServiceSource):
                         dataModelType=DataModelType.TableauDataModel.value,
                         serviceType=DashboardServiceType.Tableau.value,
                         columns=self.get_column_info(data_model),
+                        sql=self._get_datamodel_sql_query(data_model=data_model),
                     )
                     yield Either(right=data_model_request)
                     self.register_record_datamodel(datamodel_request=data_model_request)

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

@@ -12,7 +12,7 @@
 """
 Source connection helper
 """
-
+import traceback
 from typing import Any
 
 from pydantic import BaseModel
@@ -26,7 +26,14 @@ from metadata.generated.schema.security.credentials.gcpValues import (
     SingleProjectId,
 )
 from metadata.ingestion.source.connections import get_connection
+from metadata.ingestion.source.database.bigquery.queries import (
+    BIGQUERY_FOREIGN_CONSTRAINTS,
+    BIGQUERY_TABLE_CONSTRAINTS,
+)
 from metadata.utils.bigquery_utils import get_bigquery_client
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class InspectorWrapper(BaseModel):
@@ -64,3 +71,63 @@ def get_inspector_details(
     inspector = inspect(engine)
 
     return InspectorWrapper(client=client, engine=engine, inspector=inspector)
+
+
+def get_pk_constraint(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    """
+    This function overrides to get primary key constraint
+    """
+    try:
+        table_constraints = connection.engine.execute(
+            BIGQUERY_TABLE_CONSTRAINTS.format(
+                project_id=connection.engine.url.host,
+                schema_name=schema,
+                table_name=table_name,
+            )
+        )
+        col_name = []
+        for table_constraint in table_constraints:
+            col_name.append(table_constraint.column_name)
+        return {"constrained_columns": tuple(col_name)}
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(
+            f"Error while fetching primary key constraint error for table [{schema}.{table_name}]: {exc}"
+        )
+        return {"constrained_columns": []}
+
+
+def get_foreign_keys(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    """
+    This function overrides to get foreign key constraint
+    """
+    try:
+        table_constraints = connection.engine.execute(
+            BIGQUERY_FOREIGN_CONSTRAINTS.format(
+                project_id=connection.engine.url.host,
+                schema_name=schema,
+                table_name=table_name,
+            )
+        )
+        col_name = []
+        for table_constraint in table_constraints:
+            col_name.append(
+                {
+                    "name": table_constraint.name,
+                    "referred_schema": table_constraint.referred_schema,
+                    "referred_table": table_constraint.referred_table,
+                    "constrained_columns": [table_constraint.constrained_columns],
+                    "referred_columns": [table_constraint.referred_columns],
+                }
+            )
+        return col_name
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(
+            f"Error while fetching foreign key constraint error for table [{schema}.{table_name}]: {exc}"
+        )
+        return []

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

@@ -56,7 +56,11 @@ from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_test_connection_fn
-from metadata.ingestion.source.database.bigquery.helper import get_inspector_details
+from metadata.ingestion.source.database.bigquery.helper import (
+    get_foreign_keys,
+    get_inspector_details,
+    get_pk_constraint,
+)
 from metadata.ingestion.source.database.bigquery.models import (
     STORED_PROC_LANGUAGE_MAP,
     BigQueryStoredProcedure,
@@ -194,6 +198,8 @@ def _build_formatted_table_id(table):
 BigQueryDialect._build_formatted_table_id = (  # pylint: disable=protected-access
     _build_formatted_table_id
 )
+BigQueryDialect.get_pk_constraint = get_pk_constraint
+BigQueryDialect.get_foreign_keys = get_foreign_keys
 
 
 class BigquerySource(
@@ -245,7 +251,9 @@ class BigquerySource(
             test_connection_fn(
                 self.metadata, inspector_details.engine, self.service_connection
             )
-            if os.environ[GOOGLE_CREDENTIALS]:
+            # GOOGLE_CREDENTIALS may not have been set,
+            # to avoid key error, we use `get` for dict
+            if os.environ.get(GOOGLE_CREDENTIALS):
                 self.temp_credentials_file_path.append(os.environ[GOOGLE_CREDENTIALS])
 
     def query_table_names_and_types(
@@ -436,7 +444,8 @@ class BigquerySource(
         inspector_details = get_inspector_details(
             database_name=database_name, service_connection=self.service_connection
         )
-        self.temp_credentials_file_path.append(os.environ[GOOGLE_CREDENTIALS])
+        if os.environ.get(GOOGLE_CREDENTIALS):
+            self.temp_credentials_file_path.append(os.environ[GOOGLE_CREDENTIALS])
         self.client = inspector_details.client
         self.engine = inspector_details.engine
         self.inspector = inspector_details.inspector

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

@@ -58,6 +58,28 @@ BIGQUERY_TABLE_AND_TYPE = textwrap.dedent(
     """
 )
 
+BIGQUERY_TABLE_CONSTRAINTS = textwrap.dedent(
+    """
+    SELECT * 
+    FROM `{project_id}`.{schema_name}.INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE 
+    WHERE table_name = '{table_name}' AND constraint_name LIKE '%pk$';
+    """
+)
+
+BIGQUERY_FOREIGN_CONSTRAINTS = textwrap.dedent(
+    """
+    SELECT 
+      c.table_name AS referred_table, 
+      r.table_schema as referred_schema, 
+      r.constraint_name as name, 
+      c.column_name as referred_columns,
+      c.column_name as constrained_columns 
+    FROM `{project_id}`.{schema_name}.INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE c 
+    JOIN `{project_id}`.{schema_name}.INFORMATION_SCHEMA.TABLE_CONSTRAINTS r ON c.constraint_name = r.constraint_name 
+    WHERE r.constraint_type = 'FOREIGN KEY' AND r.table_name='{table_name}';
+    """
+)
+
 BIGQUERY_GET_STORED_PROCEDURES = textwrap.dedent(
     """
 SELECT 

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

@@ -0,0 +1,62 @@
+#  Copyright 2024 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  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.
+"""A client for Google Cloud Bigtable that supports multiple projects."""
+from functools import partial
+from typing import List, Optional, Type
+
+from google import auth
+from google.cloud.bigtable import Client
+
+NoProject = object()
+
+
+class MultiProjectClient:
+    """Google Cloud Client does not support ad-hoc project switching. This class wraps the client and allows
+    switching between projects. If no project is specified, the client will not have a project set and will try
+    to resolve it from ADC.
+    Example usage:
+    ```
+    from google.cloud.bigtable import Client
+    client = MultiProjectClient(Client, project_ids=["project1", "project2"])
+    instances_project1 = client.list_instances("project1")
+    instances_project2 = client.list_instances("project2")
+    """
+
+    def __init__(
+        self,
+        client_class: Type[Client],
+        project_ids: Optional[List[str]] = None,
+        **client_kwargs,
+    ):
+        if project_ids:
+            self.clients = {
+                project_id: client_class(project=project_id, **client_kwargs)
+                for project_id in project_ids
+            }
+        else:
+            self.clients = {NoProject: client_class(**client_kwargs)}
+
+    def project_ids(self):
+        if NoProject in self.clients:
+            _, project_id = auth.default()
+            return [project_id]
+        return list(self.clients.keys())
+
+    def __getattr__(self, client_method):
+        """Return the underlying client method as a partial function so we can inject the project_id."""
+        return partial(self._call, client_method)
+
+    def _call(self, method, project_id, *args, **kwargs):
+        """Call the method on the client for the given project_id. The args and kwargs are passed through."""
+        client = self.clients.get(project_id, self.clients.get(NoProject))
+        if not client:
+            raise ValueError(f"Project {project_id} not found")
+        return getattr(client, method)(*args, **kwargs)