feat(datahub-lite): adding tab completion, small serialization fixes (#7079)

This commit is contained in:
Shirshanka Das 2023-01-19 02:07:07 -08:00 committed by GitHub
parent bdcc356cc5
commit f31ff9c91e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 147 additions and 15 deletions

View File

@ -1,3 +1,6 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# DataHub Lite (Experimental)
## What is it?
@ -106,12 +109,18 @@ As a convenient short-cut, you can import metadata from any standard DataHub met
## Exploring Metadata
The `datahub lite` group of commands provides a set of capabilities for you to explore the metadata you just ingested.
The `datahub lite` group of commands provides a set of capabilities for you to explore the metadata you just ingested.
### List (ls)
Listing functions like a directory structure that is customized based on the kind of system being explored. DataHub's metadata is automatically organized into databases, tables, views, dashboards, charts, etc.
:::note
Using the `ls` command below is much more pleasant when you have tab completion enabled on your shell. Check out the [Setting up Tab Completion](#tab-completion) section at the bottom of the guide.
:::
```shell
> datahub lite ls /
databases
@ -133,6 +142,9 @@ metadata_index
metadata_aspect_v2
```
### Read (get)
Once you have located a path of interest, you can read metadata at that entity, by issuing a **get**. You can additionally filter the metadata retrieved from an entity by the aspect type of the metadata (e.g. to request the schema, filter by the **schemaMetadata** aspect).
@ -157,7 +169,7 @@ Get metadata for an entity by path
</summary>
```json
> datahub lite get /databases/mysql/instances/default/databases/datahub/tables/metadata_aspect_v2
> datahub lite get --path /databases/mysql/instances/default/databases/datahub/tables/metadata_aspect_v2
{
"urn": "urn:li:dataset:(urn:li:dataPlatform:mysql,datahub.metadata_aspect_v2,PROD)",
"container": {
@ -313,7 +325,7 @@ Get metadata for an entity by path
#### Get metadata for an entity filtered by specific aspect
```json
> datahub lite get /databases/mysql/instances/default/databases/datahub/tables/metadata_aspect_v2 --aspect status
> datahub lite get --path /databases/mysql/instances/default/databases/datahub/tables/metadata_aspect_v2 --aspect status
{
"urn": "urn:li:dataset:(urn:li:dataPlatform:mysql,datahub.metadata_aspect_v2,PROD)",
"status": {
@ -324,10 +336,17 @@ Get metadata for an entity by path
}
```
:::note
Using the `get` command by path is much more pleasant when you have tab completion enabled on your shell. Check out the [Setting up Tab Completion](#tab-completion) section at the bottom of the guide.
:::
#### Get metadata using the urn of the entity
```json
> datahub lite get "urn:li:dataset:(urn:li:dataPlatform:mysql,datahub.metadata_aspect_v2,PROD)" --aspect status
> datahub lite get --urn "urn:li:dataset:(urn:li:dataPlatform:mysql,datahub.metadata_aspect_v2,PROD)" --aspect status
{
"urn": "urn:li:dataset:(urn:li:dataPlatform:mysql,datahub.metadata_aspect_v2,PROD)",
"status": {
@ -344,7 +363,7 @@ Get metadata with additional details (systemMetadata)
</summary>
```json
> datahub lite get /databases/mysql/instances/default/databases/datahub/tables/metadata_aspect_v2 --aspect status --verbose
> datahub lite get --path /databases/mysql/instances/default/databases/datahub/tables/metadata_aspect_v2 --aspect status --verbose
{
"urn": "urn:li:dataset:(urn:li:dataPlatform:mysql,datahub.metadata_aspect_v2,PROD)",
"status": {
@ -518,4 +537,81 @@ DataHub Lite maintains a few derived tables to make access possible via both the
## Caveat Emptor!
DataHub Lite is a very new project. Do not use it for production use-cases. The API-s and storage formats are subject to change and we get feedback from early adopters. That said, we are really interested in accepting PR-s and suggestions for improvements to this fledgling project.
DataHub Lite is a very new project. Do not use it for production use-cases. The API-s and storage formats are subject to change and we get feedback from early adopters. That said, we are really interested in accepting PR-s and suggestions for improvements to this fledgling project.
## Advanced Options
### Tab Completion
Using the datahub lite commands like `ls` or `get` is much more pleasant when you have tab completion enabled on your shell. Tab completion is supported on the command line through the [Click Shell completion](https://click.palletsprojects.com/en/8.1.x/shell-completion/) module.
To set up shell completion for your shell, follow the instructions below:
#### Option 1 (inline eval)
<Tabs>
<TabItem value="zsh" label="Zsh" default>
Add this to ~/.zshrc:
```shell
eval "$(_DATAHUB_COMPLETE=zsh_source datahub)"
```
</TabItem>
<TabItem value="bash" label="Bash">
Add this to ~/.bashrc:
```shell
eval "$(_DATAHUB_COMPLETE=bash_source datahub)"
```
</TabItem>
</Tabs>
#### Option 2 (external completion script)
Using eval means that the command is invoked and evaluated every time a shell is started, which can delay shell responsiveness. To speed it up, write the generated script to a file, then source that.
<Tabs>
<TabItem value="zsh" label="Zsh" default>
Save the script somewhere.
```shell
_DATAHUB_COMPLETE=zsh_source datahub > ~/.datahub-complete.zsh
```
Source the file in ~/.zshrc.
```shell
. ~/.datahub-complete.zsh
```
</TabItem>
<TabItem value="bash" label="Bash">
```shell
_DATAHUB_COMPLETE=bash_source datahub > ~/.datahub-complete.bash
```
Source the file in ~/.bashrc.
```shell
. ~/.datahub-complete.bash
```
</TabItem>
<TabItem value="fish" label="Fish">
Save the script to ~/.config/fish/completions/datahub.fish:
```shell
_DATAHUB_COMPLETE=fish_source datahub > ~/.config/fish/completions/datahub.fish
```
</TabItem>
</Tabs>

View File

@ -6,6 +6,7 @@ from datetime import datetime
from typing import List, Optional
import click
from click.shell_completion import CompletionItem
from click_default_group import DefaultGroup
from datahub.cli.cli_utils import (
@ -73,9 +74,36 @@ def list_urns() -> None:
click.echo(result)
class CompleteablePath(click.ParamType):
name = "path"
def shell_complete(self, ctx, param, incomplete):
path = incomplete or "/"
lite = _get_datahub_lite(read_only=True)
try:
completions = lite.ls(path)
return [
CompletionItem(browseable.auto_complete.suggested_path, type="plain")
if browseable.auto_complete
else CompletionItem(
f"{incomplete}/{browseable.name}".replace("//", "/")
)
for browseable in completions
if not browseable.leaf
]
except Exception as e:
logger.debug(f"failed with {e}")
return []
@lite.command(context_settings=dict(allow_extra_args=True))
@click.option("--urn", required=False, type=str, help="Get metadata rooted at an urn")
@click.option("--path", required=False, type=str, help="Get metadata rooted at a path")
@click.option(
"--path",
required=False,
type=CompleteablePath(),
help="Get metadata rooted at a path",
)
@click.option("-a", "--aspect", required=False, multiple=True, type=str)
@click.option("--asof", required=False, type=click.DateTime(formats=["%Y-%m-%d"]))
@click.option("--verbose", required=False, is_flag=True, default=False)
@ -90,7 +118,7 @@ def get(
verbose: bool,
) -> None:
"""Get one or more metadata elements"""
start_time = time.time()
if urn is None and path is None:
if not ctx.args:
raise click.UsageError(
@ -142,6 +170,8 @@ def get(
indent=2,
)
)
end_time = time.time()
logger.debug(f"Time taken: {int((end_time - start_time)*1000.0)} millis")
@lite.command()
@ -182,16 +212,19 @@ def serve(port: int) -> None:
@lite.command(context_settings=dict(allow_extra_args=True))
@click.argument("path", required=False)
@click.argument("path", required=False, type=CompleteablePath())
@click.pass_context
@telemetry.with_telemetry
def ls(ctx: click.Context, path: Optional[str]) -> None:
"""List at a path"""
start_time = time.time()
path = path or "/"
lite = _get_datahub_lite(read_only=True)
try:
browseables = lite.ls(path)
end_time = time.time()
logger.debug(f"Time taken: {int((end_time - start_time)*1000.0)} millis")
auto_complete: List[AutoComplete] = [
b.auto_complete for b in browseables if b.auto_complete is not None
]

View File

@ -10,6 +10,7 @@ from datahub.configuration.common import ConfigModel
from datahub.emitter.aspect import ASPECT_MAP
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.emitter.mcp_builder import mcps_from_mce
from datahub.emitter.serialization_helper import post_json_transform
from datahub.lite.lite_local import (
AutoComplete,
Browseable,
@ -260,7 +261,7 @@ class DuckDBLite(DataHubLiteLocal[DuckDBLiteConfig]):
aspect: Union[dict, _Aspect] = json.loads(r[2])
if typed:
assert isinstance(aspect, dict)
aspect = ASPECT_MAP[aspect_name].from_obj(aspect)
aspect = ASPECT_MAP[aspect_name].from_obj(post_json_transform(aspect))
result_map[aspect_name] = {"value": aspect}
if details:
@ -496,7 +497,9 @@ class DuckDBLite(DataHubLiteLocal[DuckDBLiteConfig]):
aspect_name in ASPECT_MAP
), f"Missing aspect name {aspect_name} in the registry"
try:
aspect_payload = ASPECT_MAP[aspect_name].from_obj(aspect_payload)
aspect_payload = ASPECT_MAP[aspect_name].from_obj(
post_json_transform(aspect_payload)
)
except Exception as e:
logger.exception(
f"Failed to process urn: {urn}, aspect_name: {aspect_name}, metadata: {aspect_payload}",
@ -524,7 +527,7 @@ class DuckDBLite(DataHubLiteLocal[DuckDBLiteConfig]):
for r in results.fetchall():
urn = r[0]
aspect_name = r[1]
aspect_metadata = ASPECT_MAP[aspect_name].from_obj(json.loads(r[2])) # type: ignore
aspect_metadata = ASPECT_MAP[aspect_name].from_obj(post_json_transform(json.loads(r[2]))) # type: ignore
system_metadata = SystemMetadataClass.from_obj(json.loads(r[3]))
mcp = MetadataChangeProposalWrapper(
entityUrn=urn,
@ -560,17 +563,17 @@ class DuckDBLite(DataHubLiteLocal[DuckDBLiteConfig]):
"iceberg",
"trino",
],
"streaming_systems": ["kafka"],
"streaming": ["kafka"],
"orchestrators": ["airflow", "spark"],
"data_movers": ["kafka-connect", "nifi"],
"transformation_tools": ["dbt"],
"data_quality_tools": ["great-expectations"],
"data_quality": ["great-expectations"],
}
for k, v in category_to_platform_map.items():
if data_platform_urn.get_entity_id_as_string() in v:
return Urn(entity_type="systemNode", entity_id=[k])
logger.warning(
logger.debug(
f"Failed to find category for platform {data_platform_urn}, mapping to generic data_platform"
)
return Urn(entity_type="systemNode", entity_id=["data_platforms"])