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

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/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

@@ -82,7 +82,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 +129,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

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)

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

@@ -0,0 +1,116 @@
+#  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.
+"""BigTable connection"""
+from typing import List, Optional
+
+from google.cloud.bigtable import Client
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.database.bigTableConnection import (
+    BigTableConnection,
+)
+from metadata.generated.schema.security.credentials.gcpValues import (
+    GcpCredentialsValues,
+    SingleProjectId,
+)
+from metadata.ingestion.connections.test_connections import (
+    SourceConnectionException,
+    test_connection_steps,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.bigtable.client import MultiProjectClient
+from metadata.utils.credentials import set_google_credentials
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+def get_connection(connection: BigTableConnection):
+    set_google_credentials(connection.credentials)
+    project_ids = None
+    if isinstance(connection.credentials.gcpConfig, GcpCredentialsValues):
+        project_ids = (
+            [connection.credentials.gcpConfig.projectId.__root__]
+            if isinstance(connection.credentials.gcpConfig.projectId, SingleProjectId)
+            else connection.credentials.gcpConfig.projectId.__root__
+        )
+    # admin=True is required to list instances and tables
+    return MultiProjectClient(client_class=Client, project_ids=project_ids, admin=True)
+
+
+def get_nested_index(lst: list, index: List[int], default=None):
+    try:
+        for i in index:
+            lst = lst[i]
+        return lst
+    except IndexError:
+        return default
+
+
+class Tester:
+    """
+    A wrapper class that holds state. We need it because the different testing stages
+    are not independent of each other. For example, we need to list instances before we can list
+    """
+
+    def __init__(self, client: MultiProjectClient):
+        self.client = client
+        self.project_id = None
+        self.instance = None
+        self.table = None
+
+    def list_instances(self):
+        self.project_id = list(self.client.clients.keys())[0]
+        instances = list(self.client.list_instances(project_id=self.project_id))
+        self.instance = get_nested_index(instances, [0, 0])
+
+    def list_tables(self):
+        if not self.instance:
+            raise SourceConnectionException(
+                f"No instances found in project {self.project_id}"
+            )
+        tables = list(self.instance.list_tables())
+        self.table = tables[0]
+
+    def get_row(self):
+        if not self.table:
+            raise SourceConnectionException(
+                f"No tables found in project {self.instance.project_id} and instance {self.instance.instance_id}"
+            )
+        self.table.read_rows(limit=1)
+
+
+def test_connection(
+    metadata: OpenMetadata,
+    client: MultiProjectClient,
+    service_connection: BigTableConnection,
+    automation_workflow: Optional[AutomationWorkflow] = None,
+) -> None:
+    """
+    Test connection. This can be executed either as part
+    of a metadata workflow or during an Automation Workflow
+    """
+    tester = Tester(client)
+
+    test_fn = {
+        "GetInstances": tester.list_instances,
+        "GetTables": tester.list_tables,
+        "GetRows": tester.get_row,
+    }
+
+    test_connection_steps(
+        metadata=metadata,
+        test_fn=test_fn,
+        service_type=service_connection.type.value,
+        automation_workflow=automation_workflow,
+    )