acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/api/entities/forms/forms.py

@@ -5,7 +5,7 @@ from pathlib import Path
 from typing import List, Optional, Union
 
 import yaml
-from pydantic import validator
+from pydantic import Field, validator
 from ruamel.yaml import YAML
 from typing_extensions import Literal
 
@@ -26,6 +26,7 @@ from datahub.emitter.mce_builder import (
 )
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.schema_classes import (
     FormActorAssignmentClass,
     FormInfoClass,
@@ -66,7 +67,7 @@ class Prompt(ConfigModel):
     description: Optional[str] = None
     type: str
     structured_property_id: Optional[str] = None
-    structured_property_urn: Optional[str] = None
+    structured_property_urn: Optional[str] = Field(default=None, validate_default=True)
     required: Optional[bool] = None
 
     @validator("structured_property_urn", pre=True, always=True)
@@ -110,7 +111,7 @@ class Actors(ConfigModel):
 
 class Forms(ConfigModel):
     id: Optional[str] = None
-    urn: Optional[str] = None
+    urn: Optional[str] = Field(default=None, validate_default=True)
     name: str
     description: Optional[str] = None
     prompts: List[Prompt] = []
@@ -133,47 +134,46 @@ class Forms(ConfigModel):
     def create(file: str) -> None:
         emitter: DataHubGraph
 
-        with get_default_graph() as emitter:
-            with open(file) as fp:
-                forms: List[dict] = yaml.safe_load(fp)
-                for form_raw in forms:
-                    form = Forms.parse_obj(form_raw)
+        with get_default_graph(ClientMode.CLI) as emitter, open(file) as fp:
+            forms: List[dict] = yaml.safe_load(fp)
+            for form_raw in forms:
+                form = Forms.parse_obj(form_raw)
 
-                    try:
-                        if not FormType.has_value(form.type):
-                            logger.error(
-                                f"Form type {form.type} does not exist. Please try again with a valid type."
-                            )
-
-                        mcp = MetadataChangeProposalWrapper(
-                            entityUrn=form.urn,
-                            aspect=FormInfoClass(
-                                name=form.name,
-                                description=form.description,
-                                prompts=form.validate_prompts(emitter),
-                                type=form.type,
-                                actors=form.create_form_actors(form.actors),
-                            ),
+                try:
+                    if not FormType.has_value(form.type):
+                        logger.error(
+                            f"Form type {form.type} does not exist. Please try again with a valid type."
                         )
-                        emitter.emit_mcp(mcp)
 
-                        logger.info(f"Created form {form.urn}")
+                    mcp = MetadataChangeProposalWrapper(
+                        entityUrn=form.urn,
+                        aspect=FormInfoClass(
+                            name=form.name,
+                            description=form.description,
+                            prompts=form.validate_prompts(emitter),
+                            type=form.type,
+                            actors=form.create_form_actors(form.actors),
+                        ),
+                    )
+                    emitter.emit_mcp(mcp)
+
+                    logger.info(f"Created form {form.urn}")
 
-                        if form.owners or form.group_owners:
-                            form.add_owners(emitter)
+                    if form.owners or form.group_owners:
+                        form.add_owners(emitter)
 
-                        if form.entities:
-                            if form.entities.urns:
-                                # Associate specific entities with a form
-                                form.upload_entities_for_form(emitter)
+                    if form.entities:
+                        if form.entities.urns:
+                            # Associate specific entities with a form
+                            form.upload_entities_for_form(emitter)
 
-                            if form.entities.filters:
-                                # Associate groups of entities with a form based on filters
-                                form.create_form_filters(emitter)
+                        if form.entities.filters:
+                            # Associate groups of entities with a form based on filters
+                            form.create_form_filters(emitter)
 
-                    except Exception as e:
-                        logger.error(e)
-                        return
+                except Exception as e:
+                    logger.error(e)
+                    return
 
     def validate_prompts(self, emitter: DataHubGraph) -> List[FormPromptClass]:
         prompts = []

datahub/api/entities/structuredproperties/structuredproperties.py

@@ -1,10 +1,10 @@
 import logging
 from enum import Enum
 from pathlib import Path
-from typing import Iterable, List, Optional, Union
+from typing import Iterable, List, Optional, Type, Union
 
 import yaml
-from pydantic import validator
+from pydantic import Field, StrictStr, validator
 from ruamel.yaml import YAML
 
 from datahub.configuration.common import ConfigModel
@@ -38,17 +38,17 @@ class AllowedTypes(Enum):
 
 
 class AllowedValue(ConfigModel):
-    value: Union[int, float, str]
+    value: Union[StrictStr, float]
     description: Optional[str] = None
 
 
 VALID_ENTITY_TYPE_URNS = [
-    Urn.make_entity_type_urn(entity_type) for entity_type in URN_TYPES.keys()
+    Urn.make_entity_type_urn(entity_type) for entity_type in URN_TYPES
 ]
 _VALID_ENTITY_TYPES_STRING = f"Valid entity type urns are {', '.join(VALID_ENTITY_TYPE_URNS)}, etc... Ensure that the entity type is valid."
 
 
-def _validate_entity_type_urn(v: str) -> str:
+def _validate_entity_type_urn(cls: Type, v: str) -> str:
     urn = Urn.make_entity_type_urn(v)
     if urn not in VALID_ENTITY_TYPE_URNS:
         raise ValueError(
@@ -68,7 +68,7 @@ class TypeQualifierAllowedTypes(ConfigModel):
 
 class StructuredProperties(ConfigModel):
     id: Optional[str] = None
-    urn: Optional[str] = None
+    urn: Optional[str] = Field(None, validate_default=True)
     qualified_name: Optional[str] = None
     type: str
     value_entity_types: Optional[List[str]] = None

datahub/api/graphql/assertion.py

@@ -65,7 +65,7 @@ query dataset($urn: String!, $start: Int, $count: Int, $status: AssertionRunStat
 
         :param urn: The DataHub dataset unique identifier.
         :param status: The assertion status to filter for. Every status will be accepted if it is not set.
-            See valid status at https://datahubproject.io/docs/graphql/enums#assertionrunstatus
+            See valid status at https://docs.datahub.com/docs/graphql/enums#assertionrunstatus
         :param start_time_millis: The start time in milliseconds from the assertions will be queried.
         :param end_time_millis: The end time in milliseconds until the assertions will be queried.
         :param filter: Additional key value filters which will be applied as AND query

datahub/api/graphql/base.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from gql import Client
 from gql.transport.requests import RequestsHTTPTransport
@@ -39,16 +39,18 @@ class BaseApi:
 
     def gen_filter(
         self, filters: Dict[str, Optional[str]]
-    ) -> Optional[Dict[str, List[Dict[str, str]]]]:
-        filter_expression: Optional[Dict[str, List[Dict[str, str]]]] = None
+    ) -> Optional[Dict[str, List[Dict[str, Union[str, List[str]]]]]]:
+        filter_expression: Optional[
+            Dict[str, List[Dict[str, Union[str, List[str]]]]]
+        ] = None
         if not filters:
             return None
 
-        filter = []
+        filter_list: List[Dict[str, Union[str, List[str]]]] = []
         for key, value in filters.items():
             if value is None:
                 continue
-            filter.append({"field": key, "value": value})
+            filter_list.append({"field": key, "values": [value]})
 
-        filter_expression = {"and": filter}
+        filter_expression = {"and": filter_list}
         return filter_expression

datahub/api/graphql/operation.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Any, Dict, List, Optional
 
-from gql import gql
+from gql import GraphQLRequest
 
 from datahub.api.graphql.base import BaseApi
 
@@ -55,10 +55,10 @@ mutation reportOperation($urn: String!, $sourceType: OperationSourceType!, $oper
         Report operation metadata for a dataset.
         :param source_type: The source type to filter on. If not set it will accept any source type.
             Default value: DATA_PROCESS
-            See valid types here: https://datahubproject.io/docs/graphql/enums#operationsourcetype
+            See valid types here: https://docs.datahub.com/docs/graphql/enums#operationsourcetype
         :param operation_type: The operation type to filter on. If not set it will accept any source type.
             Default value: "UPDATE"
-            See valid types here: https://datahubproject.io/docs/graphql/enums/#operationtype
+            See valid types here: https://docs.datahub.com/docs/graphql/enums/#operationtype
         :param partition: The partition to set the operation.
         :param num_affected_rows: The number of rows affected by this operation.
         :param custom_properties: Key/value pair of custom propertis
@@ -79,10 +79,12 @@ mutation reportOperation($urn: String!, $sourceType: OperationSourceType!, $oper
         if custom_properties is not None:
             variable_values["customProperties"] = custom_properties
 
-        result = self.client.execute(
-            gql(Operation.REPORT_OPERATION_MUTATION), variable_values
+        request = GraphQLRequest(
+            Operation.REPORT_OPERATION_MUTATION, variable_values=variable_values
         )
 
+        result = self.client.execute(request)
+
         return result["reportOperation"]
 
     def query_operations(
@@ -103,18 +105,18 @@ mutation reportOperation($urn: String!, $sourceType: OperationSourceType!, $oper
         :param end_time_millis: The end time in milliseconds until the operations will be queried.
         :param limit: The maximum number of items to return.
         :param source_type: The source type to filter on. If not set it will accept any source type.
-            See valid types here: https://datahubproject.io/docs/graphql/enums#operationsourcetype
+            See valid types here: https://docs.datahub.com/docs/graphql/enums#operationsourcetype
         :param operation_type: The operation type to filter on. If not set it will accept any source type.
-            See valid types here: https://datahubproject.io/docs/graphql/enums#operationsourcetype
+            See valid types here: https://docs.datahub.com/docs/graphql/enums#operationsourcetype
         :param partition: The partition to check the operation.
         """
 
-        result = self.client.execute(
-            gql(Operation.QUERY_OPERATIONS),
+        request = GraphQLRequest(
+            Operation.QUERY_OPERATIONS,
             variable_values={
                 "urn": urn,
                 "startTimeMillis": start_time_millis,
-                "end_time_millis": end_time_millis,
+                "endTimeMillis": end_time_millis,
                 "limit": limit,
                 "filter": self.gen_filter(
                     {
@@ -125,6 +127,8 @@ mutation reportOperation($urn: String!, $sourceType: OperationSourceType!, $oper
                 ),
             },
         )
+
+        result = self.client.execute(request)
         if "dataset" in result and "operations" in result["dataset"]:
             operations = []
             if source_type is not None:

datahub/cli/check_cli.py

@@ -9,6 +9,7 @@ from datetime import datetime
 from typing import Any, Dict, List, Optional, Union
 
 import click
+from tabulate import tabulate
 
 from datahub._version import __package_name__
 from datahub.cli.json_file import check_mce_file
@@ -16,11 +17,12 @@ from datahub.configuration import config_loader
 from datahub.configuration.common import AllowDenyPattern
 from datahub.emitter.mce_builder import DEFAULT_ENV
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.ingestion.run.pipeline import Pipeline
 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.upgrade import upgrade
 from datahub.utilities.file_backed_collections import (
     ConnectionWrapper,
     FileBackedDict,
@@ -46,7 +48,6 @@ def check() -> None:
 @click.option(
     "--unpack-mces", default=False, is_flag=True, help="Converts MCEs into MCPs"
 )
-@telemetry.with_telemetry()
 def metadata_file(json_file: str, rewrite: bool, unpack_mces: bool) -> None:
     """Check the schema of a metadata (MCE or MCP) JSON file."""
 
@@ -104,7 +105,6 @@ def metadata_file(json_file: str, rewrite: bool, unpack_mces: bool) -> None:
     default=(),
     help="[Advanced] Paths in the deepdiff object to ignore",
 )
-@telemetry.with_telemetry()
 def metadata_diff(
     actual_file: str, expected_file: str, verbose: bool, ignore_path: List[str]
 ) -> None:
@@ -141,7 +141,6 @@ def metadata_diff(
     type=str,
     default=None,
 )
-@telemetry.with_telemetry()
 def plugins(source: Optional[str], verbose: bool) -> None:
     """List the enabled ingestion plugins."""
 
@@ -233,7 +232,7 @@ def sql_format(sql: str, platform: str) -> None:
     default=True,
     help="Run in offline mode and disable schema-aware parsing.",
 )
-@telemetry.with_telemetry()
+@upgrade.check_upgrade
 def sql_lineage(
     sql: Optional[str],
     sql_file: Optional[str],
@@ -259,7 +258,7 @@ def sql_lineage(
 
     graph = None
     if online:
-        graph = get_default_graph()
+        graph = get_default_graph(ClientMode.CLI)
 
     lineage = create_lineage_sql_parsed_result(
         sql,
@@ -296,7 +295,6 @@ def sql_lineage(
     type=str,
     help="the input to validate",
 )
-@telemetry.with_telemetry()
 def test_allow_deny(config: str, input: str, pattern_key: str) -> None:
     """Test input string against AllowDeny pattern in a DataHub recipe.
 
@@ -345,7 +343,6 @@ def test_allow_deny(config: str, input: str, pattern_key: str) -> None:
     type=str,
     help="The input to validate",
 )
-@telemetry.with_telemetry()
 def test_path_spec(config: str, input: str, path_spec_key: str) -> None:
     """Test input path string against PathSpec patterns in a DataHub recipe.
 
@@ -470,10 +467,95 @@ WHERE
 
 
 @check.command()
+@upgrade.check_upgrade
 def server_config() -> None:
     """Print the server config."""
-    graph = get_default_graph()
+    graph = get_default_graph(ClientMode.CLI)
 
     server_config = graph.get_server_config()
 
     click.echo(pprint.pformat(server_config))
+
+
+@check.command()
+@click.option(
+    "--urn", required=False, help="The urn or urn pattern (supports % for wildcard)"
+)
+@click.option("--aspect", default=None, help="Filter to a specific aspect name.")
+@click.option(
+    "--start", type=int, default=None, help="Row number of sql store to restore from."
+)
+@click.option("--batch-size", type=int, default=None, help="How many rows to restore.")
+@click.option(
+    "--file",
+    required=False,
+    type=click.Path(exists=True, dir_okay=True, readable=True),
+    help="File absolute path containing URNs (one per line) to restore indices",
+)
+@upgrade.check_upgrade
+def restore_indices(
+    urn: Optional[str],
+    aspect: Optional[str],
+    start: Optional[int],
+    batch_size: Optional[int],
+    file: Optional[str],
+) -> None:
+    """Resync metadata changes into the search and graph indices."""
+    if urn is None and file is None:
+        raise click.UsageError("Either --urn or --file must be provided")
+    graph = get_default_graph(ClientMode.CLI)
+
+    graph.restore_indices(
+        urn_pattern=urn,
+        aspect=aspect,
+        start=start,
+        batch_size=batch_size,
+        file=file,
+    )
+
+
+@check.command()
+@upgrade.check_upgrade
+def get_kafka_consumer_offsets() -> None:
+    """Get Kafka consumer offsets from the DataHub API."""
+    graph = get_default_graph(ClientMode.CLI)
+    result = graph.get_kafka_consumer_offsets()
+
+    table_data = []
+    headers = [
+        "Topic",
+        "Consumer Group",
+        "Schema",
+        "Partition",
+        "Offset",
+        "Lag",
+        "Avg Lag",
+        "Max Lag",
+        "Total Lag",
+    ]
+
+    for topic, consumers in result.items():
+        for consumer_group, schemas in consumers.items():
+            for schema, data in schemas.items():
+                metrics = data.get("metrics", {})
+                partitions = data.get("partitions", {})
+
+                for partition, partition_data in partitions.items():
+                    table_data.append(
+                        [
+                            topic,
+                            consumer_group,
+                            schema,
+                            partition,
+                            partition_data.get("offset", "N/A"),
+                            partition_data.get("lag", "N/A"),
+                            metrics.get("avgLag", "N/A"),
+                            metrics.get("maxLag", "N/A"),
+                            metrics.get("totalLag", "N/A"),
+                        ]
+                    )
+
+    if table_data:
+        click.echo(tabulate(table_data, headers=headers, tablefmt="grid"))
+    else:
+        click.echo("No Kafka consumer offset data found.")

datahub/cli/cli_utils.py

@@ -3,6 +3,7 @@ import logging
 import time
 import typing
 from datetime import datetime
+from functools import wraps
 from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union
 
 import click
@@ -424,3 +425,65 @@ def ensure_has_system_metadata(
     props = metadata.properties
     props["clientId"] = datahub_version.__package_name__
     props["clientVersion"] = datahub_version.__version__
+
+
+def enable_auto_decorators(main_group: click.Group) -> None:
+    """
+    Enable automatic decorators for all click commands.
+    This wraps existing command callback functions to add upgrade and telemetry decorators.
+    """
+
+    def has_decorator(func: Any, module_pattern: str, function_pattern: str) -> bool:
+        """Check if function already has a specific decorator"""
+        if hasattr(func, "__wrapped__"):
+            current_func = func
+            while hasattr(current_func, "__wrapped__"):
+                # Check if this wrapper matches the module and function patterns
+                if (
+                    hasattr(current_func, "__module__")
+                    and module_pattern in current_func.__module__
+                    and hasattr(current_func, "__name__")
+                    and function_pattern in current_func.__name__
+                ):
+                    return True
+                current_func = current_func.__wrapped__
+        return False
+
+    def has_telemetry_decorator(func):
+        return has_decorator(func, "telemetry", "with_telemetry")
+
+    def wrap_command_callback(command_obj):
+        """Wrap a command's callback function to add decorators"""
+        if hasattr(command_obj, "callback") and command_obj.callback:
+            original_callback = command_obj.callback
+
+            # Import here to avoid circular imports
+            from datahub.telemetry import telemetry
+
+            decorated_callback = original_callback
+
+            if not has_telemetry_decorator(decorated_callback):
+                log.debug(
+                    f"Applying telemetry decorator to {original_callback.__module__}.{original_callback.__name__}"
+                )
+                decorated_callback = telemetry.with_telemetry()(decorated_callback)
+
+            # Preserve the original function's metadata
+            decorated_callback = wraps(original_callback)(decorated_callback)
+
+            command_obj.callback = decorated_callback
+
+    def wrap_group_commands(group_obj):
+        """Recursively wrap all commands in a group"""
+        if hasattr(group_obj, "commands"):
+            for _, command_obj in group_obj.commands.items():
+                if isinstance(command_obj, click.Group):
+                    # Recursively wrap sub-groups
+                    wrap_group_commands(command_obj)
+                else:
+                    # Wrap individual commands
+                    wrap_command_callback(command_obj)
+
+    wrap_group_commands(main_group)
+
+    log.debug("Auto-decorators enabled successfully")

datahub/cli/config_utils.py

@@ -11,14 +11,23 @@ import click
 import yaml
 from pydantic import BaseModel, ValidationError
 
-from datahub.cli.env_utils import get_boolean_env_variable
+from datahub.configuration.env_vars import (
+    get_gms_host,
+    get_gms_port,
+    get_gms_protocol,
+    get_gms_token,
+    get_gms_url,
+    get_skip_config,
+    get_system_client_id,
+    get_system_client_secret,
+)
 from datahub.ingestion.graph.config import DatahubClientConfig
 
 logger = logging.getLogger(__name__)
 
 CONDENSED_DATAHUB_CONFIG_PATH = "~/.datahubenv"
-DATAHUB_CONFIG_PATH = os.path.expanduser(CONDENSED_DATAHUB_CONFIG_PATH)
-DATAHUB_ROOT_FOLDER = os.path.expanduser("~/.datahub")
+DATAHUB_CONFIG_PATH: str = os.path.expanduser(CONDENSED_DATAHUB_CONFIG_PATH)
+DATAHUB_ROOT_FOLDER: str = os.path.expanduser("~/.datahub")
 ENV_SKIP_CONFIG = "DATAHUB_SKIP_CONFIG"
 
 ENV_DATAHUB_SYSTEM_CLIENT_ID = "DATAHUB_SYSTEM_CLIENT_ID"
@@ -36,15 +45,15 @@ class MissingConfigError(Exception):
 
 
 def get_system_auth() -> Optional[str]:
-    system_client_id = os.environ.get(ENV_DATAHUB_SYSTEM_CLIENT_ID)
-    system_client_secret = os.environ.get(ENV_DATAHUB_SYSTEM_CLIENT_SECRET)
+    system_client_id = get_system_client_id()
+    system_client_secret = get_system_client_secret()
     if system_client_id is not None and system_client_secret is not None:
         return f"Basic {system_client_id}:{system_client_secret}"
     return None
 
 
 def _should_skip_config() -> bool:
-    return get_boolean_env_variable(ENV_SKIP_CONFIG, False)
+    return get_skip_config()
 
 
 def persist_raw_datahub_config(config: dict) -> None:
@@ -67,11 +76,11 @@ class DatahubConfig(BaseModel):
 
 
 def _get_config_from_env() -> Tuple[Optional[str], Optional[str]]:
-    host = os.environ.get(ENV_METADATA_HOST)
-    port = os.environ.get(ENV_METADATA_PORT)
-    token = os.environ.get(ENV_METADATA_TOKEN)
-    protocol = os.environ.get(ENV_METADATA_PROTOCOL, "http")
-    url = os.environ.get(ENV_METADATA_HOST_URL)
+    host = get_gms_host()
+    port = get_gms_port()
+    token = get_gms_token()
+    protocol = get_gms_protocol()
+    url = get_gms_url()
     if port is not None:
         url = f"{protocol}://{host}:{port}"
         return url, token
@@ -108,7 +117,6 @@ def load_client_config() -> DatahubClientConfig:
         datahub_config: DatahubClientConfig = DatahubConfig.parse_obj(
             client_config_dict
         ).gms
-
         return datahub_config
     except ValidationError as e:
         click.echo(f"Error loading your {CONDENSED_DATAHUB_CONFIG_PATH}")

datahub/cli/container_cli.py

@@ -3,6 +3,7 @@ import logging
 import click
 
 from datahub.ingestion.source.apply.datahub_apply import apply_association_to_container
+from datahub.upgrade import upgrade
 
 logger = logging.getLogger(__name__)
 
@@ -16,6 +17,7 @@ def container() -> None:
 @container.command()
 @click.option("--container-urn", required=True, type=str)
 @click.option("--tag-urn", required=True, type=str)
+@upgrade.check_upgrade
 def tag(container_urn: str, tag_urn: str) -> None:
     """Add patch to add a tag to all datasets in a container"""
     apply_association_to_container(container_urn, tag_urn, "tag")
@@ -24,6 +26,7 @@ def tag(container_urn: str, tag_urn: str) -> None:
 @container.command()
 @click.option("--container-urn", required=True, type=str)
 @click.option("--term-urn", required=True, type=str)
+@upgrade.check_upgrade
 def term(container_urn: str, term_urn: str) -> None:
     """Add patch to add a term to all datasets in a container"""
     apply_association_to_container(container_urn, term_urn, "term")
@@ -32,6 +35,7 @@ def term(container_urn: str, term_urn: str) -> None:
 @container.command()
 @click.option("--container-urn", required=True, type=str)
 @click.option("--owner-urn", required=True, type=str)
+@upgrade.check_upgrade
 def owner(container_urn: str, owner_urn: str) -> None:
     """Add patch to add a owner to all datasets in a container"""
     apply_association_to_container(container_urn, owner_urn, "owner")
@@ -40,6 +44,7 @@ def owner(container_urn: str, owner_urn: str) -> None:
 @container.command()
 @click.option("--container-urn", required=True, type=str)
 @click.option("--domain-urn", required=True, type=str)
+@upgrade.check_upgrade
 def domain(container_urn: str, domain_urn: str) -> None:
     """Add patch to add a domain to all datasets in a container"""
     apply_association_to_container(container_urn, domain_urn, "domain")