[sc-28999] Custom metadata command

aesop/app.py

@@ -10,6 +10,7 @@
 from typing_extensions import Annotated
 
 from aesop.commands import (
+    custom_metadata_app,
     datasets_app,
     info_command,
     settings_app,
@@ -26,6 +27,7 @@
 app.add_typer(settings_app, name="settings")
 app.add_typer(datasets_app, name="datasets")
 app.add_typer(webhooks_app, name="webhooks")
+app.add_typer(custom_metadata_app, name="custom-metadata")
 
 
 @app.command()

aesop/commands/__init__.py

@@ -1,4 +1,4 @@
-from .aspects import tags_app
+from .aspects import custom_metadata_app, tags_app
 from .entities import datasets_app
 from .info.info import info as info_command
 from .settings.settings import app as settings_app
@@ -7,6 +7,7 @@
 
 __all__ = [
     "upload_command",
+    "custom_metadata_app",
     "info_command",
     "settings_app",
     "tags_app",

aesop/commands/aspects/__init__.py

@@ -1,5 +1,7 @@
+from .custom_metadata import app as custom_metadata_app
 from .user_defined_resources import tags_app
 
 __all__ = [
+    "custom_metadata_app",
     "tags_app",
 ]

aesop/commands/aspects/custom_metadata.py

@@ -0,0 +1,167 @@
+import csv
+import json
+import sys
+from typing import Dict, List
+
+from pydantic import BaseModel, TypeAdapter
+from rich import print, print_json
+from rich.table import Column, Table
+from typer import Context, FileText, Typer
+
+from aesop.commands.common.arguments import InputFileArg
+from aesop.commands.common.enums.output_format import OutputFormat
+from aesop.commands.common.exception_handler import exception_handler
+from aesop.commands.common.models import InputModel
+from aesop.commands.common.options import OutputFormatOption
+from aesop.config import AesopConfig
+from aesop.graphql.generated.get_dataset_custom_metadata import (
+    GetDatasetCustomMetadataNodeDataset,
+)
+from aesop.graphql.generated.input_types import (
+    CustomMetadataItemInput,
+    UpdateCustomMetadataInput,
+)
+
+app = Typer()
+
+
+class _Item(BaseModel):
+    key: str
+    value: str
+
+
+def _format_output(
+    items: List[_Item],
+    output_format: OutputFormat,
+) -> None:
+    if output_format is OutputFormat.JSON:
+        print_json(TypeAdapter(List[_Item]).dump_json(items).decode())
+
+    if output_format is OutputFormat.TABULAR:
+        table = Table(
+            Column(header="Key", no_wrap=True, style="bold cyan"),
+            "Value",
+            show_lines=True,
+        )
+        for item in items:
+            table.add_row(item.key, item.value)
+        print(table)
+
+    if output_format is OutputFormat.CSV:
+        spamwriter = csv.writer(sys.stdout)
+        spamwriter.writerow(["Key", "Value"])
+        spamwriter.writerows([[item.key, item.value] for item in items])
+
+
+@exception_handler("get custom metadata")
+@app.command(help="Gets custom metadata attached to a dataset.")
+def get(
+    ctx: Context,
+    dataset_id: str,
+    output: OutputFormat = OutputFormatOption,
+) -> None:
+    config: AesopConfig = ctx.obj
+    node = config.get_graphql_client().get_dataset_custom_metadata(dataset_id).node
+    assert isinstance(node, GetDatasetCustomMetadataNodeDataset)
+    metadata = node.custom_metadata.metadata if node.custom_metadata else []
+    _format_output(
+        [_Item(key=item.key, value=item.value) for item in metadata],
+        output,
+    )
+
+
+@exception_handler("update custom metadata")
+def _update(
+    config: AesopConfig,
+    output: OutputFormat,
+    dataset_id: str,
+    set: Dict[str, str] = {},
+    unset: List[str] = [],
+) -> None:
+    # Custom metadata values must be valid json
+    for value in set.values():
+        try:
+            json.loads(value)
+        except Exception:
+            raise ValueError(
+                f"Custom metadata value must be valid JSON string: {value}"
+            )
+
+    metadata = (
+        config.get_graphql_client()
+        .update_dataset_custom_metadata(
+            UpdateCustomMetadataInput(
+                entityId=dataset_id,
+                set=(
+                    [
+                        CustomMetadataItemInput(key=key, value=value)
+                        for key, value in set.items()
+                    ]
+                    if set
+                    else None
+                ),
+                unset=unset if unset else None,
+            )
+        )
+        .update_custom_metadata.metadata
+    )
+    _format_output(
+        [_Item(key=item.key, value=item.value) for item in metadata],
+        output,
+    )
+
+
+@exception_handler("add custom metadata")
+@app.command(help="Adds a custom metadata to a dataset.")
+def add(
+    ctx: Context,
+    dataset_id: str,
+    key: str,
+    value: str,
+    output: OutputFormat = OutputFormatOption,
+) -> None:
+    _update(ctx.obj, output, dataset_id, {key: value})
+
+
+@exception_handler("remove custom metadata")
+@app.command(help="Removes a custom metadata from a dataset.")
+def remove(
+    ctx: Context,
+    dataset_id: str,
+    key: str,
+    output: OutputFormat = OutputFormatOption,
+) -> None:
+    _update(ctx.obj, output, dataset_id, unset=[key])
+
+
+class _UpdateInputModel(InputModel, UpdateCustomMetadataInput):
+    @staticmethod
+    def example_json(indent: int = 0) -> str:
+        return UpdateCustomMetadataInput(
+            entityId="DATASET~00000000000000000000000000000001",
+            set=[
+                CustomMetadataItemInput(key="key_to_add", value='{"key": "value"}'),
+            ],
+            unset=[
+                "key_to_remove",
+                "another_key_to_remove",
+            ],
+        ).model_dump_json(indent=indent)
+
+
+@exception_handler("update custom metadata")
+@app.command(help="Batch updates custom metadata for a dataset.")
+def update(
+    ctx: Context,
+    input_file: FileText = InputFileArg(_UpdateInputModel),
+    output: OutputFormat = OutputFormatOption,
+) -> None:
+    input = UpdateCustomMetadataInput.model_validate_json(input_file.read())
+    assert input.entity_id
+    _update(
+        ctx.obj,
+        output,
+        input.entity_id,
+        {item.key or "": item.value or "" for item in input.set} if input.set else {},
+        input.unset or [],
+    )

aesop/commands/aspects/user_defined_resources/tags/node.py

@@ -1,9 +1,8 @@
 import csv
-import json
 import sys
 from typing import List, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 from rich.table import Column, Table
 
 from aesop.commands.common.enums.output_format import OutputFormat
@@ -35,6 +34,4 @@ def display_nodes(
         spamwriter.writerow(["ID", "Name", "Description"])
         spamwriter.writerows([[node.id, node.name, node.description] for node in nodes])
     elif output is OutputFormat.JSON:
-        console.print_json(
-            json.dumps([node.model_dump(exclude_none=True) for node in nodes]), indent=2
-        )
+        console.print_json(TypeAdapter(List[GovernedTagNode]).dump_json(nodes).decode())

aesop/commands/common/arguments.py

@@ -2,6 +2,7 @@
 from typing import Any
 
 import typer
+from rich.box import SIMPLE
 from rich.markdown import Markdown
 from rich.panel import Panel
 
@@ -13,35 +14,40 @@ def _validate_input_file(
     input_model: type[InputModel], input_file: typer.FileText
 ) -> typer.FileText:
     if input_file.name == "<stdin>" and input_file.isatty():
+        console.print(Markdown("---"))
         # Got nothing, print example and exit
         example_contents = ["```json"]
         example_contents.extend(input_model.example_json(indent=2).splitlines())
         example_contents.append("```")
         example_panel = Panel(
             Markdown("\n".join(example_contents)),
-            title="[green][bold]Example input",
+            title="[green][bold][Example input]",
             title_align="left",
+            padding=1,
+            box=SIMPLE,
         )
         console.print(example_panel)
-
+        console.print(Markdown("---"))
         commands = " ".join(sys.argv[1:])
         usage_contents = [
             "Pipe the JSON body into the command:",
             "",
             "```bash",
-            f"$ cat {''.join(input_model.example_json(indent=0).splitlines())} | aesop {commands}",  # noqa E501
+            f"$ echo '{''.join(input_model.example_json(indent=0).splitlines())}' | aesop {commands}",  # noqa E501
             "```",
             "Or provide an input file to the command:",
             "```bash",
-            f"$ echo {''.join(input_model.example_json(indent=0).splitlines())} > input.json",  # noqa E501
+            f"$ echo '{''.join(input_model.example_json(indent=0).splitlines())}' > input.json",  # noqa E501
             "",
             f"$ aesop {commands} input.json",
             "```",
         ]
         usage_panel = Panel(
             Markdown("\n".join(usage_contents)),
-            title="[green][bold]Usage",
+            title="[green][bold][Usage]",
             title_align="left",
+            padding=1,
+            box=SIMPLE,
         )
         console.print(usage_panel)
 

aesop/commands/webhooks/webhooks.py

@@ -3,12 +3,14 @@
 from rich import print_json
 from typer import Argument, Context, Typer
 
+from aesop.commands.common.exception_handler import exception_handler
 from aesop.config import AesopConfig
 from aesop.graphql.generated.enums import WebhookTriggerType
 
 app = Typer(help="Manages webhooks.")
 
 
+@exception_handler("register webhook")
 @app.command(help="Registers a webhook to Metaphor.")
 def register(
     ctx: Context,
@@ -25,6 +27,7 @@ def register(
     )
 
 
+@exception_handler("unregister webhook")
 @app.command(help="Unregisters a webhook from Metaphor.")
 def unregister(
     ctx: Context,
@@ -38,6 +41,7 @@ def unregister(
     )
 
 
+@exception_handler("get webhooks")
 @app.command(help="Gets a list of webhooks that are registered to Metaphor.")
 def get(
     ctx: Context,
@@ -47,6 +51,7 @@ def get(
     print_json(config.get_graphql_client().get_webhooks(trigger).model_dump_json())
 
 
+@exception_handler("get webhook payload schema")
 @app.command(help="Gets the payload of a webhook trigger type.")
 def get_payload_schema(
     ctx: Context,