acryl-datahub 1.2.0.1__py3-none-any.whl → 1.2.0.2rc1__py3-none-any.whl

datahub/ingestion/source/sql/postgres.py

@@ -36,6 +36,15 @@ from datahub.ingestion.source.sql.sql_common import (
     register_custom_type,
 )
 from datahub.ingestion.source.sql.sql_config import BasicSQLAlchemyConfig
+from datahub.ingestion.source.sql.sql_utils import (
+    gen_database_key,
+    gen_schema_key,
+)
+from datahub.ingestion.source.sql.stored_procedures.base import (
+    BaseProcedure,
+    generate_procedure_container_workunits,
+    generate_procedure_workunits,
+)
 from datahub.metadata.com.linkedin.pegasus2avro.schema import (
     ArrayTypeClass,
     BytesTypeClass,
@@ -123,6 +132,15 @@ class PostgresConfig(BasePostgresConfig):
             "Note: this is not used if `database` or `sqlalchemy_uri` are provided."
         ),
     )
+    include_stored_procedures: bool = Field(
+        default=True,
+        description="Include ingest of stored procedures.",
+    )
+    procedure_pattern: AllowDenyPattern = Field(
+        default=AllowDenyPattern.allow_all(),
+        description="Regex patterns for stored procedures to filter in ingestion."
+        "Specify regex to match the entire procedure name in database.schema.procedure_name format. e.g. to match all procedures starting with customer in Customer database and public schema, use the regex 'Customer.public.customer.*'",
+    )
 
 
 @platform_name("Postgres")
@@ -135,7 +153,7 @@ class PostgresSource(SQLAlchemySource):
     """
     This plugin extracts the following:
 
-    - Metadata for databases, schemas, views, and tables
+    - Metadata for databases, schemas, views, tables, and stored procedures
     - Column types associated with each table
     - Also supports PostGIS extensions
     - Table, row, and column statistics via optional SQL profiling
@@ -291,3 +309,174 @@ class PostgresSource(SQLAlchemySource):
                     ] = row.table_size
         except Exception as e:
             logger.error(f"failed to fetch profile metadata: {e}")
+
+    def get_schema_level_workunits(
+        self,
+        inspector: Inspector,
+        schema: str,
+        database: str,
+    ) -> Iterable[Union[MetadataWorkUnit, SqlWorkUnit]]:
+        yield from super().get_schema_level_workunits(
+            inspector=inspector,
+            schema=schema,
+            database=database,
+        )
+
+        if self.config.include_stored_procedures:
+            try:
+                yield from self.loop_stored_procedures(inspector, schema, self.config)
+            except Exception as e:
+                self.report.failure(
+                    title="Failed to list stored procedures for schema",
+                    message="An error occurred while listing procedures for the schema.",
+                    context=f"{database}.{schema}",
+                    exc=e,
+                )
+
+    def loop_stored_procedures(
+        self,
+        inspector: Inspector,
+        schema: str,
+        config: PostgresConfig,
+    ) -> Iterable[MetadataWorkUnit]:
+        """
+        Loop schema data for get stored procedures as dataJob-s.
+        """
+        db_name = self.get_db_name(inspector)
+
+        procedures = self.fetch_procedures_for_schema(inspector, schema, db_name)
+        if procedures:
+            yield from self._process_procedures(procedures, db_name, schema)
+
+    def fetch_procedures_for_schema(
+        self, inspector: Inspector, schema: str, db_name: str
+    ) -> List[BaseProcedure]:
+        try:
+            raw_procedures: List[BaseProcedure] = self.get_procedures_for_schema(
+                inspector, schema, db_name
+            )
+            procedures: List[BaseProcedure] = []
+            for procedure in raw_procedures:
+                procedure_qualified_name = self.get_identifier(
+                    schema=schema,
+                    entity=procedure.name,
+                    inspector=inspector,
+                )
+
+                if not self.config.procedure_pattern.allowed(procedure_qualified_name):
+                    self.report.report_dropped(procedure_qualified_name)
+                else:
+                    procedures.append(procedure)
+            return procedures
+        except Exception as e:
+            self.report.warning(
+                title="Failed to get procedures for schema",
+                message="An error occurred while fetching procedures for the schema.",
+                context=f"{db_name}.{schema}",
+                exc=e,
+            )
+            return []
+
+    def get_procedures_for_schema(
+        self, inspector: Inspector, schema: str, db_name: str
+    ) -> List[BaseProcedure]:
+        """
+        Get stored procedures for a specific schema.
+        """
+        base_procedures = []
+        with inspector.engine.connect() as conn:
+            procedures = conn.execute(
+                """
+                    SELECT
+                        p.proname AS name,
+                        l.lanname AS language,
+                        pg_get_function_arguments(p.oid) AS arguments,
+                        pg_get_functiondef(p.oid) AS definition,
+                        obj_description(p.oid, 'pg_proc') AS comment
+                    FROM
+                        pg_proc p
+                    JOIN
+                        pg_namespace n ON n.oid = p.pronamespace
+                    JOIN
+                        pg_language l ON l.oid = p.prolang
+                    WHERE
+                        p.prokind = 'p'
+                        AND n.nspname = '"""
+                + schema
+                + """';
+                    """
+            )
+
+            procedure_rows = list(procedures)
+            for row in procedure_rows:
+                base_procedures.append(
+                    BaseProcedure(
+                        name=row.name,
+                        language=row.language,
+                        argument_signature=row.arguments,
+                        return_type=None,
+                        procedure_definition=row.definition,
+                        created=None,
+                        last_altered=None,
+                        comment=row.comment,
+                        extra_properties=None,
+                    )
+                )
+            return base_procedures
+
+    def _process_procedures(
+        self,
+        procedures: List[BaseProcedure],
+        db_name: str,
+        schema: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        if procedures:
+            yield from generate_procedure_container_workunits(
+                database_key=gen_database_key(
+                    database=db_name,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_key=gen_schema_key(
+                    db_name=db_name,
+                    schema=schema,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+            )
+        for procedure in procedures:
+            yield from self._process_procedure(procedure, schema, db_name)
+
+    def _process_procedure(
+        self,
+        procedure: BaseProcedure,
+        schema: str,
+        db_name: str,
+    ) -> Iterable[MetadataWorkUnit]:
+        try:
+            yield from generate_procedure_workunits(
+                procedure=procedure,
+                database_key=gen_database_key(
+                    database=db_name,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_key=gen_schema_key(
+                    db_name=db_name,
+                    schema=schema,
+                    platform=self.platform,
+                    platform_instance=self.config.platform_instance,
+                    env=self.config.env,
+                ),
+                schema_resolver=self.get_schema_resolver(),
+            )
+        except Exception as e:
+            self.report.warning(
+                title="Failed to emit stored procedure",
+                message="An error occurred while emitting stored procedure",
+                context=procedure.name,
+                exc=e,
+            )

datahub/ingestion/source/sql_queries.py

@@ -4,7 +4,7 @@ import os
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from functools import partial
-from typing import Iterable, List, Optional, Set
+from typing import Iterable, List, Optional, Union
 
 from pydantic import Field
 
@@ -14,9 +14,8 @@ from datahub.configuration.source_common import (
 )
 from datahub.emitter.mce_builder import (
     make_dataset_urn_with_platform_instance,
-    make_user_urn,
 )
-from datahub.emitter.sql_parsing_builder import SqlParsingBuilder
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -35,8 +34,15 @@ from datahub.ingestion.api.source_helpers import auto_workunit_reporter
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
+from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
+from datahub.metadata.urns import CorpUserUrn
 from datahub.sql_parsing.schema_resolver import SchemaResolver
-from datahub.sql_parsing.sqlglot_lineage import sqlglot_lineage
+from datahub.sql_parsing.sql_parsing_aggregator import (
+    KnownQueryLineageInfo,
+    ObservedQuery,
+    SqlAggregatorReport,
+    SqlParsingAggregator,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -67,28 +73,19 @@ class SqlQueriesSourceConfig(PlatformInstanceConfigMixin, EnvConfigMixin):
         default=None,
     )
     override_dialect: Optional[str] = Field(
-        description="The SQL dialect to use when parsing queries. Overrides automatic dialect detection.",
+        description="DEPRECATED: This field is ignored. SQL dialect detection is now handled automatically by the SQL parsing aggregator based on the platform.",
         default=None,
+        hidden_from_docs=True,
     )
 
 
-class SqlQueriesSourceReport(SourceReport):
-    num_queries_parsed: int = 0
-    num_table_parse_failures: int = 0
-    num_column_parse_failures: int = 0
+@dataclass
+class SqlQueriesSourceReport(SourceReport, IngestionStageReport):
+    num_entries_processed: int = 0
+    num_entries_failed: int = 0
+    num_queries_aggregator_failures: int = 0
 
-    def compute_stats(self) -> None:
-        super().compute_stats()
-        self.table_failure_rate = (
-            f"{self.num_table_parse_failures / self.num_queries_parsed:.4f}"
-            if self.num_queries_parsed
-            else "0"
-        )
-        self.column_failure_rate = (
-            f"{self.num_column_parse_failures / self.num_queries_parsed:.4f}"
-            if self.num_queries_parsed
-            else "0"
-        )
+    sql_aggregator: Optional[SqlAggregatorReport] = None
 
 
 @platform_name("SQL Queries")
@@ -107,15 +104,15 @@ class SqlQueriesSource(Source):
     - user (optional): string - The user who ran the query.
     This user value will be directly converted into a DataHub user urn.
     - operation_type (optional): string - Platform-specific operation type, used if the operation type can't be parsed.
+    - session_id (optional): string - Session identifier for temporary table resolution across queries.
     - downstream_tables (optional): string[] - Fallback list of tables that the query writes to,
      used if the query can't be parsed.
     - upstream_tables (optional): string[] - Fallback list of tables the query reads from,
      used if the query can't be parsed.
     """
 
-    urns: Optional[Set[str]]
-    schema_resolver: SchemaResolver
-    builder: SqlParsingBuilder
+    schema_resolver: Optional[SchemaResolver]
+    aggregator: SqlParsingAggregator
 
     def __init__(self, ctx: PipelineContext, config: SqlQueriesSourceConfig):
         if not ctx.graph:
@@ -128,22 +125,36 @@ class SqlQueriesSource(Source):
         self.config = config
         self.report = SqlQueriesSourceReport()
 
-        self.builder = SqlParsingBuilder(usage_config=self.config.usage)
-
         if self.config.use_schema_resolver:
+            # TODO: `initialize_schema_resolver_from_datahub` does a  bulk initialization by fetching all schemas
+            # for the given platform, platform instance, and env. Instead this should be configurable:
+            # bulk initialization vs lazy on-demand schema fetching.
             self.schema_resolver = self.graph.initialize_schema_resolver_from_datahub(
                 platform=self.config.platform,
                 platform_instance=self.config.platform_instance,
                 env=self.config.env,
             )
-            self.urns = self.schema_resolver.get_urns()
         else:
-            self.schema_resolver = self.graph._make_schema_resolver(
-                platform=self.config.platform,
-                platform_instance=self.config.platform_instance,
-                env=self.config.env,
-            )
-            self.urns = None
+            self.schema_resolver = None
+
+        self.aggregator = SqlParsingAggregator(
+            platform=self.config.platform,
+            platform_instance=self.config.platform_instance,
+            env=self.config.env,
+            schema_resolver=self.schema_resolver,
+            eager_graph_load=False,
+            generate_lineage=True,  # TODO: make this configurable
+            generate_queries=True,  # TODO: make this configurable
+            generate_query_subject_fields=True,  # TODO: make this configurable
+            generate_query_usage_statistics=True,  # This enables publishing SELECT query entities, otherwise only mutation queries are published
+            generate_usage_statistics=True,
+            generate_operations=True,  # TODO: make this configurable
+            usage_config=self.config.usage,
+            is_temp_table=None,
+            is_allowed_table=None,
+            format_queries=False,
+        )
+        self.report.sql_aggregator = self.aggregator.report
 
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> "SqlQueriesSource":
@@ -156,68 +167,91 @@ class SqlQueriesSource(Source):
     def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
         return [partial(auto_workunit_reporter, self.get_report())]
 
-    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+    def get_workunits_internal(
+        self,
+    ) -> Iterable[Union[MetadataWorkUnit, MetadataChangeProposalWrapper]]:
         logger.info(f"Parsing queries from {os.path.basename(self.config.query_file)}")
+
+        with self.report.new_stage("Collecting queries from file"):
+            queries = list(self._parse_query_file())
+            logger.info(f"Collected {len(queries)} queries for processing")
+
+        with self.report.new_stage("Processing queries through SQL parsing aggregator"):
+            for query_entry in queries:
+                self._add_query_to_aggregator(query_entry)
+
+        with self.report.new_stage("Generating metadata work units"):
+            logger.info("Generating workunits from SQL parsing aggregator")
+            yield from self.aggregator.gen_metadata()
+
+    def _parse_query_file(self) -> Iterable["QueryEntry"]:
+        """Parse the query file and yield QueryEntry objects."""
         with open(self.config.query_file) as f:
             for line in f:
                 try:
                     query_dict = json.loads(line, strict=False)
                     entry = QueryEntry.create(query_dict, config=self.config)
-                    yield from self._process_query(entry)
+                    self.report.num_entries_processed += 1
+                    if self.report.num_entries_processed % 1000 == 0:
+                        logger.info(
+                            f"Processed {self.report.num_entries_processed} query entries"
+                        )
+                    yield entry
                 except Exception as e:
-                    logger.warning("Error processing query", exc_info=True)
-                    self.report.report_warning("process-query", str(e))
-
-        logger.info("Generating workunits")
-        yield from self.builder.gen_workunits()
+                    self.report.num_entries_failed += 1
+                    self.report.warning(
+                        title="Error processing query",
+                        message="Query skipped due to parsing error",
+                        context=line.strip(),
+                        exc=e,
+                    )
 
-    def _process_query(self, entry: "QueryEntry") -> Iterable[MetadataWorkUnit]:
-        self.report.num_queries_parsed += 1
-        if self.report.num_queries_parsed % 1000 == 0:
-            logger.info(f"Parsed {self.report.num_queries_parsed} queries")
-
-        result = sqlglot_lineage(
-            sql=entry.query,
-            schema_resolver=self.schema_resolver,
-            default_db=self.config.default_db,
-            default_schema=self.config.default_schema,
-            override_dialect=self.config.override_dialect,
-        )
-        if result.debug_info.table_error:
-            logger.info(f"Error parsing table lineage, {result.debug_info.table_error}")
-            self.report.num_table_parse_failures += 1
-            for downstream_urn in set(entry.downstream_tables):
-                self.builder.add_lineage(
-                    downstream_urn=downstream_urn,
-                    upstream_urns=entry.upstream_tables,
-                    timestamp=entry.timestamp,
-                    user=entry.user,
+    def _add_query_to_aggregator(self, query_entry: "QueryEntry") -> None:
+        """Add a query to the SQL parsing aggregator."""
+        try:
+            # If we have explicit lineage, use it directly
+            if query_entry.upstream_tables or query_entry.downstream_tables:
+                logger.debug("Using explicit lineage from query file")
+                for downstream_table in query_entry.downstream_tables:
+                    known_lineage = KnownQueryLineageInfo(
+                        query_text=query_entry.query,
+                        downstream=downstream_table,
+                        upstreams=query_entry.upstream_tables,
+                        timestamp=query_entry.timestamp,
+                        session_id=query_entry.session_id,
+                    )
+                    self.aggregator.add_known_query_lineage(known_lineage)
+            else:
+                # No explicit lineage, rely on parsing
+                observed_query = ObservedQuery(
+                    query=query_entry.query,
+                    timestamp=query_entry.timestamp,
+                    user=query_entry.user,
+                    session_id=query_entry.session_id,
+                    default_db=self.config.default_db,
+                    default_schema=self.config.default_schema,
                 )
-            return
-        elif result.debug_info.column_error:
-            logger.debug(
-                f"Error parsing column lineage, {result.debug_info.column_error}"
+                self.aggregator.add_observed_query(observed_query)
+
+        except Exception as e:
+            self.report.num_queries_aggregator_failures += 1
+            self.report.warning(
+                title="Error adding query to aggregator",
+                message="Query skipped due to failure when adding query to SQL parsing aggregator",
+                context=query_entry.query,
+                exc=e,
             )
-            self.report.num_column_parse_failures += 1
-
-        yield from self.builder.process_sql_parsing_result(
-            result,
-            query=entry.query,
-            query_timestamp=entry.timestamp,
-            user=entry.user,
-            custom_operation_type=entry.operation_type,
-            include_urns=self.urns,
-        )
 
 
 @dataclass
 class QueryEntry:
     query: str
     timestamp: Optional[datetime]
-    user: Optional[str]
+    user: Optional[CorpUserUrn]
     operation_type: Optional[str]
     downstream_tables: List[str]
     upstream_tables: List[str]
+    session_id: Optional[str] = None
 
     @classmethod
     def create(
@@ -230,7 +264,7 @@ class QueryEntry:
                 if "timestamp" in entry_dict
                 else None
             ),
-            user=make_user_urn(entry_dict["user"]) if "user" in entry_dict else None,
+            user=CorpUserUrn(entry_dict["user"]) if "user" in entry_dict else None,
             operation_type=entry_dict.get("operation_type"),
             downstream_tables=[
                 make_dataset_urn_with_platform_instance(
@@ -250,4 +284,5 @@ class QueryEntry:
                 )
                 for table in entry_dict.get("upstream_tables", [])
             ],
+            session_id=entry_dict.get("session_id"),
         )

datahub/metadata/_internal_schema_classes.py

@@ -19964,17 +19964,20 @@ class DataHubPageModuleParamsClass(DictWrapper):
         linkParams: Union[None, "LinkModuleParamsClass"]=None,
         richTextParams: Union[None, "RichTextModuleParamsClass"]=None,
         assetCollectionParams: Union[None, "AssetCollectionModuleParamsClass"]=None,
+        hierarchyViewParams: Union[None, "HierarchyModuleParamsClass"]=None,
     ):
         super().__init__()
         
         self.linkParams = linkParams
         self.richTextParams = richTextParams
         self.assetCollectionParams = assetCollectionParams
+        self.hierarchyViewParams = hierarchyViewParams
     
     def _restore_defaults(self) -> None:
         self.linkParams = self.RECORD_SCHEMA.fields_dict["linkParams"].default
         self.richTextParams = self.RECORD_SCHEMA.fields_dict["richTextParams"].default
         self.assetCollectionParams = self.RECORD_SCHEMA.fields_dict["assetCollectionParams"].default
+        self.hierarchyViewParams = self.RECORD_SCHEMA.fields_dict["hierarchyViewParams"].default
     
     
     @property
@@ -20007,6 +20010,16 @@ class DataHubPageModuleParamsClass(DictWrapper):
         self._inner_dict['assetCollectionParams'] = value
     
     
+    @property
+    def hierarchyViewParams(self) -> Union[None, "HierarchyModuleParamsClass"]:
+        """The params required if the module is type HIERARCHY_VIEW"""
+        return self._inner_dict.get('hierarchyViewParams')  # type: ignore
+    
+    @hierarchyViewParams.setter
+    def hierarchyViewParams(self, value: Union[None, "HierarchyModuleParamsClass"]) -> None:
+        self._inner_dict['hierarchyViewParams'] = value
+    
+    
 class DataHubPageModulePropertiesClass(_Aspect):
     """The main properties of a DataHub page module"""
 
@@ -20149,6 +20162,46 @@ class DataHubPageModuleVisibilityClass(DictWrapper):
         self._inner_dict['scope'] = value
     
     
+class HierarchyModuleParamsClass(DictWrapper):
+    """The params required if the module is type HIERARCHY_VIEW
+    TODO: add filters
+    relatedEntitiesFilter: optional Filter"""
+    
+    RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.module.HierarchyModuleParams")
+    def __init__(self,
+        showRelatedEntities: bool,
+        assetUrns: Union[None, List[str]]=None,
+    ):
+        super().__init__()
+        
+        self.assetUrns = assetUrns
+        self.showRelatedEntities = showRelatedEntities
+    
+    def _restore_defaults(self) -> None:
+        self.assetUrns = self.RECORD_SCHEMA.fields_dict["assetUrns"].default
+        self.showRelatedEntities = bool()
+    
+    
+    @property
+    def assetUrns(self) -> Union[None, List[str]]:
+        # No docs available.
+        return self._inner_dict.get('assetUrns')  # type: ignore
+    
+    @assetUrns.setter
+    def assetUrns(self, value: Union[None, List[str]]) -> None:
+        self._inner_dict['assetUrns'] = value
+    
+    
+    @property
+    def showRelatedEntities(self) -> bool:
+        # No docs available.
+        return self._inner_dict.get('showRelatedEntities')  # type: ignore
+    
+    @showRelatedEntities.setter
+    def showRelatedEntities(self, value: bool) -> None:
+        self._inner_dict['showRelatedEntities'] = value
+    
+    
 class LinkModuleParamsClass(DictWrapper):
     # No docs available.
     
@@ -20772,6 +20825,8 @@ class SystemMetadataClass(_Aspect):
         registryVersion: Union[None, str]=None,
         properties: Union[None, Dict[str, str]]=None,
         version: Union[None, str]=None,
+        aspectCreated: Union[None, "AuditStampClass"]=None,
+        aspectModified: Union[None, "AuditStampClass"]=None,
     ):
         super().__init__()
         
@@ -20795,6 +20850,8 @@ class SystemMetadataClass(_Aspect):
         self.registryVersion = registryVersion
         self.properties = properties
         self.version = version
+        self.aspectCreated = aspectCreated
+        self.aspectModified = aspectModified
     
     def _restore_defaults(self) -> None:
         self.lastObserved = self.RECORD_SCHEMA.fields_dict["lastObserved"].default
@@ -20805,6 +20862,8 @@ class SystemMetadataClass(_Aspect):
         self.registryVersion = self.RECORD_SCHEMA.fields_dict["registryVersion"].default
         self.properties = self.RECORD_SCHEMA.fields_dict["properties"].default
         self.version = self.RECORD_SCHEMA.fields_dict["version"].default
+        self.aspectCreated = self.RECORD_SCHEMA.fields_dict["aspectCreated"].default
+        self.aspectModified = self.RECORD_SCHEMA.fields_dict["aspectModified"].default
     
     
     @property
@@ -20889,6 +20948,26 @@ class SystemMetadataClass(_Aspect):
         self._inner_dict['version'] = value
     
     
+    @property
+    def aspectCreated(self) -> Union[None, "AuditStampClass"]:
+        """When the aspect was initially created and who created it, detected by version 0 -> 1 change"""
+        return self._inner_dict.get('aspectCreated')  # type: ignore
+    
+    @aspectCreated.setter
+    def aspectCreated(self, value: Union[None, "AuditStampClass"]) -> None:
+        self._inner_dict['aspectCreated'] = value
+    
+    
+    @property
+    def aspectModified(self) -> Union[None, "AuditStampClass"]:
+        """When the aspect was last modified and the actor that performed the modification"""
+        return self._inner_dict.get('aspectModified')  # type: ignore
+    
+    @aspectModified.setter
+    def aspectModified(self, value: Union[None, "AuditStampClass"]) -> None:
+        self._inner_dict['aspectModified'] = value
+    
+    
 class ChartCellClass(DictWrapper):
     """Chart cell in a notebook, which will present content in chart format"""
     
@@ -27171,6 +27250,7 @@ __SCHEMA_TYPES = {
     'com.linkedin.pegasus2avro.module.DataHubPageModuleProperties': DataHubPageModulePropertiesClass,
     'com.linkedin.pegasus2avro.module.DataHubPageModuleType': DataHubPageModuleTypeClass,
     'com.linkedin.pegasus2avro.module.DataHubPageModuleVisibility': DataHubPageModuleVisibilityClass,
+    'com.linkedin.pegasus2avro.module.HierarchyModuleParams': HierarchyModuleParamsClass,
     'com.linkedin.pegasus2avro.module.LinkModuleParams': LinkModuleParamsClass,
     'com.linkedin.pegasus2avro.module.PageModuleScope': PageModuleScopeClass,
     'com.linkedin.pegasus2avro.module.RichTextModuleParams': RichTextModuleParamsClass,
@@ -27680,6 +27760,7 @@ __SCHEMA_TYPES = {
     'DataHubPageModuleProperties': DataHubPageModulePropertiesClass,
     'DataHubPageModuleType': DataHubPageModuleTypeClass,
     'DataHubPageModuleVisibility': DataHubPageModuleVisibilityClass,
+    'HierarchyModuleParams': HierarchyModuleParamsClass,
     'LinkModuleParams': LinkModuleParamsClass,
     'PageModuleScope': PageModuleScopeClass,
     'RichTextModuleParams': RichTextModuleParamsClass,

datahub/metadata/com/linkedin/pegasus2avro/module/__init__.py

@@ -12,6 +12,7 @@ from .....schema_classes import DataHubPageModuleParamsClass
 from .....schema_classes import DataHubPageModulePropertiesClass
 from .....schema_classes import DataHubPageModuleTypeClass
 from .....schema_classes import DataHubPageModuleVisibilityClass
+from .....schema_classes import HierarchyModuleParamsClass
 from .....schema_classes import LinkModuleParamsClass
 from .....schema_classes import PageModuleScopeClass
 from .....schema_classes import RichTextModuleParamsClass
@@ -22,6 +23,7 @@ DataHubPageModuleParams = DataHubPageModuleParamsClass
 DataHubPageModuleProperties = DataHubPageModulePropertiesClass
 DataHubPageModuleType = DataHubPageModuleTypeClass
 DataHubPageModuleVisibility = DataHubPageModuleVisibilityClass
+HierarchyModuleParams = HierarchyModuleParamsClass
 LinkModuleParams = LinkModuleParamsClass
 PageModuleScope = PageModuleScopeClass
 RichTextModuleParams = RichTextModuleParamsClass