diff --git a/metadata-ingestion/scripts/avro_codegen.py b/metadata-ingestion/scripts/avro_codegen.py index c895e5fabfd37..c7e4318346f1d 100644 --- a/metadata-ingestion/scripts/avro_codegen.py +++ b/metadata-ingestion/scripts/avro_codegen.py @@ -526,6 +526,36 @@ def get_notebook_id(self) -> str: """ ], "tag": [_create_from_id.format(class_name="TagUrn")], + "chart": [ + """ +@classmethod +def create_from_ids( + cls, + platform: str, + name: str, + platform_instance: Optional[str] = None, +) -> "ChartUrn": + return ChartUrn( + dashboard_tool=platform, + chart_id=f"{platform_instance}.{name}" if platform_instance else name, + ) + """ + ], + "dashboard": [ + """ +@classmethod +def create_from_ids( + cls, + platform: str, + name: str, + platform_instance: Optional[str] = None, +) -> "DashboardUrn": + return DashboardUrn( + dashboard_tool=platform, + dashboard_id=f"{platform_instance}.{name}" if platform_instance else name, + ) + """ + ], } diff --git a/metadata-ingestion/src/datahub/emitter/mce_builder.py b/metadata-ingestion/src/datahub/emitter/mce_builder.py index 51003ed5700d6..7d6f0bdcd8475 100644 --- a/metadata-ingestion/src/datahub/emitter/mce_builder.py +++ b/metadata-ingestion/src/datahub/emitter/mce_builder.py @@ -53,6 +53,8 @@ _Aspect as AspectAbstract, ) from datahub.metadata.urns import ( + ChartUrn, + DashboardUrn, DataFlowUrn, DataJobUrn, DataPlatformUrn, @@ -335,10 +337,11 @@ def make_dashboard_urn( platform: str, name: str, platform_instance: Optional[str] = None ) -> str: # FIXME: dashboards don't currently include data platform urn prefixes. - if platform_instance: - return f"urn:li:dashboard:({UrnEncoder.encode_string(platform)},{UrnEncoder.encode_string(platform_instance)}.{UrnEncoder.encode_string(name)})" - else: - return f"urn:li:dashboard:({UrnEncoder.encode_string(platform)},{UrnEncoder.encode_string(name)})" + return DashboardUrn.create_from_ids( + platform=platform, + name=name, + platform_instance=platform_instance, + ).urn() def dashboard_urn_to_key(dashboard_urn: str) -> Optional[DashboardKeyClass]: @@ -353,10 +356,11 @@ def make_chart_urn( platform: str, name: str, platform_instance: Optional[str] = None ) -> str: # FIXME: charts don't currently include data platform urn prefixes. - if platform_instance: - return f"urn:li:chart:({UrnEncoder.encode_string(platform)},{UrnEncoder.encode_string(platform_instance)}.{UrnEncoder.encode_string(name)})" - else: - return f"urn:li:chart:({UrnEncoder.encode_string(platform)},{UrnEncoder.encode_string(name)})" + return ChartUrn.create_from_ids( + platform=platform, + name=name, + platform_instance=platform_instance, + ).urn() def chart_urn_to_key(chart_urn: str) -> Optional[ChartKeyClass]: