acryl-datahub 0.15.0.6rc3__py3-none-any.whl → 1.0.0__py3-none-any.whl

datahub/api/entities/structuredproperties/structuredproperties.py

@@ -1,7 +1,7 @@
 import logging
 from enum import Enum
 from pathlib import Path
-from typing import Iterable, List, Optional
+from typing import Iterable, List, Optional, Union
 
 import yaml
 from pydantic import validator
@@ -38,7 +38,7 @@ class AllowedTypes(Enum):
 
 
 class AllowedValue(ConfigModel):
-    value: str
+    value: Union[int, float, str]
     description: Optional[str] = None
 
 

datahub/cli/check_cli.py

@@ -5,7 +5,8 @@ import pathlib
 import pprint
 import shutil
 import tempfile
-from typing import Dict, List, Optional, Union
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Union
 
 import click
 
@@ -20,7 +21,10 @@ from datahub.ingestion.sink.sink_registry import sink_registry
 from datahub.ingestion.source.source_registry import source_registry
 from datahub.ingestion.transformer.transform_registry import transform_registry
 from datahub.telemetry import telemetry
-from datahub.utilities.file_backed_collections import ConnectionWrapper, FileBackedList
+from datahub.utilities.file_backed_collections import (
+    ConnectionWrapper,
+    FileBackedDict,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -391,29 +395,78 @@ def test_path_spec(config: str, input: str, path_spec_key: str) -> None:
         raise e
 
 
+def _jsonify(data: Any) -> Any:
+    if dataclasses.is_dataclass(data):
+        # dataclasses.asdict() is recursive. We're doing the recursion
+        # manually here via _jsonify calls, so we can't use
+        # dataclasses.asdict() here.
+        return {
+            f.name: _jsonify(getattr(data, f.name)) for f in dataclasses.fields(data)
+        }
+    elif isinstance(data, list):
+        return [_jsonify(item) for item in data]
+    elif isinstance(data, dict):
+        return {_jsonify(k): _jsonify(v) for k, v in data.items()}
+    elif isinstance(data, datetime):
+        return data.isoformat()
+    else:
+        return data
+
+
 @check.command()
-@click.argument("query-log-file", type=click.Path(exists=True, dir_okay=False))
-@click.option("--output", type=click.Path())
-def extract_sql_agg_log(query_log_file: str, output: Optional[str]) -> None:
+@click.argument("db-file", type=click.Path(exists=True, dir_okay=False))
+def extract_sql_agg_log(db_file: str) -> None:
     """Convert a sqlite db generated by the SqlParsingAggregator into a JSON."""
 
-    from datahub.sql_parsing.sql_parsing_aggregator import LoggedQuery
+    if pathlib.Path(db_file).suffix != ".db":
+        raise click.UsageError("DB file must be a sqlite db")
+
+    output_dir = pathlib.Path(db_file).with_suffix("")
+    output_dir.mkdir(exist_ok=True)
+
+    shared_connection = ConnectionWrapper(pathlib.Path(db_file))
+
+    tables: List[str] = [
+        row[0]
+        for row in shared_connection.execute(
+            """\
+SELECT
+    name
+FROM
+    sqlite_schema
+WHERE
+    type ='table' AND
+    name NOT LIKE 'sqlite_%';
+""",
+            parameters={},
+        )
+    ]
+    logger.info(f"Extracting {len(tables)} tables from {db_file}: {tables}")
+
+    for table in tables:
+        table_output_path = output_dir / f"{table}.json"
+        if table_output_path.exists():
+            logger.info(f"Skipping {table_output_path} because it already exists")
+            continue
 
-    assert dataclasses.is_dataclass(LoggedQuery)
+        # Some of the tables might actually be FileBackedList. Because
+        # the list is built on top of the FileBackedDict, we don't
+        # need to distinguish between the two cases.
 
-    shared_connection = ConnectionWrapper(pathlib.Path(query_log_file))
-    query_log = FileBackedList[LoggedQuery](
-        shared_connection=shared_connection, tablename="stored_queries"
-    )
-    logger.info(f"Extracting {len(query_log)} queries from {query_log_file}")
-    queries = [dataclasses.asdict(query) for query in query_log]
+        table_data: FileBackedDict[Any] = FileBackedDict(
+            shared_connection=shared_connection, tablename=table
+        )
 
-    if output:
-        with open(output, "w") as f:
-            json.dump(queries, f, indent=2, default=str)
-        logger.info(f"Extracted {len(queries)} queries to {output}")
-    else:
-        click.echo(json.dumps(queries, indent=2))
+        data = {}
+        with click.progressbar(
+            table_data.items(), length=len(table_data), label=f"Extracting {table}"
+        ) as items:
+            for k, v in items:
+                data[k] = _jsonify(v)
+
+        with open(table_output_path, "w") as f:
+            json.dump(data, f, indent=2, default=str)
+        logger.info(f"Extracted {len(data)} entries to {table_output_path}")
 
 
 @check.command()

datahub/cli/docker_cli.py

@@ -231,7 +231,7 @@ def _docker_compose_v2() -> List[str]:
             # docker-compose v1 is not installed either.
             raise DockerComposeVersionError(
                 "You don't have Docker Compose installed. Please install Docker Compose. See https://docs.docker.com/compose/install/.",
-            )
+            ) from None
 
 
 def _attempt_stop(quickstart_compose_file: List[pathlib.Path]) -> None:
@@ -430,7 +430,7 @@ def detect_quickstart_arch(arch: Optional[str]) -> Architectures:
     return quickstart_arch
 
 
-@docker.command()  # noqa: C901
+@docker.command()
 @click.option(
     "--version",
     type=str,
@@ -592,7 +592,7 @@ def detect_quickstart_arch(arch: Optional[str]) -> Architectures:
         "arch",
     ]
 )
-def quickstart(  # noqa: C901
+def quickstart(
     version: Optional[str],
     build_locally: bool,
     pull_images: bool,

datahub/cli/iceberg_cli.py

@@ -645,7 +645,7 @@ def get_all_warehouses(client: DataHubGraph) -> Iterator[str]:
     graph_query = """
         query getIcebergWarehouses($start: Int, $count: Int) {
           search(
-            input: {type: DATA_PLATFORM_INSTANCE, query: "*", start: $start, count: $count}
+            input: {type: DATA_PLATFORM_INSTANCE, query: "dataPlatform:iceberg", start: $start, count: $count}
           ) {
             start
             total

datahub/cli/ingest_cli.py

@@ -15,14 +15,14 @@ from tabulate import tabulate
 from datahub._version import nice_version_name
 from datahub.cli import cli_utils
 from datahub.cli.config_utils import CONDENSED_DATAHUB_CONFIG_PATH
-from datahub.configuration.common import ConfigModel, GraphError
+from datahub.configuration.common import GraphError
 from datahub.configuration.config_loader import load_config_file
-from datahub.emitter.mce_builder import datahub_guid
 from datahub.ingestion.graph.client import get_default_graph
 from datahub.ingestion.run.connection import ConnectionManager
 from datahub.ingestion.run.pipeline import Pipeline
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
+from datahub.utilities.ingest_utils import deploy_source_vars
 from datahub.utilities.perf_timer import PerfTimer
 
 logger = logging.getLogger(__name__)
@@ -191,23 +191,6 @@ def run(
     # don't raise SystemExit if there's no error
 
 
-def _make_ingestion_urn(name: str) -> str:
-    guid = datahub_guid(
-        {
-            "name": name,
-        }
-    )
-    return f"urn:li:dataHubIngestionSource:deploy-{guid}"
-
-
-class DeployOptions(ConfigModel):
-    name: str
-    schedule: Optional[str] = None
-    time_zone: str = "UTC"
-    cli_version: Optional[str] = None
-    executor_id: str = "default"
-
-
 @ingest.command()
 @upgrade.check_upgrade
 @telemetry.with_telemetry()
@@ -258,6 +241,16 @@ class DeployOptions(ConfigModel):
     required=False,
     default="UTC",
 )
+@click.option(
+    "--debug", type=bool, help="Should we debug.", required=False, default=False
+)
+@click.option(
+    "--extra-pip",
+    type=str,
+    help='Extra pip packages. e.g. ["memray"]',
+    required=False,
+    default=None,
+)
 def deploy(
     name: Optional[str],
     config: str,
@@ -266,6 +259,8 @@ def deploy(
     cli_version: Optional[str],
     schedule: Optional[str],
     time_zone: str,
+    extra_pip: Optional[str],
+    debug: bool = False,
 ) -> None:
     """
     Deploy an ingestion recipe to your DataHub instance.
@@ -276,83 +271,23 @@ def deploy(
 
     datahub_graph = get_default_graph()
 
-    pipeline_config = load_config_file(
-        config,
-        allow_stdin=True,
-        allow_remote=True,
-        resolve_env_vars=False,
+    variables = deploy_source_vars(
+        name=name,
+        config=config,
+        urn=urn,
+        executor_id=executor_id,
+        cli_version=cli_version,
+        schedule=schedule,
+        time_zone=time_zone,
+        extra_pip=extra_pip,
+        debug=debug,
     )
 
-    deploy_options_raw = pipeline_config.pop("deployment", None)
-    if deploy_options_raw is not None:
-        deploy_options = DeployOptions.parse_obj(deploy_options_raw)
-
-        if name:
-            logger.info(f"Overriding deployment name {deploy_options.name} with {name}")
-            deploy_options.name = name
-    else:
-        if not name:
-            raise click.UsageError(
-                "Either --name must be set or deployment_name specified in the config"
-            )
-        deploy_options = DeployOptions(name=name)
-
-    # Use remaining CLI args to override deploy_options
-    if schedule:
-        deploy_options.schedule = schedule
-    if time_zone:
-        deploy_options.time_zone = time_zone
-    if cli_version:
-        deploy_options.cli_version = cli_version
-    if executor_id:
-        deploy_options.executor_id = executor_id
-
-    logger.info(f"Using {repr(deploy_options)}")
-
-    if not urn:
-        # When urn/name is not specified, we will generate a unique urn based on the deployment name.
-        urn = _make_ingestion_urn(deploy_options.name)
-        logger.info(f"Using recipe urn: {urn}")
-
-    # Invariant - at this point, both urn and deploy_options are set.
-
-    variables: dict = {
-        "urn": urn,
-        "name": deploy_options.name,
-        "type": pipeline_config["source"]["type"],
-        "recipe": json.dumps(pipeline_config),
-        "executorId": deploy_options.executor_id,
-        "version": deploy_options.cli_version,
-    }
-
-    if deploy_options.schedule is not None:
-        variables["schedule"] = {
-            "interval": deploy_options.schedule,
-            "timezone": deploy_options.time_zone,
-        }
-
     # The updateIngestionSource endpoint can actually do upserts as well.
     graphql_query: str = textwrap.dedent(
         """
-        mutation updateIngestionSource(
-            $urn: String!,
-            $name: String!,
-            $type: String!,
-            $schedule: UpdateIngestionSourceScheduleInput,
-            $recipe: String!,
-            $executorId: String!
-            $version: String) {
-
-            updateIngestionSource(urn: $urn, input: {
-                name: $name,
-                type: $type,
-                schedule: $schedule,
-                config: {
-                    recipe: $recipe,
-                    executorId: $executorId,
-                    version: $version,
-                }
-            })
+        mutation updateIngestionSource($urn: String!, $input: UpdateIngestionSourceInput!) {
+            updateIngestionSource(urn: $urn, input: $input)
         }
         """
     )
@@ -372,7 +307,7 @@ def deploy(
         sys.exit(1)
 
     click.echo(
-        f"✅ Successfully wrote data ingestion source metadata for recipe {deploy_options.name}:"
+        f"✅ Successfully wrote data ingestion source metadata for recipe {variables['input']['name']}:"
     )
     click.echo(response)
 
@@ -414,7 +349,9 @@ def parse_restli_response(response):
 
 
 @ingest.command()
-@click.argument("path", type=click.Path(exists=True))
+@click.argument(
+    "path", type=click.Path(exists=False)
+)  # exists=False since it only supports local filesystems
 def mcps(path: str) -> None:
     """
     Ingest metadata from a mcp json file or directory of files.

datahub/cli/lite_cli.py

@@ -285,10 +285,12 @@ def search(
     ctx: click.Context,
     query: str = "",
     flavor: str = SearchFlavor.FREE_TEXT.name.lower(),
-    aspect: List[str] = [],
+    aspect: Optional[List[str]] = None,
     details: bool = True,
 ) -> None:
     """Search with a free text or exact query string"""
+    if aspect is None:
+        aspect = []
 
     # query flavor should be sanitized by now, but we still need to convert it to a SearchFlavor
     try:
@@ -296,7 +298,7 @@ def search(
     except KeyError:
         raise click.UsageError(
             f"Failed to find a matching query flavor for {flavor}. Valid values are {[x.lower() for x in SearchFlavor._member_names_]}"
-        )
+        ) from None
     catalog = _get_datahub_lite(read_only=True)
     # sanitize query
     result_ids = set()

datahub/cli/specific/dataproduct_cli.py

@@ -49,7 +49,7 @@ def _abort_if_non_existent_urn(graph: DataHubGraph, urn: str, operation: str) ->
         entity_type = parsed_urn.get_type()
     except Exception:
         click.secho(f"Provided urn {urn} does not seem valid", fg="red")
-        raise click.Abort()
+        raise click.Abort() from None
     else:
         if not graph.exists(urn):
             click.secho(

datahub/cli/specific/dataset_cli.py

@@ -1,12 +1,15 @@
+import filecmp
 import json
 import logging
+import os
+import shutil
 from pathlib import Path
-from typing import Set, Tuple
+from typing import List, Set, Tuple
 
 import click
 from click_default_group import DefaultGroup
 
-from datahub.api.entities.dataset.dataset import Dataset
+from datahub.api.entities.dataset.dataset import Dataset, DatasetRetrievalConfig
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
 from datahub.metadata.com.linkedin.pegasus2avro.common import Siblings
@@ -30,18 +33,9 @@ def dataset() -> None:
 @telemetry.with_telemetry()
 def upsert(file: Path) -> None:
     """Upsert attributes to a Dataset in DataHub."""
-
-    with get_default_graph() as graph:
-        for dataset in Dataset.from_yaml(str(file)):
-            try:
-                for mcp in dataset.generate_mcp():
-                    graph.emit(mcp)
-                click.secho(f"Update succeeded for urn {dataset.urn}.", fg="green")
-            except Exception as e:
-                click.secho(
-                    f"Update failed for id {id}. due to {e}",
-                    fg="red",
-                )
+    # 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)
 
 
 @dataset.command(
@@ -111,3 +105,123 @@ def _get_existing_siblings(graph: DataHubGraph, urn: str) -> Set[str]:
         return set(existing.siblings)
     else:
         return set()
+
+
+@dataset.command(
+    name="file",
+)
+@click.option("--lintCheck", required=False, is_flag=True)
+@click.option("--lintFix", required=False, is_flag=True)
+@click.argument("file", type=click.Path(exists=True))
+@upgrade.check_upgrade
+@telemetry.with_telemetry()
+def file(lintcheck: bool, lintfix: bool, file: str) -> None:
+    """Operate on a Dataset file"""
+
+    if lintcheck or lintfix:
+        import tempfile
+        from pathlib import Path
+
+        # Create a temporary file in a secure way
+        # The file will be automatically deleted when the context manager exits
+        with tempfile.NamedTemporaryFile(suffix=".yml", delete=False) as temp:
+            temp_path = Path(temp.name)
+            try:
+                # Copy content to the temporary file
+                shutil.copyfile(file, temp_path)
+
+                # Run the linting
+                datasets = Dataset.from_yaml(temp.name)
+                for dataset in datasets:
+                    dataset.to_yaml(temp_path)
+
+                # Compare the files
+                files_match = filecmp.cmp(file, temp_path)
+
+                if files_match:
+                    click.secho("No differences found", fg="green")
+                else:
+                    # Show diff for visibility
+                    os.system(f"diff {file} {temp_path}")
+
+                    if lintfix:
+                        shutil.copyfile(temp_path, file)
+                        click.secho(f"Fixed linting issues in {file}", fg="green")
+                    else:
+                        click.secho(
+                            f"To fix these differences, run 'datahub dataset file --lintFix {file}'",
+                            fg="yellow",
+                        )
+            finally:
+                # Ensure the temporary file is removed
+                if temp_path.exists():
+                    temp_path.unlink()
+    else:
+        click.secho(
+            "No operation specified. Choose from --lintCheck or --lintFix", fg="yellow"
+        )
+
+
+@dataset.command(
+    name="sync",
+)
+@click.option("-f", "--file", required=True, type=click.Path(exists=True))
+@click.option("--to-datahub/--from-datahub", required=True, is_flag=True)
+@upgrade.check_upgrade
+@telemetry.with_telemetry()
+def sync(file: str, to_datahub: bool) -> None:
+    """Sync a Dataset file to/from DataHub"""
+
+    failures: List[str] = []
+    with get_default_graph() as graph:
+        datasets = Dataset.from_yaml(file)
+        for dataset in datasets:
+            assert (
+                dataset.urn is not None
+            )  # Validator should have ensured this is filled. Tell mypy it's not None
+            if to_datahub:
+                missing_entity_references = [
+                    entity_reference
+                    for entity_reference in dataset.entity_references()
+                    if not graph.exists(entity_reference)
+                ]
+                if missing_entity_references:
+                    click.secho(
+                        "\n\t- ".join(
+                            [
+                                f"Skipping Dataset {dataset.urn} due to missing entity references: "
+                            ]
+                            + missing_entity_references
+                        ),
+                        fg="red",
+                    )
+                    failures.append(dataset.urn)
+                    continue
+                try:
+                    for mcp in dataset.generate_mcp():
+                        graph.emit(mcp)
+                    click.secho(f"Update succeeded for urn {dataset.urn}.", fg="green")
+                except Exception as e:
+                    click.secho(
+                        f"Update failed for id {id}. due to {e}",
+                        fg="red",
+                    )
+            else:
+                # Sync from DataHub
+                if graph.exists(dataset.urn):
+                    dataset_get_config = DatasetRetrievalConfig()
+                    if dataset.downstreams:
+                        dataset_get_config.include_downstreams = True
+                    existing_dataset: Dataset = Dataset.from_datahub(
+                        graph=graph, urn=dataset.urn, config=dataset_get_config
+                    )
+                    existing_dataset.to_yaml(Path(file))
+                else:
+                    click.secho(f"Dataset {dataset.urn} does not exist")
+                    failures.append(dataset.urn)
+    if failures:
+        click.secho(
+            f"\nFailed to sync the following Datasets: {', '.join(failures)}",
+            fg="red",
+        )
+        raise click.Abort()

datahub/configuration/common.py

@@ -20,7 +20,7 @@ from pydantic import BaseModel, Extra, ValidationError
 from pydantic.fields import Field
 from typing_extensions import Protocol, Self
 
-from datahub.configuration._config_enum import ConfigEnum as ConfigEnum  # noqa: I250
+from datahub.configuration._config_enum import ConfigEnum as ConfigEnum
 from datahub.configuration.pydantic_migration_helpers import PYDANTIC_VERSION_2
 from datahub.utilities.dedup_list import deduplicate_list
 
@@ -130,7 +130,7 @@ class PermissiveConfigModel(ConfigModel):
     # It is usually used for argument bags that are passed through to third-party libraries.
 
     class Config:
-        if PYDANTIC_VERSION_2:
+        if PYDANTIC_VERSION_2:  # noqa: SIM108
             extra = "allow"
         else:
             extra = Extra.allow
@@ -198,6 +198,14 @@ class IgnorableError(MetaError):
     """An error that can be ignored."""
 
 
+class TraceTimeoutError(OperationalError):
+    """Failure to complete an API Trace within the timeout."""
+
+
+class TraceValidationError(OperationalError):
+    """Failure to complete the expected write operation."""
+
+
 @runtime_checkable
 class ExceptionWithProps(Protocol):
     def get_telemetry_props(self) -> Dict[str, Any]: ...

datahub/configuration/git.py

@@ -43,9 +43,7 @@ class GitReference(ConfigModel):
 
     @validator("repo", pre=True)
     def simplify_repo_url(cls, repo: str) -> str:
-        if repo.startswith("github.com/"):
-            repo = f"https://{repo}"
-        elif repo.startswith("gitlab.com"):
+        if repo.startswith("github.com/") or repo.startswith("gitlab.com"):
             repo = f"https://{repo}"
         elif repo.count("/") == 1:
             repo = f"https://github.com/{repo}"

datahub/configuration/kafka.py

@@ -44,7 +44,7 @@ class KafkaConsumerConnectionConfig(_KafkaConnectionConfig):
             try:
                 value = CallableConsumerConfig(value).callable_config()
             except Exception as e:
-                raise ConfigurationError(e)
+                raise ConfigurationError(e) from e
         return value