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

datahub/testing/pytest_hooks.py

@@ -0,0 +1,56 @@
+import dataclasses
+from typing import Optional
+
+import pytest
+
+__all__ = [
+    "load_golden_flags",
+    "get_golden_settings",
+    "pytest_addoption",
+    "GoldenFileSettings",
+]
+
+
+@dataclasses.dataclass
+class GoldenFileSettings:
+    update_golden: bool
+    copy_output: bool
+
+
+_registered: bool = False
+_settings: Optional[GoldenFileSettings] = None
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption(
+        "--update-golden-files",
+        action="store_true",
+        default=False,
+    )
+
+    # TODO: Deprecate and remove this flag.
+    parser.addoption("--copy-output-files", action="store_true", default=False)
+
+    global _registered
+    _registered = True
+
+
+@pytest.fixture(scope="session", autouse=True)
+def load_golden_flags(pytestconfig: pytest.Config) -> None:
+    global _settings
+    _settings = GoldenFileSettings(
+        update_golden=pytestconfig.getoption("--update-golden-files"),
+        copy_output=pytestconfig.getoption("--copy-output-files"),
+    )
+
+
+def get_golden_settings() -> GoldenFileSettings:
+    if not _registered:
+        raise ValueError(
+            "Golden files aren't set up properly. Call register_golden_flags from a conftest pytest_addoptions method."
+        )
+    if not _settings:
+        raise ValueError(
+            "Golden files aren't set up properly. Ensure load_golden_flags is imported in your conftest."
+        )
+    return _settings

datahub/upgrade/upgrade.py

@@ -293,9 +293,9 @@ def is_client_server_compatible(client: VersionStats, server: VersionStats) -> i
         return server.version.micro - client.version.micro
 
 
-def _maybe_print_upgrade_message(  # noqa: C901
+def _maybe_print_upgrade_message(
     version_stats: Optional[DataHubVersionStats],
-) -> None:  # noqa: C901
+) -> None:
     days_before_cli_stale = 7
     days_before_quickstart_stale = 7
 

datahub/utilities/file_backed_collections.py

@@ -10,13 +10,11 @@ import tempfile
 import threading
 from dataclasses import dataclass, field
 from datetime import datetime
-from enum import Enum
 from types import TracebackType
 from typing import (
     Any,
     Callable,
     Dict,
-    Final,
     Generic,
     Iterator,
     List,
@@ -31,6 +29,7 @@ from typing import (
 )
 
 from datahub.ingestion.api.closeable import Closeable
+from datahub.utilities.sentinels import Unset, unset
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -59,16 +58,6 @@ SqliteValue = Union[int, float, str, bytes, datetime, None]
 _VT = TypeVar("_VT")
 
 
-class Unset(Enum):
-    token = 0
-
-
-# It's pretty annoying to create a true sentinel that works with typing.
-# https://peps.python.org/pep-0484/#support-for-singleton-types-in-unions
-# Can't wait for https://peps.python.org/pep-0661/
-_unset: Final = Unset.token
-
-
 class ConnectionWrapper:
     """
     Wraps a SQlite connection, allowing connection reuse across multiple FileBacked* objects.
@@ -372,7 +361,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
         self,
         /,
         key: str,
-        default: Union[_VT, Unset] = _unset,
+        default: Union[_VT, Unset] = unset,
     ) -> _VT:
         # If key is in the dictionary, this is similar to __getitem__ + mark_dirty.
         # If key is not in the dictionary, this is similar to __setitem__.
@@ -383,7 +372,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
             self.mark_dirty(key)
             return value
         except KeyError:
-            if default is _unset:
+            if default is unset:
                 raise
 
             self[key] = default

datahub/utilities/ingest_utils.py

@@ -0,0 +1,106 @@
+import json
+import logging
+from typing import Optional
+
+import click
+
+from datahub.configuration.common import ConfigModel
+from datahub.configuration.config_loader import load_config_file
+from datahub.emitter.mce_builder import datahub_guid
+
+logger = logging.getLogger(__name__)
+
+
+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"
+
+
+def deploy_source_vars(
+    name: Optional[str],
+    config: str,
+    urn: Optional[str],
+    executor_id: str,
+    cli_version: Optional[str],
+    schedule: Optional[str],
+    time_zone: str,
+    extra_pip: Optional[str],
+    debug: bool = False,
+) -> dict:
+    pipeline_config = load_config_file(
+        config,
+        allow_stdin=True,
+        allow_remote=True,
+        resolve_env_vars=False,
+    )
+
+    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}")
+
+    variables: dict = {
+        "urn": urn,
+        "input": {
+            "name": deploy_options.name,
+            "type": pipeline_config["source"]["type"],
+            "config": {
+                "recipe": json.dumps(pipeline_config),
+                "executorId": deploy_options.executor_id,
+                "debugMode": debug,
+                "version": deploy_options.cli_version,
+            },
+        },
+    }
+
+    if deploy_options.schedule is not None:
+        variables["input"]["schedule"] = {
+            "interval": deploy_options.schedule,
+            "timezone": deploy_options.time_zone,
+        }
+    if extra_pip is not None:
+        extra_args_list = (
+            variables.get("input", {}).get("config", {}).get("extraArgs", [])
+        )
+        extra_args_list.append({"key": "extra_pip_requirements", "value": extra_pip})
+        variables["input"]["config"]["extraArgs"] = extra_args_list
+
+    return variables

datahub/utilities/mapping.py

@@ -171,7 +171,7 @@ class OperationProcessor:
         self.owner_source_type = owner_source_type
         self.match_nested_props = match_nested_props
 
-    def process(self, raw_props: Mapping[str, Any]) -> Dict[str, Any]:  # noqa: C901
+    def process(self, raw_props: Mapping[str, Any]) -> Dict[str, Any]:
         # Defining the following local variables -
         # operations_map - the final resulting map when operations are processed.
         # Against each operation the values to be applied are stored.

datahub/utilities/memory_footprint.py

@@ -1,10 +1,10 @@
 from collections import deque
 from itertools import chain
 from sys import getsizeof
-from typing import Any, Iterator
+from typing import Any, Iterator, Optional
 
 
-def total_size(o: Any, handlers: Any = {}) -> int:
+def total_size(o: Any, handlers: Optional[Any] = None) -> int:
     """Returns the approximate memory footprint an object and all of its contents.
     Automatically finds the contents of the following builtin containers and
     their subclasses:  tuple, list, deque, dict, set and frozenset.
@@ -14,6 +14,7 @@ def total_size(o: Any, handlers: Any = {}) -> int:
 
     Based on https://github.com/ActiveState/recipe-577504-compute-mem-footprint/blob/master/recipe.py
     """
+    handlers = handlers or {}
 
     def dict_handler(d: dict) -> Iterator[Any]:
         return chain.from_iterable(d.items())

datahub/utilities/sentinels.py

@@ -0,0 +1,22 @@
+from enum import Enum
+from typing import Final
+
+# It's pretty annoying to create a true sentinel that works with typing.
+# This approach using enums is inspired by:
+# https://peps.python.org/pep-0484/#support-for-singleton-types-in-unions
+#
+# Can't wait for https://peps.python.org/pep-0661/
+
+
+class Unset(Enum):
+    token = 0
+
+
+unset: Final = Unset.token
+
+
+class Auto(Enum):
+    token = 0
+
+
+auto: Final = Auto.token

datahub/utilities/unified_diff.py

@@ -2,8 +2,12 @@ import logging
 from dataclasses import dataclass
 from typing import List, Tuple
 
+from datahub.cli.env_utils import get_boolean_env_variable
+
+_debug_diff = get_boolean_env_variable("DATAHUB_DEBUG_DIFF_PATCHER")
+
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
+logger.setLevel(logging.DEBUG if _debug_diff else logging.INFO)
 
 _LOOKAROUND_LINES = 300