feat(cli): Add run-id option to put sub-command

metadata-ingestion/src/datahub/cli/cli_utils.py

@@ -13,7 +13,7 @@
 from datahub.emitter.aspect import ASPECT_MAP, TIMESERIES_ASPECT_MAP
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.request_helper import make_curl_command
-from datahub.emitter.serialization_helper import post_json_transform
+from datahub.emitter.serialization_helper import post_json_transform, pre_json_transform
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
  MetadataChangeEvent,
  MetadataChangeProposal,
@@ -153,10 +153,11 @@ def post_entity(
  aspect_value: Dict,
  cached_session_host: Optional[Tuple[Session, str]] = None,
  is_async: Optional[str] = "false",
+ system_metadata: Union[None, SystemMetadataClass] = None,
 ) -> int:
  endpoint: str = "/aspects/?action=ingestProposal"
 
- proposal = {
+ proposal: Dict[str, Any] = {
  "proposal": {
  "entityType": entity_type,
  "entityUrn": urn,
@@ -169,6 +170,12 @@ def post_entity(
  },
  "async": is_async,
  }
+
+ if system_metadata is not None:
+ proposal["proposal"]["systemMetadata"] = json.dumps(
+ pre_json_transform(system_metadata.to_obj())
+ )
+
  payload = json.dumps(proposal)
  url = gms_host + endpoint
  curl_command = make_curl_command(session, "POST", url, payload)

metadata-ingestion/src/datahub/cli/put_cli.py

@@ -6,7 +6,7 @@
 
 from datahub.cli.cli_utils import post_entity
 from datahub.configuration.config_loader import load_config_file
-from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.mcp import MetadataChangeProposalWrapper, SystemMetadataClass
 from datahub.ingestion.graph.client import get_default_graph
 from datahub.metadata.schema_classes import (
  DataPlatformInfoClass as DataPlatformInfo,
@@ -36,9 +36,10 @@ def put() -> None:
 @click.option("--urn", required=True, type=str)
 @click.option("-a", "--aspect", required=True, type=str)
 @click.option("-d", "--aspect-data", required=True, type=str)
+@click.option("--run-id", type=str, help="Run ID into which we should log the aspect.")
 @upgrade.check_upgrade
 @telemetry.with_telemetry()
-def aspect(urn: str, aspect: str, aspect_data: str) -> None:
+def aspect(urn: str, aspect: str, aspect_data: str, run_id: str) -> None:
  """Update a single aspect of an entity"""
 
  entity_type = guess_entity_type(urn)
@@ -56,6 +57,7 @@ def aspect(urn: str, aspect: str, aspect_data: str) -> None:
  aspect_name=aspect,
  entity_type=entity_type,
  aspect_value=aspect_obj,
+ system_metadata=SystemMetadataClass(runId=run_id),
  )
  click.secho(f"Update succeeded with status {status}", fg="green")
 
@@ -82,8 +84,11 @@ def aspect(urn: str, aspect: str, aspect_data: str) -> None:
  help="Logo URL that must be reachable from the DataHub UI.",
  required=True,
 )
+@click.option(
+ "--run-id", type=str, help="Run ID into which we should log the platform."
+)
 def platform(
- ctx: click.Context, name: str, display_name: Optional[str], logo: str
+ ctx: click.Context, name: str, display_name: Optional[str], logo: str, run_id: str
 ) -> None:
  """
  Create or update a dataplatform entity in DataHub
@@ -104,11 +109,12 @@ def platform(
  logoUrl=logo,
  )
  datahub_graph = get_default_graph()
- datahub_graph.emit(
- MetadataChangeProposalWrapper(
-  entityUrn=str(platform_urn), aspect=data_platform_info
- )
+ mcp = MetadataChangeProposalWrapper(
+ entityUrn=str(platform_urn),
+ aspect=data_platform_info,
+ systemMetadata=SystemMetadataClass(runId=run_id),
  )
+ datahub_graph.emit(mcp)
  click.echo(
  f"✅ Successfully wrote data platform metadata for {platform_urn} to DataHub ({datahub_graph})"
  )