openmetadata-ingestion 1.3.0.0.dev0__py3-none-any.whl → 1.3.1.0__py3-none-any.whl

metadata/ingestion/source/database/extended_sample_data.py

@@ -24,6 +24,10 @@ from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequ
 from metadata.generated.schema.api.data.createDatabaseSchema import (
     CreateDatabaseSchemaRequest,
 )
+from metadata.generated.schema.api.data.createGlossary import CreateGlossaryRequest
+from metadata.generated.schema.api.data.createGlossaryTerm import (
+    CreateGlossaryTermRequest,
+)
 from metadata.generated.schema.api.data.createTable import CreateTableRequest
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.dashboard import Dashboard
@@ -91,10 +95,20 @@ class ExtendedSampleDataSource(Source):  # pylint: disable=too-many-instance-att
         self.store_table_fqn = set()
         self.store_data_model_fqn = []
         self.store_dashboard_fqn = []
+        self.main_glossary = None
+        self.glossary_term_list = []
 
         sample_data_folder = self.service_connection.connectionOptions.__root__.get(
             "sampleDataFolder"
         )
+        self.include_glossary = self.service_connection.connectionOptions.__root__.get(
+            "includeGlossary"
+        )
+        self.include_lineage_stress_testing = (
+            self.service_connection.connectionOptions.__root__.get(
+                "includeLineageStressTesting"
+            )
+        )
         extneded_sample_data_folder = (
             self.service_connection.connectionOptions.__root__.get(
                 "extendedSampleDataFolder"
@@ -225,43 +239,71 @@ class ExtendedSampleDataSource(Source):  # pylint: disable=too-many-instance-att
             )
         )
 
-    def generate_sample_data(self):  # pylint: disable=too-many-locals
+    def generate_sample_data(
+        self,
+    ):  # pylint: disable=too-many-locals,too-many-statements
         """
         Generate sample data for dashboard and database service,
         with lineage between them, having long names, special characters and description
         """
-
-        db = self.create_database_request("extended_sample_data", self.generate_text())
-        yield Either(right=db)
-
-        schema = self.create_database_schema_request(
-            "extended_sample_database_schema", self.generate_text(), db
-        )
-        yield Either(right=schema)
-        for table in self.extended_tables["tables"]:
-            text = self.generate_text()
-            table_request = self.create_table_request(
-                table["name"], text, schema, table
-            )
-            yield Either(right=table_request)
-            downstream_node_fqn_table = fqn.build(
-                self.metadata,
-                entity_type=Table,
-                service_name=self.database_service.name.__root__,
-                database_name=db.name.__root__,
-                schema_name=schema.name.__root__,
-                table_name=table_request.name.__root__,
+        if self.include_glossary:
+            yield from self.create_glossary()
+            yield from self.create_glossary_term()
+        if self.include_lineage_stress_testing:
+            db = self.create_database_request(
+                "extended_sample_data", self.generate_text()
             )
-            to_table = self.metadata.get_by_name(
-                entity=Table, fqn=downstream_node_fqn_table
-            )
-            main_table = to_table
-            self.store_table_fqn.add(downstream_node_fqn_table)
+            yield Either(right=db)
 
+            schema = self.create_database_schema_request(
+                "extended_sample_database_schema", self.generate_text(), db
+            )
+            yield Either(right=schema)
+            for table in self.extended_tables["tables"]:
+                text = self.generate_text()
+                table_request = self.create_table_request(
+                    table["name"], text, schema, table
+                )
+                yield Either(right=table_request)
+                downstream_node_fqn_table = fqn.build(
+                    self.metadata,
+                    entity_type=Table,
+                    service_name=self.database_service.name.__root__,
+                    database_name=db.name.__root__,
+                    schema_name=schema.name.__root__,
+                    table_name=table_request.name.__root__,
+                )
+                to_table = self.metadata.get_by_name(
+                    entity=Table, fqn=downstream_node_fqn_table
+                )
+                main_table = to_table
+                self.store_table_fqn.add(downstream_node_fqn_table)
+
+                for _ in range(40):
+                    # 40 Dynamic Lineage Depths
+                    name = self.generate_name()
+                    table_request = self.create_table_request(name, text, schema, table)
+                    yield Either(right=table_request)
+                    upstream_node_fqn_table = fqn.build(
+                        self.metadata,
+                        entity_type=Table,
+                        service_name=self.database_service.name.__root__,
+                        database_name=db.name.__root__,
+                        schema_name=schema.name.__root__,
+                        table_name=table_request.name.__root__,
+                    )
+                    from_table = self.metadata.get_by_name(
+                        entity=Table, fqn=upstream_node_fqn_table
+                    )
+                    yield from self.create_depth_nodes(
+                        from_table=from_table, to_table=to_table
+                    )
+                    to_table = from_table
             for _ in range(40):
-                # 40 Dynamic Lineage Depths
                 name = self.generate_name()
-                table_request = self.create_table_request(name, text, schema, table)
+                table_request = self.create_table_request(
+                    name, text, schema, self.extended_tables["tables"][-1]
+                )
                 yield Either(right=table_request)
                 upstream_node_fqn_table = fqn.build(
                     self.metadata,
@@ -275,127 +317,75 @@ class ExtendedSampleDataSource(Source):  # pylint: disable=too-many-instance-att
                     entity=Table, fqn=upstream_node_fqn_table
                 )
                 yield from self.create_depth_nodes(
-                    from_table=from_table, to_table=to_table
+                    from_table=from_table, to_table=main_table
                 )
-                to_table = from_table
-        for _ in range(40):
-            name = self.generate_name()
-            table_request = self.create_table_request(
-                name, text, schema, self.extended_tables["tables"][-1]
-            )
-            yield Either(right=table_request)
-            upstream_node_fqn_table = fqn.build(
-                self.metadata,
-                entity_type=Table,
-                service_name=self.database_service.name.__root__,
-                database_name=db.name.__root__,
-                schema_name=schema.name.__root__,
-                table_name=table_request.name.__root__,
-            )
-            from_table = self.metadata.get_by_name(
-                entity=Table, fqn=upstream_node_fqn_table
-            )
-            yield from self.create_depth_nodes(
-                from_table=from_table, to_table=main_table
-            )
 
-        for table in self.tables["tables"]:
-            text = self.generate_text()
-            table_request = self.create_table_request(
-                table["name"], text, schema, table
-            )
-            yield Either(right=table_request)
-            table_entity_fqn = fqn.build(
-                self.metadata,
-                entity_type=Table,
-                service_name=self.database_service.name.__root__,
-                database_name=db.name.__root__,
-                schema_name=schema.name.__root__,
-                table_name=table_request.name.__root__,
-            )
-            from_table = self.metadata.get_by_name(entity=Table, fqn=table_entity_fqn)
-
-            yield from self.create_depth_nodes(from_table=from_table, to_table=to_table)
-
-        self.dashboard_service_json["name"] = name
-        self.dashboard_service_json["description"] = text
-        for chart in self.charts["charts"]:
-            yield Either(
-                right=CreateChartRequest(
-                    name=chart["name"],
-                    displayName=chart["displayName"],
-                    description=chart["description"],
-                    chartType=get_standard_chart_type(chart["chartType"]),
-                    sourceUrl=chart["sourceUrl"],
-                    service=self.dashboard_service.fullyQualifiedName,
+            for table in self.tables["tables"]:
+                text = self.generate_text()
+                table_request = self.create_table_request(
+                    table["name"], text, schema, table
+                )
+                yield Either(right=table_request)
+                table_entity_fqn = fqn.build(
+                    self.metadata,
+                    entity_type=Table,
+                    service_name=self.database_service.name.__root__,
+                    database_name=db.name.__root__,
+                    schema_name=schema.name.__root__,
+                    table_name=table_request.name.__root__,
+                )
+                from_table = self.metadata.get_by_name(
+                    entity=Table, fqn=table_entity_fqn
                 )
-            )
-        for data_model in self.data_models["datamodels"]:
-            name = self.generate_name()
-            text = self.generate_text()
-            data_model_request = self.create_dashboard_data_model_request(
-                name, text, data_model
-            )
-            yield Either(right=data_model_request)
-            data_model_entity_fqn = fqn.build(
-                self.metadata,
-                entity_type=DashboardDataModel,
-                service_name=self.dashboard_service.name.__root__,
-                data_model_name=data_model_request.name.__root__,
-            )
-            self.store_data_model_fqn.append(data_model_entity_fqn)
 
-        for table_fqn in self.store_table_fqn:
-            from_table = self.metadata.get_by_name(entity=Table, fqn=table_fqn)
-            for dashboard_datamodel_fqn in self.store_data_model_fqn:
-                to_datamodel = self.metadata.get_by_name(
-                    entity=DashboardDataModel, fqn=dashboard_datamodel_fqn
+                yield from self.create_depth_nodes(
+                    from_table=from_table, to_table=to_table
                 )
 
+            self.dashboard_service_json["name"] = name
+            self.dashboard_service_json["description"] = text
+            for chart in self.charts["charts"]:
                 yield Either(
-                    right=AddLineageRequest(
-                        edge=EntitiesEdge(
-                            fromEntity=EntityReference(
-                                id=from_table.id.__root__, type="table"
-                            ),
-                            toEntity=EntityReference(
-                                id=to_datamodel.id.__root__,
-                                type="dashboardDataModel",
-                            ),
-                            lineageDetails=LineageDetails(
-                                source=LineageSource.DashboardLineage
-                            ),
-                        )
-                    )
-                )
-                for dashboard in self.dashboards["dashboards"]:
-                    dashboard_request = CreateDashboardRequest(
-                        name=dashboard["name"],
-                        displayName=dashboard["displayName"],
-                        description=dashboard["description"],
-                        sourceUrl=dashboard["sourceUrl"],
-                        charts=dashboard["charts"],
+                    right=CreateChartRequest(
+                        name=chart["name"],
+                        displayName=chart["displayName"],
+                        description=chart["description"],
+                        chartType=get_standard_chart_type(chart["chartType"]),
+                        sourceUrl=chart["sourceUrl"],
                         service=self.dashboard_service.fullyQualifiedName,
                     )
-                    yield Either(right=dashboard_request)
-                    dashboard_fqn = fqn.build(
-                        self.metadata,
-                        entity_type=Dashboard,
-                        service_name=self.dashboard_service.name.__root__,
-                        dashboard_name=dashboard_request.name.__root__,
-                    )
-                    to_dashboard = self.metadata.get_by_name(
-                        entity=Dashboard, fqn=dashboard_fqn
+                )
+            for data_model in self.data_models["datamodels"]:
+                name = self.generate_name()
+                text = self.generate_text()
+                data_model_request = self.create_dashboard_data_model_request(
+                    name, text, data_model
+                )
+                yield Either(right=data_model_request)
+                data_model_entity_fqn = fqn.build(
+                    self.metadata,
+                    entity_type=DashboardDataModel,
+                    service_name=self.dashboard_service.name.__root__,
+                    data_model_name=data_model_request.name.__root__,
+                )
+                self.store_data_model_fqn.append(data_model_entity_fqn)
+
+            for table_fqn in self.store_table_fqn:
+                from_table = self.metadata.get_by_name(entity=Table, fqn=table_fqn)
+                for dashboard_datamodel_fqn in self.store_data_model_fqn:
+                    to_datamodel = self.metadata.get_by_name(
+                        entity=DashboardDataModel, fqn=dashboard_datamodel_fqn
                     )
+
                     yield Either(
                         right=AddLineageRequest(
                             edge=EntitiesEdge(
                                 fromEntity=EntityReference(
-                                    id=to_datamodel.id.__root__,
-                                    type="dashboardDataModel",
+                                    id=from_table.id.__root__, type="table"
                                 ),
                                 toEntity=EntityReference(
-                                    id=to_dashboard.id.__root__, type="dashboard"
+                                    id=to_datamodel.id.__root__,
+                                    type="dashboardDataModel",
                                 ),
                                 lineageDetails=LineageDetails(
                                     source=LineageSource.DashboardLineage
@@ -403,6 +393,84 @@ class ExtendedSampleDataSource(Source):  # pylint: disable=too-many-instance-att
                             )
                         )
                     )
+                    for dashboard in self.dashboards["dashboards"]:
+                        dashboard_request = CreateDashboardRequest(
+                            name=dashboard["name"],
+                            displayName=dashboard["displayName"],
+                            description=dashboard["description"],
+                            sourceUrl=dashboard["sourceUrl"],
+                            charts=dashboard["charts"],
+                            service=self.dashboard_service.fullyQualifiedName,
+                        )
+                        yield Either(right=dashboard_request)
+                        dashboard_fqn = fqn.build(
+                            self.metadata,
+                            entity_type=Dashboard,
+                            service_name=self.dashboard_service.name.__root__,
+                            dashboard_name=dashboard_request.name.__root__,
+                        )
+                        to_dashboard = self.metadata.get_by_name(
+                            entity=Dashboard, fqn=dashboard_fqn
+                        )
+                        yield Either(
+                            right=AddLineageRequest(
+                                edge=EntitiesEdge(
+                                    fromEntity=EntityReference(
+                                        id=to_datamodel.id.__root__,
+                                        type="dashboardDataModel",
+                                    ),
+                                    toEntity=EntityReference(
+                                        id=to_dashboard.id.__root__, type="dashboard"
+                                    ),
+                                    lineageDetails=LineageDetails(
+                                        source=LineageSource.DashboardLineage
+                                    ),
+                                )
+                            )
+                        )
+
+    def create_glossary(self):
+        self.main_glossary = CreateGlossaryRequest(
+            name="NestedGlossaryTest",
+            displayName="NestedGlossaryTest",
+            description="Description of test glossary",
+        )
+        yield Either(right=self.main_glossary)
+
+    def create_glossary_term(self):
+        """
+        Create Glossary Terms
+        """
+        for _ in range(20):
+            random_name = self.fake.first_name()
+            yield Either(
+                right=CreateGlossaryTermRequest(
+                    glossary="NestedGlossaryTest",
+                    name=random_name,
+                    displayName=random_name,
+                    description="Test glossary term ",
+                    parent=self.glossary_term_list[-1]
+                    if len(self.glossary_term_list) > 3
+                    else None,
+                )
+            )
+            if len(self.glossary_term_list) > 3 and self.glossary_term_list[-1]:
+                self.glossary_term_list.append(
+                    f"{self.glossary_term_list[-1]}.{random_name}"
+                )
+            else:
+                self.glossary_term_list.append(f"NestedGlossaryTest.{random_name}")
+
+        for _ in range(500):
+            random_name = self.fake.first_name()
+            yield Either(
+                right=CreateGlossaryTermRequest(
+                    glossary="NestedGlossaryTest",
+                    name=random_name,
+                    displayName=random_name,
+                    description="Test glossary term 1",
+                )
+            )
 
     def generate_name(self):
         return f"Sample-@!3_(%t3st@)%_^{self.fake.name()}"

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

@@ -126,7 +126,6 @@ def get_metastore_connection(connection: Any) -> Engine:
 
 @get_metastore_connection.register
 def _(connection: PostgresConnection):
-
     # import required to load sqlalchemy plugin
     # pylint: disable=import-outside-toplevel,unused-import
     from metadata.ingestion.source.database.hive.metastore_dialects.postgres import (  # nopycln: import
@@ -153,7 +152,6 @@ def _(connection: PostgresConnection):
 
 @get_metastore_connection.register
 def _(connection: MysqlConnection):
-
     # import required to load sqlalchemy plugin
     # pylint: disable=import-outside-toplevel,unused-import
     from metadata.ingestion.source.database.hive.metastore_dialects.mysql import (  # nopycln: import

metadata/ingestion/source/database/mssql/metadata.py

@@ -114,7 +114,6 @@ class MssqlSource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource):
         yield from self._execute_database_query(MSSQL_GET_DATABASE)
 
     def get_database_names(self) -> Iterable[str]:
-
         if not self.config.serviceConnection.__root__.config.ingestAllDatabases:
             configured_db = self.config.serviceConnection.__root__.config.database
             self.set_inspector(database_name=configured_db)

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

@@ -236,8 +236,7 @@ Q_HISTORY (database_name, query_text, start_time, end_time, duration,query_type,
   CROSS APPLY sys.dm_exec_sql_text(p.plan_handle) AS t
   INNER JOIN sys.databases db
     ON db.database_id = t.dbid
-  WHERE s.last_execution_time between '2024-01-13' and '2024-01-20'
-    AND t.text NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
+  WHERE t.text NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
     AND t.text NOT LIKE '/* {{"app": "dbt", %%}} */%%'
     AND p.objtype NOT IN ('Prepared', 'Proc')
     AND s.last_execution_time > '{start_date}'

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

@@ -172,14 +172,14 @@ SELECT
   SQL_FULLTEXT AS query_text,
   TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') AS start_time,
   ELAPSED_TIME / 1000 AS duration,
-  TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') + NUMTODSINTERVAL(ELAPSED_TIME / 1000, 'SECOND') AS end_time
+  TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') + NUMTODSINTERVAL(ELAPSED_TIME / 1000000, 'SECOND') AS end_time
 FROM gv$sql
 WHERE OBJECT_STATUS = 'VALID' 
   {filters}
   AND SQL_FULLTEXT NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
   AND SQL_FULLTEXT NOT LIKE '/* {{"app": "dbt", %%}} */%%'
   AND TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') >= TO_TIMESTAMP('{start_time}', 'yy-MM-dd HH24:MI:SS')
-  AND TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') + NUMTODSINTERVAL(ELAPSED_TIME / 1000, 'SECOND') 
+  AND TO_TIMESTAMP(FIRST_LOAD_TIME, 'yy-MM-dd/HH24:MI:SS') + NUMTODSINTERVAL(ELAPSED_TIME / 1000000, 'SECOND') 
   < TO_TIMESTAMP('{end_time}', 'yy-MM-dd HH24:MI:SS')
 ORDER BY FIRST_LOAD_TIME DESC 
 OFFSET 0 ROWS FETCH NEXT {result_limit} ROWS ONLY

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

@@ -62,7 +62,6 @@ def get_view_definition(
     dblink="",
     **kw,
 ):
-
     return get_view_definition_wrapper(
         self,
         connection,

metadata/ingestion/source/database/postgres/pgspider/lineage.py

@@ -75,7 +75,6 @@ def get_lineage_from_multi_tenant_table(
     connection: any,
     service_name: str,
 ) -> Iterator[Either[AddLineageRequest]]:
-
     """
     For PGSpider, firstly, get list of multi-tenant tables.
     Next, get child foreign tables of each multi-tenant tables.

metadata/ingestion/source/database/sample_data.py

@@ -800,7 +800,6 @@ class SampleDataSource(
 
         # Create table and stored procedure lineage
         for lineage_entities in self.stored_procedures["lineage"]:
-
             from_table = self.metadata.get_by_name(
                 entity=Table, fqn=lineage_entities["from_table_fqn"]
             )

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

@@ -465,7 +465,6 @@ class SasSource(
                 or table_entity.extension.__root__.get("analysisTimeStamp")
                 != table_extension.get("analysisTimeStamp")
             ):
-
                 # create the columns of the table
                 columns, col_profile_list = self.create_columns_and_profiles(
                     col_entity_instances, table_entity_instance
@@ -711,10 +710,10 @@ class SasSource(
                 if "state" in table_resource and table_resource["state"] == "unloaded":
                     self.sas_client.load_table(table_uri + "/state?value=loaded")
 
-            except HTTPError as e:
+            except HTTPError as exc:
                 # append http error to table description if it can't be found
                 logger.error(f"table_uri: {table_uri}")
-                self.report_description.append(str(e))
+                self.report_description.append(str(exc))
                 name_index = table_uri.rindex("/")
                 table_name = table_uri[name_index + 1 :]
                 param = f"filter=eq(name,'{table_name}')"

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

@@ -39,7 +39,6 @@ from metadata.utils.sqlalchemy_utils import (
 
 
 def _quoted_name(entity_name: Optional[str]) -> Optional[str]:
-
     if entity_name:
         return fqn.quote_name(entity_name)
 
@@ -121,13 +120,10 @@ def get_view_definition(  # pylint: disable=unused-argument
     schema = schema or self.default_schema_name
     if schema:
         cursor = connection.execute(
-            "SHOW /* sqlalchemy:get_view_definition */ VIEWS "
-            f"LIKE '{view_name}' IN {schema}"
+            f"SELECT GET_DDL('VIEW','{schema}.{view_name}') AS \"text\""
         )
     else:
-        cursor = connection.execute(
-            "SHOW /* sqlalchemy:get_view_definition */ VIEWS " f"LIKE '{view_name}'"
-        )
+        cursor = connection.execute(f"SELECT GET_DDL('VIEW','{view_name}') AS \"text\"")
     n2i = self.__class__._map_name_to_idx(cursor)  # pylint: disable=protected-access
     try:
         ret = cursor.fetchone()

metadata/ingestion/source/database/stored_procedures_mixin.py

@@ -153,7 +153,6 @@ class StoredProcedureMixin(ABC):
             query_type=query_by_procedure.query_type,
             query_text=query_by_procedure.query_text,
         ):
-
             self.context.stored_procedure_query_lineage = True
             for either_lineage in get_lineage_by_query(
                 self.metadata,

metadata/ingestion/source/database/unitycatalog/metadata.py

@@ -13,11 +13,10 @@ Databricks Unity Catalog Source source methods.
 """
 import json
 import traceback
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Iterable, List, Optional, Tuple, Union
 
 from databricks.sdk.service.catalog import ColumnInfo
 from databricks.sdk.service.catalog import TableConstraint as DBTableConstraint
-from databricks.sdk.service.catalog import TableConstraintList
 
 from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequest
 from metadata.generated.schema.api.data.createDatabaseSchema import (
@@ -75,19 +74,6 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 
-# pylint: disable=not-callable
-@classmethod
-def from_dict(cls, dct: Dict[str, Any]) -> "TableConstraintList":
-    return cls(
-        table_constraints=[
-            DBTableConstraint.from_dict(constraint) for constraint in dct
-        ]
-    )
-
-
-TableConstraintList.from_dict = from_dict
-
-
 class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
     """
     Implements the necessary methods to extract
@@ -351,7 +337,7 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
             )
 
     def get_table_constraints(
-        self, constraints: TableConstraintList
+        self, constraints: List[DBTableConstraint]
     ) -> Tuple[List[TableConstraint], List[ForeignConstrains]]:
         """
         Function to handle table constraint for the current table and add it to context
@@ -359,23 +345,22 @@ class UnitycatalogSource(DatabaseServiceSource, MultiDBSource):
 
         primary_constraints = []
         foreign_constraints = []
-        if constraints and constraints.table_constraints:
-            for constraint in constraints.table_constraints:
-                if constraint.primary_key_constraint:
-                    primary_constraints.append(
-                        TableConstraint(
-                            constraintType=ConstraintType.PRIMARY_KEY,
-                            columns=constraint.primary_key_constraint.child_columns,
-                        )
+        for constraint in constraints:
+            if constraint.primary_key_constraint:
+                primary_constraints.append(
+                    TableConstraint(
+                        constraintType=ConstraintType.PRIMARY_KEY,
+                        columns=constraint.primary_key_constraint.child_columns,
                     )
-                if constraint.foreign_key_constraint:
-                    foreign_constraints.append(
-                        ForeignConstrains(
-                            child_columns=constraint.foreign_key_constraint.child_columns,
-                            parent_columns=constraint.foreign_key_constraint.parent_columns,
-                            parent_table=constraint.foreign_key_constraint.parent_table,
-                        )
+                )
+            if constraint.foreign_key_constraint:
+                foreign_constraints.append(
+                    ForeignConstrains(
+                        child_columns=constraint.foreign_key_constraint.child_columns,
+                        parent_columns=constraint.foreign_key_constraint.parent_columns,
+                        parent_table=constraint.foreign_key_constraint.parent_table,
                     )
+                )
         return primary_constraints, foreign_constraints
 
     def _get_foreign_constraints(self, foreign_columns) -> List[TableConstraint]:

metadata/ingestion/source/messaging/common_broker_source.py

@@ -226,14 +226,16 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
         if topic_entity and self.generate_sample_data:
             topic_name = topic_details.topic_name
             sample_data = []
+            messages = None
             try:
-                self.consumer_client.subscribe(
-                    [topic_name], on_assign=on_partitions_assignment_to_consumer
-                )
-                logger.info(
-                    f"Broker consumer polling for sample messages in topic {topic_name}"
-                )
-                messages = self.consumer_client.consume(num_messages=10, timeout=10)
+                if self.consumer_client:
+                    self.consumer_client.subscribe(
+                        [topic_name], on_assign=on_partitions_assignment_to_consumer
+                    )
+                    logger.info(
+                        f"Broker consumer polling for sample messages in topic {topic_name}"
+                    )
+                    messages = self.consumer_client.consume(num_messages=10, timeout=10)
             except Exception as exc:
                 yield Either(
                     left=StackTraceError(
@@ -246,7 +248,6 @@ class CommonBrokerSource(MessagingServiceSource, ABC):
                 if messages:
                     for message in messages:
                         try:
-
                             value = message.value()
                             sample_data.append(
                                 self.decode_message(