openmetadata-ingestion 1.3.4.0__py3-none-any.whl → 1.4.0.0__py3-none-any.whl

metadata/ingestion/source/database/redshift/utils.py

@@ -33,6 +33,7 @@ from metadata.ingestion.source.database.redshift.queries import (
     REDSHIFT_GET_SCHEMA_COLUMN_INFO,
     REDSHIFT_TABLE_COMMENTS,
 )
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.sqlalchemy_utils import get_table_comment_wrapper
 
 sa_version = Version(sa.__version__)
@@ -45,6 +46,7 @@ ischema_names.update(REDSHIFT_ISCHEMA_NAMES)
 
 
 # pylint: disable=protected-access
+@calculate_execution_time()
 @reflection.cache
 def get_columns(self, connection, table_name, schema=None, **kw):
     """
@@ -80,6 +82,7 @@ def get_columns(self, connection, table_name, schema=None, **kw):
     return columns
 
 
+@calculate_execution_time()
 def _get_column_info(self, *args, **kwargs):
     """
     Get column info
@@ -112,6 +115,7 @@ def _get_column_info(self, *args, **kwargs):
     return column_info
 
 
+@calculate_execution_time()
 @reflection.cache
 def _get_schema_column_info(
     self, connection, schema=None, **kw
@@ -130,13 +134,12 @@ def _get_schema_column_info(
     """
     schema_clause = f"AND schema = '{schema if schema else ''}'"
     all_columns = defaultdict(list)
-    with connection.connect() as cnct:
-        result = cnct.execute(
-            REDSHIFT_GET_SCHEMA_COLUMN_INFO.format(schema_clause=schema_clause)
-        )
-        for col in result:
-            key = RelationKey(col.table_name, col.schema, connection)
-            all_columns[key].append(col)
+    result = connection.execute(
+        REDSHIFT_GET_SCHEMA_COLUMN_INFO.format(schema_clause=schema_clause)
+    )
+    for col in result:
+        key = RelationKey(col.table_name, col.schema, connection)
+        all_columns[key].append(col)
     return dict(all_columns)
 
 
@@ -270,6 +273,7 @@ def _get_charlen(format_type):
     return charlen
 
 
+@calculate_execution_time()
 @reflection.cache
 def _get_pg_column_info(  # pylint: disable=too-many-locals,too-many-arguments, unused-argument
     self,
@@ -353,6 +357,7 @@ def _get_pg_column_info(  # pylint: disable=too-many-locals,too-many-arguments,
     return column_info
 
 
+@calculate_execution_time()
 @reflection.cache
 def get_table_comment(
     self, connection, table_name, schema=None, **kw  # pylint: disable=unused-argument
@@ -366,6 +371,7 @@ def get_table_comment(
     )
 
 
+@calculate_execution_time()
 @reflection.cache
 def _get_all_relation_info(self, connection, **kw):  # pylint: disable=unused-argument
     # pylint: disable=consider-using-f-string

metadata/ingestion/source/database/salesforce/metadata.py

@@ -81,7 +81,9 @@ class SalesforceSource(DatabaseServiceSource):
         self.database_source_state = set()
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: SalesforceConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, SalesforceConnection):
@@ -112,7 +114,7 @@ class SalesforceSource(DatabaseServiceSource):
         yield Either(
             right=CreateDatabaseRequest(
                 name=database_name,
-                service=self.context.database_service,
+                service=self.context.get().database_service,
             )
         )
 
@@ -135,8 +137,8 @@ class SalesforceSource(DatabaseServiceSource):
                 database=fqn.build(
                     metadata=self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
                 ),
             )
         )
@@ -150,7 +152,7 @@ class SalesforceSource(DatabaseServiceSource):
 
         :return: tables or views, depending on config
         """
-        schema_name = self.context.database_schema
+        schema_name = self.context.get().database_schema
 
         try:
             if self.service_connection.sobjectName:
@@ -165,9 +167,9 @@ class SalesforceSource(DatabaseServiceSource):
                     table_fqn = fqn.build(
                         self.metadata,
                         entity_type=Table,
-                        service_name=self.context.database_service,
-                        database_name=self.context.database,
-                        schema_name=self.context.database_schema,
+                        service_name=self.context.get().database_service,
+                        database_name=self.context.get().database,
+                        schema_name=self.context.get().database_schema,
                         table_name=table_name,
                     )
                     if filter_by_table(
@@ -240,9 +242,9 @@ class SalesforceSource(DatabaseServiceSource):
                 databaseSchema=fqn.build(
                     metadata=self.metadata,
                     entity_type=DatabaseSchema,
-                    service_name=self.context.database_service,
-                    database_name=self.context.database,
-                    schema_name=self.context.database_schema,
+                    service_name=self.context.get().database_service,
+                    database_name=self.context.get().database,
+                    schema_name=self.context.get().database_schema,
                 ),
                 sourceUrl=self.get_source_url(
                     table_name=table_name,

metadata/ingestion/source/database/sample_data.py

@@ -281,6 +281,11 @@ class SampleDataSource(
             entity=DatabaseService, config=WorkflowSource(**self.database_service_json)
         )
 
+        self.glue_database_service = self.metadata.get_service_or_create(
+            entity=DatabaseService,
+            config=WorkflowSource(**self.glue_database_service_json),
+        )
+
         self.kafka_service_json = json.load(
             open(  # pylint: disable=consider-using-with
                 sample_data_folder + "/topics/service.json",
@@ -533,7 +538,9 @@ class SampleDataSource(
         )
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: CustomDatabaseConnection = config.serviceConnection.__root__.config
@@ -549,8 +556,8 @@ class SampleDataSource(
     def _iter(self, *_, **__) -> Iterable[Entity]:
         yield from self.ingest_teams()
         yield from self.ingest_users()
-        yield from self.ingest_glue()
         yield from self.ingest_tables()
+        yield from self.ingest_glue()
         yield from self.ingest_stored_procedures()
         yield from self.ingest_topics()
         yield from self.ingest_charts()
@@ -604,9 +611,9 @@ class SampleDataSource(
         """Ingest Sample Data for glue database source"""
 
         db = CreateDatabaseRequest(
-            name=self.database["name"],
-            description=self.database["description"],
-            service=self.database_service.fullyQualifiedName,
+            name=self.glue_database["name"],
+            description=self.glue_database["description"],
+            service=self.glue_database_service.fullyQualifiedName,
         )
 
         yield Either(right=db)
@@ -614,7 +621,7 @@ class SampleDataSource(
         database_entity = fqn.build(
             self.metadata,
             entity_type=Database,
-            service_name=self.database_service.name.__root__,
+            service_name=self.glue_database_service.fullyQualifiedName.__root__,
             database_name=db.name.__root__,
         )
 
@@ -622,8 +629,8 @@ class SampleDataSource(
             entity=Database, fqn=database_entity
         )
         schema = CreateDatabaseSchemaRequest(
-            name=self.database_schema["name"],
-            description=self.database_schema["description"],
+            name=self.glue_database_schema["name"],
+            description=self.glue_database_schema["description"],
             database=database_object.fullyQualifiedName,
         )
         yield Either(right=schema)
@@ -631,7 +638,7 @@ class SampleDataSource(
         database_schema_entity = fqn.build(
             self.metadata,
             entity_type=DatabaseSchema,
-            service_name=self.database_service.name.__root__,
+            service_name=self.glue_database_service.fullyQualifiedName.__root__,
             database_name=db.name.__root__,
             schema_name=schema.name.__root__,
         )
@@ -651,6 +658,29 @@ class SampleDataSource(
             )
             yield Either(right=table_request)
 
+        database_schema_entity = fqn.build(
+            self.metadata,
+            entity_type=DatabaseSchema,
+            service_name=self.database_service.fullyQualifiedName.__root__,
+            database_name=self.database["name"],
+            schema_name=self.database_schema["name"],
+        )
+
+        database_schema_object = self.metadata.get_by_name(
+            entity=DatabaseSchema, fqn=database_schema_entity
+        )
+
+        for table in self.glue_tables["tables"]:
+            table_request = CreateTableRequest(
+                name=table["name"],
+                description=table["description"],
+                columns=table["columns"],
+                databaseSchema=database_schema_object.fullyQualifiedName,
+                tableConstraints=table.get("tableConstraints"),
+                tableType=table["tableType"],
+            )
+            yield Either(right=table_request)
+
     def ingest_tables(self) -> Iterable[Either[Entity]]:
         """Ingest Sample Tables"""
 
@@ -703,6 +733,7 @@ class SampleDataSource(
                 tableType=table["tableType"],
                 tableConstraints=table.get("tableConstraints"),
                 tags=table["tags"],
+                schemaDefinition=table.get("schemaDefinition"),
             )
 
             yield Either(right=table_and_db)
@@ -1491,6 +1522,19 @@ class SampleDataSource(
                         test_case_name=case.fullyQualifiedName.__root__,
                     )
                     yield Either(right=test_case_result_req)
+            if test_case_results.get("failedRowsSample"):
+                self.metadata.ingest_failed_rows_sample(
+                    case,
+                    TableData(
+                        rows=test_case_results["failedRowsSample"]["rows"],
+                        columns=test_case_results["failedRowsSample"]["columns"],
+                    ),
+                )
+            if test_case_results.get("inspectionQuery"):
+                self.metadata.ingest_inspection_query(
+                    case,
+                    test_case_results["inspectionQuery"],
+                )
 
     def ingest_data_insights(self) -> Iterable[Either[OMetaDataInsightSample]]:
         """Iterate over all the data insights and ingest them"""
@@ -1524,9 +1568,6 @@ class SampleDataSource(
     def ingest_life_cycle(self) -> Iterable[Either[OMetaLifeCycleData]]:
         """Iterate over all the life cycle data and ingest them"""
         for table_life_cycle in self.life_cycle_data["lifeCycleData"]:
-            table = self.metadata.get_by_name(
-                entity=Table, fqn=table_life_cycle["fqn"], fields=["lifeCycle"]
-            )
             life_cycle = table_life_cycle["lifeCycle"]
             life_cycle_data = LifeCycle()
             life_cycle_data.created = AccessDetails(
@@ -1581,7 +1622,9 @@ class SampleDataSource(
                 )
 
             life_cycle_request = OMetaLifeCycleData(
-                entity=table, life_cycle=life_cycle_data
+                entity=Table,
+                entity_fqn=table_life_cycle["fqn"],
+                life_cycle=life_cycle_data,
             )
             yield Either(right=life_cycle_request)
 

metadata/ingestion/source/database/sample_usage.py

@@ -69,7 +69,9 @@ class SampleUsageSource(UsageSource):
         )
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         """Create class instance"""
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: CustomDatabaseConnection = config.serviceConnection.__root__.config

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/saphana/metadata.py

@@ -11,9 +11,7 @@
 """
 SAP Hana source module
 """
-from typing import Iterable
-
-from sqlalchemy import inspect
+from typing import Iterable, Optional
 
 from metadata.generated.schema.entity.services.connections.database.sapHanaConnection import (
     SapHanaConnection,
@@ -36,7 +34,9 @@ class SaphanaSource(CommonDbSourceService):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: SapHanaConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, SapHanaConnection):
@@ -49,7 +49,8 @@ class SaphanaSource(CommonDbSourceService):
         """
         Check if the db is configured, or query the name
         """
-        self.inspector = inspect(self.engine)
+        self._connection_map = {}  # Lazy init as well
+        self._inspector_map = {}
 
         if getattr(self.service_connection.connection, "database"):
             yield self.service_connection.connection.database

metadata/ingestion/source/database/sas/metadata.py

@@ -122,7 +122,12 @@ class SasSource(
         self.database_schemas = None
 
     @classmethod
-    def create(cls, config_dict: dict, metadata: OpenMetadata):
+    def create(
+        cls,
+        config_dict: dict,
+        metadata: OpenMetadata,
+        pipeline_name: Optional[str] = None,
+    ):
         logger.info(f"running create {config_dict}")
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: SASConnection = config.serviceConnection.__root__.config
@@ -853,7 +858,7 @@ class SasSource(
         yield Either(
             right=CreateDatabaseRequest(
                 name=database_name,
-                service=self.context.database_service,
+                service=self.context.get().database_service,
             )
         )
 
@@ -871,7 +876,7 @@ class SasSource(
                 database=fqn.build(
                     metadata=self.metadata,
                     entity_type=Database,
-                    service_name=self.context.database_service,
+                    service_name=self.context.get().database_service,
                     database_name=schema_name[0],
                 ),
             )

metadata/ingestion/source/database/singlestore/metadata.py

@@ -11,6 +11,8 @@
 """
 Singlestore source ingestion
 """
+from typing import Optional
+
 from sqlalchemy.dialects.mysql.base import ischema_names
 from sqlalchemy.dialects.mysql.reflection import MySQLTableDefinitionParser
 
@@ -40,7 +42,9 @@ class SinglestoreSource(CommonDbSourceService):
     """
 
     @classmethod
-    def create(cls, config_dict, metadata: OpenMetadata):
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
         config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
         connection: SingleStoreConnection = config.serviceConnection.__root__.config
         if not isinstance(connection, SingleStoreConnection):