fix(ingest/vertica): Fixing missing container properties (datahub-pro…

…ject#7197)

metadata-ingestion/src/datahub/ingestion/source/sql/athena.py

@@ -1,7 +1,7 @@
 import json
 import logging
 import typing
-from typing import Dict, Iterable, List, Optional, Tuple, cast
+from typing import Any, Dict, Iterable, List, Optional, Tuple, cast
 
 import pydantic
 from pyathena.common import BaseCursor
@@ -186,6 +186,7 @@ def get_table_properties(
     def gen_database_containers(
         self,
         database: str,
+        extra_properties: Optional[Dict[str, Any]] = None,
     ) -> Iterable[MetadataWorkUnit]:
         # In Athena the schema is the database and database is not existing
         return []
@@ -194,6 +195,7 @@ def gen_schema_containers(
         self,
         schema: str,
         database: str,
+        extra_properties: Optional[Dict[str, Any]] = None,
     ) -> Iterable[MetadataWorkUnit]:
         database_container_key = gen_database_key(
             database,
@@ -209,6 +211,7 @@ def gen_schema_containers(
             domain_registry=self.domain_registry,
             domain_config=self.config.domain,
             report=self.report,
+            extra_properties=extra_properties,
         )
 
     def get_database_container_key(self, db_name: str, schema: str) -> PlatformKey:

metadata-ingestion/src/datahub/ingestion/source/sql/presto_on_hive.py

@@ -325,6 +325,7 @@ def create(cls, config_dict, ctx):
     def gen_database_containers(
         self,
         database: str,
+        extra_properties: Optional[Dict[str, Any]] = None,
     ) -> Iterable[MetadataWorkUnit]:
         database_container_key = gen_database_key(
             database,
@@ -340,12 +341,14 @@ def gen_database_containers(
             domain_registry=self.domain_registry,
             domain_config=self.config.domain,
             report=self.report,
+            extra_properties=extra_properties,
         )
 
     def gen_schema_containers(
         self,
         schema: str,
         database: str,
+        extra_properties: Optional[Dict[str, Any]] = None,
     ) -> Iterable[MetadataWorkUnit]:
         assert isinstance(self.config, PrestoOnHiveConfig)
         where_clause_suffix: str = ""
@@ -394,6 +397,7 @@ def gen_schema_containers(
                 domain_registry=self.domain_registry,
                 domain_config=self.config.domain,
                 report=self.report,
+                extra_properties=extra_properties,
             )
 
     def get_default_ingestion_job_id(self) -> JobId:

metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py

@@ -417,6 +417,7 @@ def get_allowed_schemas(self, inspector: Inspector, db_name: str) -> Iterable[st
     def gen_database_containers(
         self,
         database: str,
+        extra_properties: Optional[Dict[str, Any]] = None,
     ) -> Iterable[MetadataWorkUnit]:
         database_container_key = gen_database_key(
             database,
@@ -432,12 +433,14 @@ def gen_database_containers(
             domain_registry=self.domain_registry,
             domain_config=self.config.domain,
             report=self.report,
+            extra_properties=extra_properties,
         )
 
     def gen_schema_containers(
         self,
         schema: str,
         database: str,
+        extra_properties: Optional[Dict[str, Any]] = None,
     ) -> Iterable[MetadataWorkUnit]:
 
         database_container_key = gen_database_key(
@@ -464,6 +467,7 @@ def gen_schema_containers(
             domain_registry=self.domain_registry,
             domain_config=self.config.domain,
             report=self.report,
+            extra_properties=extra_properties,
         )
 
     def add_table_to_schema_container(
@@ -517,6 +521,9 @@ def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit
                 yield from self.gen_schema_containers(
                     database=db_name,
                     schema=schema,
+                    extra_properties=self.get_schema_properties(
+                        inspector=inspector, schema=schema, database=db_name
+                    ),
                 )
 
                 if sql_config.include_tables:

metadata-ingestion/src/datahub/ingestion/source/sql/two_tier_sql_source.py

@@ -1,5 +1,5 @@
 import typing
-from typing import Iterable, Optional
+from typing import Any, Dict, Iterable, Optional
 
 from pydantic.fields import Field
 from sqlalchemy import create_engine, inspect
@@ -116,6 +116,7 @@ def gen_schema_containers(
         self,
         schema: str,
         database: str,
+        extra_properties: Optional[Dict[str, Any]] = None,
     ) -> Iterable[MetadataWorkUnit]:
         return []
 

metadata-ingestion/src/datahub/ingestion/source/sql/vertica.py

@@ -163,12 +163,23 @@ def get_workunits(self) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
                 profiler = self.get_profiler_instance(inspector)
 
             db_name = self.get_db_name(inspector)
-            yield from self.gen_database_containers(database=db_name)
+            yield from self.gen_database_containers(
+                database=db_name,
+                extra_properties=self.get_database_properties(
+                    inspector=inspector, database=db_name
+                ),
+            )
 
             for schema in self.get_allowed_schemas(inspector, db_name):
                 self.add_information_for_schema(inspector, schema)
 
-                yield from self.gen_schema_containers(schema=schema, database=db_name)
+                yield from self.gen_schema_containers(
+                    schema=schema,
+                    database=db_name,
+                    extra_properties=self.get_schema_properties(
+                        inspector=inspector, schema=schema, database=db_name
+                    ),
+                )
 
                 if sql_config.include_tables:
                     yield from self.loop_tables(inspector, schema, sql_config)