feat(dataset_cli): add dry-run support

Author: chakru-r

metadata-ingestion/src/datahub/cli/specific/dataset_cli.py

@@ -29,13 +29,16 @@ def dataset() -> None:
     name="upsert",
 )
 @click.option("-f", "--file", required=True, type=click.Path(exists=True))
+@click.option(
+    "-n", "--dry-run", type=bool, is_flag=True, default=False, help="Perform a dry run"
+)
 @upgrade.check_upgrade
 @telemetry.with_telemetry()
-def upsert(file: Path) -> None:
+def upsert(file: Path, dry_run: bool) -> None:
     """Upsert attributes to a Dataset in DataHub."""
     # Call the sync command with to_datahub=True to perform the upsert operation
     ctx = click.get_current_context()
-    ctx.invoke(sync, file=str(file), to_datahub=True)
+    ctx.invoke(sync, file=str(file), dry_run=dry_run, to_datahub=True)
 
 
 @dataset.command(
@@ -167,11 +170,16 @@ def file(lintcheck: bool, lintfix: bool, file: str) -> None:
 )
 @click.option("-f", "--file", required=True, type=click.Path(exists=True))
 @click.option("--to-datahub/--from-datahub", required=True, is_flag=True)
+@click.option(
+    "-n", "--dry-run", type=bool, is_flag=True, default=False, help="Perform a dry run"
+)
 @upgrade.check_upgrade
 @telemetry.with_telemetry()
-def sync(file: str, to_datahub: bool) -> None:
+def sync(file: str, to_datahub: bool, dry_run: bool) -> None:
     """Sync a Dataset file to/from DataHub"""
 
+    dry_run_prefix = "[dry-run]: " if dry_run else ""  # prefix to use in messages
+
     failures: List[str] = []
     with get_default_graph() as graph:
         datasets = Dataset.from_yaml(file)
@@ -189,7 +197,7 @@ def sync(file: str, to_datahub: bool) -> None:
                     click.secho(
                         "\n\t- ".join(
                             [
-                                f"Skipping Dataset {dataset.urn} due to missing entity references: "
+                                f"{dry_run_prefix}Skipping Dataset {dataset.urn} due to missing entity references: "
                             ]
                             + missing_entity_references
                         ),
@@ -199,13 +207,18 @@ def sync(file: str, to_datahub: bool) -> None:
                     continue
                 try:
                     for mcp in dataset.generate_mcp():
-                        graph.emit(mcp)
-                    click.secho(f"Update succeeded for urn {dataset.urn}.", fg="green")
+                        if not dry_run:
+                            graph.emit(mcp)
+                    click.secho(
+                        f"{dry_run_prefix}Update succeeded for urn {dataset.urn}.",
+                        fg="green",
+                    )
                 except Exception as e:
                     click.secho(
-                        f"Update failed for id {id}. due to {e}",
+                        f"{dry_run_prefix}Update failed for id {id}. due to {e}",
                         fg="red",
                     )
+                    failures.append(dataset.urn)
             else:
                 # Sync from DataHub
                 if graph.exists(dataset.urn):
@@ -215,13 +228,16 @@ def sync(file: str, to_datahub: bool) -> None:
                     existing_dataset: Dataset = Dataset.from_datahub(
                         graph=graph, urn=dataset.urn, config=dataset_get_config
                     )
-                    existing_dataset.to_yaml(Path(file))
+                    if not dry_run:
+                        existing_dataset.to_yaml(Path(file))
+                    else:
+                        click.secho(f"{dry_run_prefix}Will update file {file}")
                 else:
-                    click.secho(f"Dataset {dataset.urn} does not exist")
+                    click.secho(f"{dry_run_prefix}Dataset {dataset.urn} does not exist")
                     failures.append(dataset.urn)
     if failures:
         click.secho(
-            f"\nFailed to sync the following Datasets: {', '.join(failures)}",
+            f"\n{dry_run_prefix}Failed to sync the following Datasets: {', '.join(failures)}",
             fg="red",
         )
         raise click.Abort()

metadata-ingestion/tests/unit/cli/dataset/test_dataset_cmd.py

@@ -217,3 +217,63 @@ def test_multiple_datasets_in_file(self, mock_dataset, test_yaml_file):
             # Verify both dataset instances had to_yaml called
             mock_dataset1.to_yaml.assert_called_once()
             mock_dataset2.to_yaml.assert_called_once()
+
+    @patch("datahub.ingestion.graph.client.get_default_graph")
+    def test_dry_run_sync(self, mock_get_default_graph):
+        mock_graph = MagicMock()
+        mock_get_default_graph.return_value.__enter__.return_value = mock_graph
+
+        test_file = TEST_RESOURCES_DIR / "dataset_dry_run.yaml"
+
+        runner = CliRunner()
+        result = runner.invoke(
+            dataset, ["sync", "--dry-run", "--to-datahub", "-f", str(test_file)]
+        )
+
+        # Verify
+        assert result.exit_code == 0
+        assert not mock_get_default_graph.emit.called
+
+    @patch("datahub.ingestion.graph.client.get_default_graph")
+    def test_dry_run_sync_fail(self, mock_get_default_graph):
+        mock_graph = MagicMock()
+        mock_get_default_graph.return_value.__enter__.return_value = mock_graph
+
+        test_file = TEST_RESOURCES_DIR / "dataset_dry_run_bad_type.yaml"
+        runner = CliRunner()
+        result = runner.invoke(
+            dataset, ["sync", "--dry-run", "--to-datahub", "-f", str(test_file)]
+        )
+
+        # Verify
+        assert result.exit_code != 0
+        assert not mock_get_default_graph.emit.called
+        assert "Type bad_type is not a valid primitive type" in result.output
+
+    @patch("datahub.cli.specific.dataset_cli.get_default_graph")
+    def test_run_sync(self, mock_get_default_graph):
+        mock_graph = MagicMock()
+        mock_get_default_graph.return_value.__enter__.return_value = mock_graph
+
+        test_file = TEST_RESOURCES_DIR / "dataset_dry_run.yaml"
+
+        runner = CliRunner()
+        result = runner.invoke(dataset, ["sync", "--to-datahub", "-f", str(test_file)])
+
+        # Verify
+        assert result.exit_code == 0
+        assert mock_graph.emit.called
+
+    @patch("datahub.cli.specific.dataset_cli.get_default_graph")
+    def test_run_sync_fail(self, mock_get_default_graph):
+        mock_graph = MagicMock()
+        mock_get_default_graph.return_value.__enter__.return_value = mock_graph
+
+        test_file = TEST_RESOURCES_DIR / "dataset_dry_run_bad_type.yaml"
+        runner = CliRunner()
+        result = runner.invoke(dataset, ["sync", "--to-datahub", "-f", str(test_file)])
+
+        # Verify
+        assert result.exit_code != 0
+        assert not mock_get_default_graph.emit.called
+        assert "is not a valid primitive type" in result.output

metadata-ingestion/tests/unit/cli/dataset/test_resources/dataset_dry_run.yaml

@@ -0,0 +1,13 @@
+## This file is used to define the dataset schema for the dataset `foobar`
+- id: user.clicksv3
+  platform: hive
+  urn: urn:li:dataset:(urn:li:dataPlatform:hive,user.clicksv3,PROD) # use urn instead of id and platform
+  subtype: View
+  schema:
+    fields:
+    - id: ip
+      type: string
+      description: The IP address of the user
+    - id: user_id
+      type: string
+      description: The user ID of the user

metadata-ingestion/tests/unit/cli/dataset/test_resources/dataset_dry_run_bad_type.yaml

@@ -0,0 +1,13 @@
+## This file is used to define the dataset schema for the dataset `foobar`
+- id: user.clicksv3
+  platform: hive
+  urn: urn:li:dataset:(urn:li:dataPlatform:hive,user.clicksv3,PROD) # use urn instead of id and platform
+  subtype: View
+  schema:
+    fields:
+    - id: ip
+      type: string
+      description: The IP address of the user
+    - id: user_id
+      type: bad_type
+      description: The user ID of the user