acryl-datahub 1.0.0rc3__py3-none-any.whl → 1.0.0rc5__py3-none-any.whl

datahub/sdk/main_client.py

@@ -41,10 +41,24 @@ class DataHubClient:
 
     @classmethod
     def from_env(cls) -> "DataHubClient":
+        """Initialize a DataHubClient from the environment variables or ~/.datahubenv file.
+
+        This will first check DATAHUB_GMS_URL and DATAHUB_GMS_TOKEN. If not present,
+        it will read credentials from ~/.datahubenv. That file can be created using
+        the `datahub init` command.
+
+        If you're looking to specify the server/token in code, use the
+        DataHubClient(server=..., token=...) constructor instead.
+
+        Returns:
+            A DataHubClient instance.
+        """
+
         # Inspired by the DockerClient.from_env() method.
         # TODO: This one also reads from ~/.datahubenv, so the "from_env" name might be a bit confusing.
         # That file is part of the "environment", but is not a traditional "env variable".
         graph = get_default_graph()
+
         return cls(graph=graph)
 
     @property
@@ -54,3 +68,6 @@ class DataHubClient:
     @property
     def resolve(self) -> ResolverClient:
         return ResolverClient(self)
+
+    # TODO: search client
+    # TODO: lineage client

datahub/specific/dataset.py

@@ -15,6 +15,7 @@ from datahub.metadata.schema_classes import (
     UpstreamClass as Upstream,
     UpstreamLineageClass as UpstreamLineage,
 )
+from datahub.metadata.urns import DatasetUrn, TagUrn, Urn
 from datahub.specific.aspect_helpers.custom_properties import HasCustomPropertiesPatch
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
 from datahub.specific.aspect_helpers.structured_properties import (
@@ -22,8 +23,6 @@ from datahub.specific.aspect_helpers.structured_properties import (
 )
 from datahub.specific.aspect_helpers.tags import HasTagsPatch
 from datahub.specific.aspect_helpers.terms import HasTermsPatch
-from datahub.utilities.urns.tag_urn import TagUrn
-from datahub.utilities.urns.urn import Urn
 
 _Parent = TypeVar("_Parent", bound=MetadataPatchProposal)
 
@@ -104,12 +103,12 @@ class DatasetPatchBuilder(
 ):
     def __init__(
         self,
-        urn: str,
+        urn: Union[str, DatasetUrn],
         system_metadata: Optional[SystemMetadataClass] = None,
         audit_header: Optional[KafkaAuditHeaderClass] = None,
     ) -> None:
         super().__init__(
-            urn, system_metadata=system_metadata, audit_header=audit_header
+            str(urn), system_metadata=system_metadata, audit_header=audit_header
         )
 
     @classmethod

datahub/sql_parsing/split_statements.py

@@ -8,11 +8,11 @@ END_KEYWORD = "END"
 
 CONTROL_FLOW_KEYWORDS = [
     "GO",
-    r"BEGIN\w+TRY",
-    r"BEGIN\w+CATCH",
+    r"BEGIN\s+TRY",
+    r"BEGIN\s+CATCH",
     "BEGIN",
-    r"END\w+TRY",
-    r"END\w+CATCH",
+    r"END\s+TRY",
+    r"END\s+CATCH",
     # This isn't strictly correct, but we assume that IF | (condition) | (block) should all be split up
     # This mainly ensures that IF statements don't get tacked onto the previous statement incorrectly
     "IF",
@@ -73,25 +73,31 @@ class _StatementSplitter:
         # what a given END is closing.
         self.current_case_statements = 0
 
-    def _is_keyword_at_position(self, pos: int, keyword: str) -> bool:
+    def _is_keyword_at_position(self, pos: int, keyword: str) -> Tuple[bool, str]:
         """
         Check if a keyword exists at the given position using regex word boundaries.
         """
         sql = self.sql
 
-        if pos + len(keyword) > len(sql):
-            return False
+        keyword_length = len(keyword.replace(r"\s+", " "))
+
+        if pos + keyword_length > len(sql):
+            return False, ""
 
         # If we're not at a word boundary, we can't generate a keyword.
         if pos > 0 and not (
             bool(re.match(r"\w\W", sql[pos - 1 : pos + 1]))
             or bool(re.match(r"\W\w", sql[pos - 1 : pos + 1]))
         ):
-            return False
+            return False, ""
 
-        pattern = rf"^{re.escape(keyword)}\b"
+        pattern = rf"^{keyword}\b"
         match = re.match(pattern, sql[pos:], re.IGNORECASE)
-        return bool(match)
+        is_match = bool(match)
+        actual_match = (
+            sql[pos:][match.start() : match.end()] if match is not None else ""
+        )
+        return is_match, actual_match
 
     def _look_ahead_for_keywords(self, keywords: List[str]) -> Tuple[bool, str, int]:
         """
@@ -99,7 +105,8 @@ class _StatementSplitter:
         """
 
         for keyword in keywords:
-            if self._is_keyword_at_position(self.i, keyword):
+            is_match, keyword = self._is_keyword_at_position(self.i, keyword)
+            if is_match:
                 return True, keyword, len(keyword)
         return False, "", 0
 
@@ -118,7 +125,7 @@ class _StatementSplitter:
 
     def process(self) -> Iterator[str]:
         if not self.sql or not self.sql.strip():
-            return
+            yield from ()
 
         prev_real_char = "\0"  # the most recent non-whitespace, non-comment character
         while self.i < len(self.sql):
@@ -181,7 +188,7 @@ class _StatementSplitter:
     def _process_normal(self, most_recent_real_char: str) -> Iterator[str]:
         c = self.sql[self.i]
 
-        if self._is_keyword_at_position(self.i, CASE_KEYWORD):
+        if self._is_keyword_at_position(self.i, CASE_KEYWORD)[0]:
             self.current_case_statements += 1
 
         is_control_keyword, keyword, keyword_len = self._look_ahead_for_keywords(

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/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