openmetadata-ingestion 1.4.0.0rc3__py3-none-any.whl → 1.4.0.0rc4__py3-none-any.whl

metadata/generated/schema/type/votes.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/votes.json
-#   timestamp: 2024-05-17T18:24:48+00:00
+#   timestamp: 2024-05-20T10:39:41+00:00
 
 from __future__ import annotations
 

metadata/ingestion/source/dashboard/quicksight/metadata.py

@@ -11,7 +11,7 @@
 """QuickSight source module"""
 
 import traceback
-from typing import Any, Iterable, List, Optional
+from typing import Iterable, List, Optional
 
 from pydantic import ValidationError
 
@@ -34,7 +34,12 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.dashboard.dashboard_service import DashboardServiceSource
-from metadata.ingestion.source.dashboard.quicksight.models import DataSourceResp
+from metadata.ingestion.source.dashboard.quicksight.models import (
+    DashboardDetail,
+    DashboardResp,
+    DataSourceResp,
+    DescribeDataSourceResponse,
+)
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_chart
 from metadata.utils.logger import ingestion_logger
@@ -109,36 +114,40 @@ class QuicksightSource(DashboardServiceSource):
             dashboard["DashboardId"] for dashboard in dashboard_summary_list
         }
         dashboards = [
-            self.client.describe_dashboard(
-                AwsAccountId=self.aws_account_id, DashboardId=dashboard_id
-            )["Dashboard"]
+            DashboardResp(
+                **self.client.describe_dashboard(
+                    AwsAccountId=self.aws_account_id, DashboardId=dashboard_id
+                )
+            ).Dashboard
             for dashboard_id in dashboard_set
         ]
         return dashboards
 
-    def get_dashboard_name(self, dashboard: dict) -> str:
+    def get_dashboard_name(self, dashboard: DashboardDetail) -> str:
         """
         Get Dashboard Name
         """
-        return dashboard["Name"]
+        return dashboard.Name
 
-    def get_dashboard_details(self, dashboard: dict) -> dict:
+    def get_dashboard_details(self, dashboard: dict) -> DashboardDetail:
         """
         Get Dashboard Details
         """
-        return dashboard
+        return DashboardDetail(**dashboard)
 
     def yield_dashboard(
-        self, dashboard_details: dict
+        self, dashboard_details: DashboardDetail
     ) -> Iterable[Either[CreateDashboardRequest]]:
         """
         Method to Get Dashboard Entity
         """
         dashboard_request = CreateDashboardRequest(
-            name=dashboard_details["DashboardId"],
+            name=dashboard_details.DashboardId,
             sourceUrl=self.dashboard_url,
-            displayName=dashboard_details["Name"],
-            description=dashboard_details["Version"].get("Description"),
+            displayName=dashboard_details.Name,
+            description=dashboard_details.Version.Description
+            if dashboard_details.Version
+            else None,
             charts=[
                 fqn.build(
                     self.metadata,
@@ -155,45 +164,45 @@ class QuicksightSource(DashboardServiceSource):
         self.register_record(dashboard_request=dashboard_request)
 
     def yield_dashboard_chart(
-        self, dashboard_details: Any
+        self, dashboard_details: DashboardDetail
     ) -> Iterable[Either[CreateChartRequest]]:
         """Get chart method"""
         # Each dashboard is guaranteed to have at least one sheet, which represents
         # a chart in the context of QuickSight
         # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/quicksight.html#QuickSight.Client.describe_dashboard
-        charts = dashboard_details["Version"]["Sheets"]
-        for chart in charts:
-            try:
-                if filter_by_chart(
-                    self.source_config.chartFilterPattern, chart["Name"]
-                ):
-                    self.status.filter(chart["Name"], "Chart Pattern not allowed")
-                    continue
+        if dashboard_details.Version:
+            for chart in dashboard_details.Version.Charts or []:
+                try:
+                    if filter_by_chart(
+                        self.source_config.chartFilterPattern, chart.Name
+                    ):
+                        self.status.filter(chart.Name, "Chart Pattern not allowed")
+                        continue
 
-                self.dashboard_url = (
-                    f"https://{self.aws_region}.quicksight.aws.amazon.com/sn/dashboards"
-                    f'/{dashboard_details.get("DashboardId")}'
-                )
-                yield Either(
-                    right=CreateChartRequest(
-                        name=chart["SheetId"],
-                        displayName=chart["Name"],
-                        chartType=ChartType.Other.value,
-                        sourceUrl=self.dashboard_url,
-                        service=self.context.get().dashboard_service,
+                    self.dashboard_url = (
+                        f"https://{self.aws_region}.quicksight.aws.amazon.com/sn/dashboards"
+                        f"/{dashboard_details.DashboardId}"
                     )
-                )
-            except Exception as exc:
-                yield Either(
-                    left=StackTraceError(
-                        name="Chart",
-                        error=f"Error creating chart [{chart}]: {exc}",
-                        stackTrace=traceback.format_exc(),
+                    yield Either(
+                        right=CreateChartRequest(
+                            name=chart.ChartId,
+                            displayName=chart.Name,
+                            chartType=ChartType.Other.value,
+                            sourceUrl=self.dashboard_url,
+                            service=self.context.get().dashboard_service,
+                        )
+                    )
+                except Exception as exc:
+                    yield Either(
+                        left=StackTraceError(
+                            name="Chart",
+                            error=f"Error creating chart [{chart}]: {exc}",
+                            stackTrace=traceback.format_exc(),
+                        )
                     )
-                )
 
     def yield_dashboard_lineage_details(  # pylint: disable=too-many-locals
-        self, dashboard_details: dict, db_service_name: str
+        self, dashboard_details: DashboardDetail, db_service_name: str
     ) -> Iterable[Either[AddLineageRequest]]:
         """
         Get lineage between dashboard and data sources
@@ -208,28 +217,26 @@ class QuicksightSource(DashboardServiceSource):
             )
             dataset_ids = {
                 dataset["DataSetId"]
-                for dataset in data_set_summary_list
-                if dataset.get("Arn") in dashboard_details["Version"]["DataSetArns"]
+                for dataset in data_set_summary_list or []
+                if dataset.get("Arn") in dashboard_details.Version.DataSetArns
             }
 
-            for dataset_id in dataset_ids:
-                for data_source in list(
-                    self.client.describe_data_set(
-                        AwsAccountId=self.aws_account_id, DataSetId=dataset_id
-                    )["DataSet"]["PhysicalTableMap"].values()
+            for dataset_id in dataset_ids or []:
+                for data_source in (
+                    list(
+                        self.client.describe_data_set(
+                            AwsAccountId=self.aws_account_id, DataSetId=dataset_id
+                        )["DataSet"]["PhysicalTableMap"].values()
+                    )
+                    or []
                 ):
                     try:
                         if not data_source.get("RelationalTable"):
                             raise KeyError(
                                 f"We currently don't support lineage to {list(data_source.keys())}"
                             )
-                        data_source_relational_table = data_source["RelationalTable"]
                         data_source_resp = DataSourceResp(
-                            datasource_arn=data_source_relational_table[
-                                "DataSourceArn"
-                            ],
-                            schema_name=data_source_relational_table["Schema"],
-                            table_name=data_source_relational_table["Name"],
+                            **data_source["RelationalTable"]
                         )
                     except (KeyError, ValidationError) as err:
                         yield Either(
@@ -257,42 +264,47 @@ class QuicksightSource(DashboardServiceSource):
 
                     data_source_ids = [
                         data_source_arn["DataSourceId"]
-                        for data_source_arn in data_source_summary_list
+                        for data_source_arn in data_source_summary_list or []
                         if data_source_arn["Arn"] in data_source_resp.datasource_arn
                     ]
 
-                    for data_source_id in data_source_ids:
-                        data_source_dict = self.client.describe_data_source(
-                            AwsAccountId=self.aws_account_id,
-                            DataSourceId=data_source_id,
-                        )["DataSource"]["DataSourceParameters"]
-                        for db in data_source_dict.keys():
-                            from_fqn = fqn.build(
-                                self.metadata,
-                                entity_type=Table,
-                                service_name=db_service_name,
-                                database_name=data_source_dict[db]["Database"],
-                                schema_name=schema_name,
-                                table_name=table_name,
-                                skip_es_search=True,
-                            )
-                            from_entity = self.metadata.get_by_name(
-                                entity=Table,
-                                fqn=from_fqn,
-                            )
-                            to_fqn = fqn.build(
-                                self.metadata,
-                                entity_type=Dashboard,
-                                service_name=self.config.serviceName,
-                                dashboard_name=dashboard_details["DashboardId"],
-                            )
-                            to_entity = self.metadata.get_by_name(
-                                entity=Dashboard,
-                                fqn=to_fqn,
-                            )
-                            yield self._get_add_lineage_request(
-                                to_entity=to_entity, from_entity=from_entity
+                    for data_source_id in data_source_ids or []:
+                        data_source_resp = DescribeDataSourceResponse(
+                            **self.client.describe_data_source(
+                                AwsAccountId=self.aws_account_id,
+                                DataSourceId=data_source_id,
                             )
+                        ).DataSource
+                        if data_source_resp and data_source_resp.DataSourceParameters:
+                            data_source_dict = data_source_resp.DataSourceParameters
+                            for db in data_source_dict.keys() or []:
+                                from_fqn = fqn.build(
+                                    self.metadata,
+                                    entity_type=Table,
+                                    service_name=db_service_name,
+                                    database_name=data_source_dict[db].get("Database"),
+                                    schema_name=schema_name,
+                                    table_name=table_name,
+                                    skip_es_search=True,
+                                )
+                                from_entity = self.metadata.get_by_name(
+                                    entity=Table,
+                                    fqn=from_fqn,
+                                )
+                                to_fqn = fqn.build(
+                                    self.metadata,
+                                    entity_type=Dashboard,
+                                    service_name=self.config.serviceName,
+                                    dashboard_name=dashboard_details.DashboardId,
+                                )
+                                to_entity = self.metadata.get_by_name(
+                                    entity=Dashboard,
+                                    fqn=to_fqn,
+                                )
+                                if from_entity is not None and to_entity is not None:
+                                    yield self._get_add_lineage_request(
+                                        to_entity=to_entity, from_entity=from_entity
+                                    )
         except Exception as exc:  # pylint: disable=broad-except
             yield Either(
                 left=StackTraceError(

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

@@ -12,11 +12,50 @@
 Pydantic Model to validate Quick Sight responses
 """
 
+from typing import List, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 class DataSourceResp(BaseModel):
-    datasource_arn: str
-    schema_name: str
-    table_name: str
+    datasource_arn: str = Field(alias="DataSourceArn")
+    schema_name: str = Field(alias="Schema")
+    table_name: str = Field(alias="Name")
+
+
+class VersionSheet(BaseModel):
+    ChartId: Optional[str] = Field(alias="SheetId")
+    Name: Optional[str]
+
+
+class DashboardVersion(BaseModel):
+    Status: Optional[str]
+    Arn: Optional[str]
+    SourceEntityArn: Optional[str]
+    DataSetArns: Optional[List]
+    Description: Optional[str]
+    Charts: Optional[List[VersionSheet]] = Field(alias="Sheets")
+
+
+class DashboardDetail(BaseModel):
+    DashboardId: str
+    Arn: Optional[str]
+    Name: str
+    Version: Optional[DashboardVersion]
+
+
+class DashboardResp(BaseModel):
+    Dashboard: DashboardDetail
+    Status: Optional[int]
+    RequestId: Optional[str]
+
+
+class DataSource(BaseModel):
+    DataSourceId: str
+    DataSourceParameters: Optional[dict]
+
+
+class DescribeDataSourceResponse(BaseModel):
+    DataSource: Optional[DataSource]
+    RequestId: Optional[str]
+    Status: Optional[int]

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

@@ -41,7 +41,10 @@ def get_connection_url(connection: Union[AzureSQLConnection, MssqlConnection]) -
     Build the connection URL
     """
 
-    if connection.authenticationMode:
+    if (
+        hasattr(connection.authenticationMode, "authentication")
+        and connection.authenticationMode.authentication is not None
+    ):
         connection_string = f"Driver={connection.driver};Server={connection.hostPort};Database={connection.database};"
         connection_string += f"Uid={connection.username};"
         if (

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

@@ -630,7 +630,9 @@ class BigquerySource(
                 self.connection, table_name, schema_name
             )
             schema_definition = (
-                str(schema_definition) if schema_definition is not None else None
+                str(schema_definition.strip())
+                if schema_definition is not None
+                else None
             )
             return schema_definition
         except NotImplementedError:

metadata/ingestion/source/database/clickhouse/metadata.py

@@ -161,7 +161,9 @@ class ClickhouseSource(CommonDbSourceService):
                     self.connection, table_name, schema_name
                 )
             schema_definition = (
-                str(schema_definition) if schema_definition is not None else None
+                str(schema_definition.strip())
+                if schema_definition is not None
+                else None
             )
             return schema_definition
 

metadata/ingestion/source/database/common_db_source.py

@@ -366,16 +366,19 @@ class CommonDbSourceService(
         Get the DDL statement or View Definition for a table
         """
         try:
+            schema_definition = None
             if table_type in (TableType.View, TableType.MaterializedView):
                 schema_definition = inspector.get_view_definition(
                     table_name, schema_name
                 )
-            else:
+            elif hasattr(inspector, "get_table_ddl"):
                 schema_definition = inspector.get_table_ddl(
                     self.connection, table_name, schema_name
                 )
             schema_definition = (
-                str(schema_definition) if schema_definition is not None else None
+                str(schema_definition).strip()
+                if schema_definition is not None
+                else None
             )
             return schema_definition
 

metadata/ingestion/source/database/hive/metadata.py

@@ -108,7 +108,9 @@ class HiveSource(CommonDbSourceService):
         try:
             schema_definition = inspector.get_view_definition(table_name, schema_name)
             schema_definition = (
-                str(schema_definition) if schema_definition is not None else None
+                str(schema_definition).strip()
+                if schema_definition is not None
+                else None
             )
             return schema_definition
 

metadata/ingestion/source/database/oracle/metadata.py

@@ -163,7 +163,9 @@ class OracleSource(StoredProcedureMixin, CommonDbSourceService):
                 schema_definition = definition_fn(table_name, schema_name)
 
             schema_definition = (
-                str(schema_definition) if schema_definition is not None else None
+                str(schema_definition.strip())
+                if schema_definition is not None
+                else None
             )
             return schema_definition
 

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

@@ -97,7 +97,7 @@ WITH SP_HISTORY AS (SELECT
     TO_TIMESTAMP(LAST_LOAD_TIME, 'YYYY-MM-DD HH24:MI:SS') + NUMTODSINTERVAL(ELAPSED_TIME / 1000, 'SECOND') AS end_time,
     PARSING_SCHEMA_NAME as user_name
   FROM gv$sql
-  WHERE UPPER(sql_text) LIKE 'CALL%%'
+  WHERE UPPER(sql_text) LIKE '%%CALL%%' or UPPER(sql_text) LIKE '%%BEGIN%%'
   AND TO_TIMESTAMP(FIRST_LOAD_TIME, 'YYYY-MM-DD HH24:MI:SS') >= TO_TIMESTAMP('{start_date}', 'YYYY-MM-DD HH24:MI:SS')
  ),
  Q_HISTORY AS (SELECT
@@ -114,7 +114,7 @@ WITH SP_HISTORY AS (SELECT
       PARSING_SCHEMA_NAME AS SCHEMA_NAME,
       NULL AS DATABASE_NAME
     FROM gv$sql
-    WHERE UPPER(sql_text) NOT LIKE '%CALL%'
+    WHERE UPPER(sql_text) NOT LIKE '%%CALL%%' AND UPPER(sql_text) NOT LIKE '%%BEGIN%%'
       AND SQL_FULLTEXT NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
       AND SQL_FULLTEXT NOT LIKE '/* {{"app": "dbt", %%}} */%%'
       AND TO_TIMESTAMP(FIRST_LOAD_TIME, 'YYYY-MM-DD HH24:MI:SS') 
@@ -135,6 +135,7 @@ JOIN Q_HISTORY Q
   ON Q.start_time between SP.start_time and SP.end_time
   AND Q.end_time between SP.start_time and SP.end_time
   AND Q.user_name = SP.user_name
+  AND Q.QUERY_TYPE <> 'SELECT'
 ORDER BY PROCEDURE_START_TIME DESC
 """
 )

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

@@ -21,10 +21,14 @@ from sqlalchemy.engine import Engine
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
 )
+from metadata.generated.schema.entity.services.connections.database.sapHana.sapHanaHDBConnection import (
+    SapHanaHDBConnection,
+)
+from metadata.generated.schema.entity.services.connections.database.sapHana.sapHanaSQLConnection import (
+    SapHanaSQLConnection,
+)
 from metadata.generated.schema.entity.services.connections.database.sapHanaConnection import (
-    HdbUserStoreConnection,
     SapHanaConnection,
-    SqlConnection,
 )
 from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
@@ -46,7 +50,7 @@ def get_database_connection_url(connection: SapHanaConnection) -> str:
 
     conn = connection.connection
 
-    if not isinstance(conn, SqlConnection):
+    if not isinstance(conn, SapHanaSQLConnection):
         raise ValueError("Database Connection requires the SQL connection details")
 
     url = (
@@ -75,7 +79,7 @@ def get_hdb_connection_url(connection: SapHanaConnection) -> str:
     Build the SQLConnection URL for the database connection
     """
 
-    if not isinstance(connection.connection, HdbUserStoreConnection):
+    if not isinstance(connection.connection, SapHanaHDBConnection):
         raise ValueError("Database Connection requires the SQL connection details")
 
     return f"{connection.scheme.value}://userkey={connection.connection.userKey}"
@@ -86,14 +90,14 @@ def get_connection(connection: SapHanaConnection) -> Engine:
     Create connection
     """
 
-    if isinstance(connection.connection, SqlConnection):
+    if isinstance(connection.connection, SapHanaSQLConnection):
         return create_generic_db_connection(
             connection=connection,
             get_connection_url_fn=get_database_connection_url,
             get_connection_args_fn=get_connection_args_common,
         )
 
-    if isinstance(connection.connection, HdbUserStoreConnection):
+    if isinstance(connection.connection, SapHanaHDBConnection):
         return create_generic_db_connection(
             connection=connection,
             get_connection_url_fn=get_hdb_connection_url,
@@ -128,7 +132,7 @@ def _build_test_fn_dict(
                 inspector_fn(schema)
                 break
 
-    if isinstance(service_connection.connection, SqlConnection):
+    if isinstance(service_connection.connection, SapHanaSQLConnection):
         return {
             "CheckAccess": partial(test_connection_engine_step, engine),
             "GetSchemas": partial(execute_inspector_func, engine, "get_schema_names"),
@@ -136,7 +140,7 @@ def _build_test_fn_dict(
             "GetViews": partial(custom_executor, engine, "get_view_names"),
         }
 
-    if isinstance(service_connection.connection, HdbUserStoreConnection):
+    if isinstance(service_connection.connection, SapHanaHDBConnection):
         return {
             "CheckAccess": partial(test_connection_engine_step, engine),
             "GetSchemas": partial(execute_inspector_func, engine, "get_schema_names"),

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

@@ -49,7 +49,10 @@ def get_connection_url(connection: TrinoConnection) -> str:
     """
     url = f"{connection.scheme.value}://"
     if connection.username:
-        url += f"{quote_plus(connection.username)}"
+        # we need to encode twice because trino dialect internally
+        # url decodes the username and if there is an special char in username
+        # it will fail to authenticate
+        url += f"{quote_plus(quote_plus(connection.username))}"
         if (
             isinstance(connection.authType, basicAuth.BasicAuth)
             and connection.authType.password

metadata/ingestion/source/messaging/kafka/connection.py

@@ -12,6 +12,7 @@
 """
 Source connection handler
 """
+from copy import deepcopy
 from dataclasses import dataclass
 from typing import Optional, Union
 
@@ -64,45 +65,39 @@ def get_connection(
     """
     Create connection
     """
+    consumer_config = deepcopy(connection.consumerConfig) or {}
+    schema_registry_config = deepcopy(connection.schemaRegistryConfig) or {}
     if connection.saslUsername or connection.saslPassword or connection.saslMechanism:
-        connection.consumerConfig = connection.consumerConfig or {}
         if connection.saslUsername:
-            connection.consumerConfig["sasl.username"] = connection.saslUsername
+            consumer_config["sasl.username"] = connection.saslUsername
         if connection.saslPassword:
-            connection.consumerConfig[
+            consumer_config[
                 "sasl.password"
             ] = connection.saslPassword.get_secret_value()
         if connection.saslMechanism:
-            connection.consumerConfig["sasl.mechanism"] = connection.saslMechanism.value
+            consumer_config["sasl.mechanism"] = connection.saslMechanism.value
 
         if (
             connection.consumerConfig.get("security.protocol") is None
             and connection.securityProtocol
         ):
-            connection.consumerConfig[
-                "security.protocol"
-            ] = connection.securityProtocol.value
+            consumer_config["security.protocol"] = connection.securityProtocol.value
 
     if connection.basicAuthUserInfo:
-        connection.schemaRegistryConfig = connection.schemaRegistryConfig or {}
-        connection.schemaRegistryConfig[
+        schema_registry_config[
             "basic.auth.user.info"
         ] = connection.basicAuthUserInfo.get_secret_value()
 
-    admin_client_config = connection.consumerConfig
+    admin_client_config = consumer_config
     admin_client_config["bootstrap.servers"] = connection.bootstrapServers
     admin_client = AdminClient(admin_client_config)
 
     schema_registry_client = None
     consumer_client = None
     if connection.schemaRegistryURL:
-        connection.schemaRegistryConfig["url"] = connection.schemaRegistryURL
-        schema_registry_client = SchemaRegistryClient(connection.schemaRegistryConfig)
-        connection.schemaRegistryConfig["url"] = str(connection.schemaRegistryURL)
-        consumer_config = {
-            **connection.consumerConfig,
-            "bootstrap.servers": connection.bootstrapServers,
-        }
+        schema_registry_config["url"] = connection.schemaRegistryURL
+        schema_registry_client = SchemaRegistryClient(schema_registry_config)
+        consumer_config["bootstrap.servers"] = connection.bootstrapServers
         if "group.id" not in consumer_config:
             consumer_config["group.id"] = "openmetadata-consumer"
         if "auto.offset.reset" not in consumer_config: