acryl-datahub 0.14.1.13rc9__py3-none-any.whl → 0.15.0__py3-none-any.whl

{acryl_datahub-0.14.1.13rc9.dist-info → acryl_datahub-0.15.0.dist-info}/entry_points.txt

@@ -68,12 +68,13 @@ mode = datahub.ingestion.source.mode:ModeSource
 mongodb = datahub.ingestion.source.mongodb:MongoDBSource
 mssql = datahub.ingestion.source.sql.mssql:SQLServerSource
 mysql = datahub.ingestion.source.sql.mysql:MySQLSource
+neo4j = datahub.ingestion.source.neo4j.neo4j_source:Neo4jSource
 nifi = datahub.ingestion.source.nifi:NifiSource
 okta = datahub.ingestion.source.identity.okta:OktaSource
 openapi = datahub.ingestion.source.openapi:OpenApiSource
 oracle = datahub.ingestion.source.sql.oracle:OracleSource
 postgres = datahub.ingestion.source.sql.postgres:PostgresSource
-powerbi = datahub.ingestion.source.powerbi:PowerBiDashboardSource
+powerbi = datahub.ingestion.source.powerbi.powerbi:PowerBiDashboardSource
 powerbi-report-server = datahub.ingestion.source.powerbi_report_server:PowerBiReportServerDashboardSource
 preset = datahub.ingestion.source.preset:PresetSource
 presto = datahub.ingestion.source.sql.presto:PrestoSource

datahub/__init__.py

@@ -3,7 +3,7 @@ import warnings
 
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "0.14.1.13rc9"
+__version__ = "0.15.0"
 
 
 def is_dev_mode() -> bool:

datahub/api/entities/structuredproperties/structuredproperties.py

@@ -1,8 +1,7 @@
 import logging
-from contextlib import contextmanager
 from enum import Enum
 from pathlib import Path
-from typing import Generator, List, Optional
+from typing import List, Optional
 
 import yaml
 from pydantic import validator
@@ -10,39 +9,18 @@ from ruamel.yaml import YAML
 
 from datahub.configuration.common import ConfigModel
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.schema_classes import (
     PropertyValueClass,
     StructuredPropertyDefinitionClass,
 )
-from datahub.utilities.urns.urn import Urn
+from datahub.metadata.urns import StructuredPropertyUrn, Urn
+from datahub.utilities.urns._urn_base import URN_TYPES
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-class StructuredPropertiesConfig:
-    """Configuration class to hold the graph client"""
-
-    _graph: Optional[DataHubGraph] = None
-
-    @classmethod
-    @contextmanager
-    def use_graph(cls, graph: DataHubGraph) -> Generator[None, None, None]:
-        """Context manager to temporarily set a custom graph"""
-        previous_graph = cls._graph
-        cls._graph = graph
-        try:
-            yield
-        finally:
-            cls._graph = previous_graph
-
-    @classmethod
-    def get_graph(cls) -> DataHubGraph:
-        """Get the current graph, falling back to default if none set"""
-        return cls._graph if cls._graph is not None else get_default_graph()
-
-
 class AllowedTypes(Enum):
     STRING = "string"
     RICH_TEXT = "rich_text"
@@ -64,29 +42,28 @@ class AllowedValue(ConfigModel):
     description: Optional[str] = None
 
 
-VALID_ENTITY_TYPES_PREFIX_STRING = ", ".join(
-    [
-        f"urn:li:entityType:datahub.{x}"
-        for x in ["dataset", "dashboard", "dataFlow", "schemaField"]
-    ]
-)
-VALID_ENTITY_TYPES_STRING = f"Valid entity type urns are {VALID_ENTITY_TYPES_PREFIX_STRING}, etc... Ensure that the entity type is valid."
+VALID_ENTITY_TYPE_URNS = [
+    Urn.make_entity_type_urn(entity_type) for entity_type in URN_TYPES.keys()
+]
+_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:
+    urn = Urn.make_entity_type_urn(v)
+    if urn not in VALID_ENTITY_TYPE_URNS:
+        raise ValueError(
+            f"Input {v} is not a valid entity type urn. {_VALID_ENTITY_TYPES_STRING}"
+        )
+    v = str(urn)
+    return v
 
 
 class TypeQualifierAllowedTypes(ConfigModel):
     allowed_types: List[str]
 
-    @validator("allowed_types", each_item=True)
-    def validate_allowed_types(cls, v):
-        if v:
-            graph = StructuredPropertiesConfig.get_graph()
-            validated_urn = Urn.make_entity_type_urn(v)
-            if not graph.exists(validated_urn):
-                raise ValueError(
-                    f"Input {v} is not a valid entity type urn. {VALID_ENTITY_TYPES_STRING}"
-                )
-            v = str(validated_urn)
-        return v
+    _check_allowed_types = validator("allowed_types", each_item=True, allow_reuse=True)(
+        _validate_entity_type_urn
+    )
 
 
 class StructuredProperties(ConfigModel):
@@ -103,26 +80,36 @@ class StructuredProperties(ConfigModel):
     type_qualifier: Optional[TypeQualifierAllowedTypes] = None
     immutable: Optional[bool] = False
 
-    @validator("entity_types", each_item=True)
-    def validate_entity_types(cls, v):
-        if v:
-            graph = StructuredPropertiesConfig.get_graph()
-            validated_urn = Urn.make_entity_type_urn(v)
-            if not graph.exists(validated_urn):
-                raise ValueError(
-                    f"Input {v} is not a valid entity type urn. {VALID_ENTITY_TYPES_STRING}"
-                )
-            v = str(validated_urn)
+    _check_entity_types = validator("entity_types", each_item=True, allow_reuse=True)(
+        _validate_entity_type_urn
+    )
+
+    @validator("type")
+    def validate_type(cls, v: str) -> str:
+        # Convert to lowercase if needed
+        if not v.islower():
+            logger.warning(
+                f"Structured property type should be lowercase. Updated to {v.lower()}"
+            )
+            v = v.lower()
+
+        # Check if type is allowed
+        if not AllowedTypes.check_allowed_type(v):
+            raise ValueError(
+                f"Type {v} is not allowed. Allowed types are {AllowedTypes.values()}"
+            )
         return v
 
     @property
     def fqn(self) -> str:
         assert self.urn is not None
-        return (
-            self.qualified_name
-            or self.id
-            or Urn.create_from_string(self.urn).get_entity_id()[0]
-        )
+        id = StructuredPropertyUrn.from_string(self.urn).id
+        if self.qualified_name is not None:
+            # ensure that qualified name and ID match
+            assert (
+                self.qualified_name == id
+            ), "ID in the urn and the qualified_name must match"
+        return id
 
     @validator("urn", pre=True, always=True)
     def urn_must_be_present(cls, v, values):
@@ -133,100 +120,90 @@ class StructuredProperties(ConfigModel):
         return v
 
     @staticmethod
-    def create(file: str, graph: Optional[DataHubGraph] = None) -> None:
-        emitter: DataHubGraph = graph if graph else get_default_graph()
-        with StructuredPropertiesConfig.use_graph(emitter):
-            print("Using graph")
-            with open(file) as fp:
-                structuredproperties: List[dict] = yaml.safe_load(fp)
-                for structuredproperty_raw in structuredproperties:
-                    structuredproperty = StructuredProperties.parse_obj(
-                        structuredproperty_raw
-                    )
-                    if not structuredproperty.type.islower():
-                        structuredproperty.type = structuredproperty.type.lower()
-                        logger.warn(
-                            f"Structured property type should be lowercase. Updated to {structuredproperty.type}"
-                        )
-                    if not AllowedTypes.check_allowed_type(structuredproperty.type):
-                        raise ValueError(
-                            f"Type {structuredproperty.type} is not allowed. Allowed types are {AllowedTypes.values()}"
-                        )
-                    mcp = MetadataChangeProposalWrapper(
-                        entityUrn=structuredproperty.urn,
-                        aspect=StructuredPropertyDefinitionClass(
-                            qualifiedName=structuredproperty.fqn,
-                            valueType=Urn.make_data_type_urn(structuredproperty.type),
-                            displayName=structuredproperty.display_name,
-                            description=structuredproperty.description,
-                            entityTypes=[
-                                Urn.make_entity_type_urn(entity_type)
-                                for entity_type in structuredproperty.entity_types or []
-                            ],
-                            cardinality=structuredproperty.cardinality,
-                            immutable=structuredproperty.immutable,
-                            allowedValues=(
-                                [
-                                    PropertyValueClass(
-                                        value=v.value, description=v.description
-                                    )
-                                    for v in structuredproperty.allowed_values
-                                ]
-                                if structuredproperty.allowed_values
-                                else None
-                            ),
-                            typeQualifier=(
-                                {
-                                    "allowedTypes": structuredproperty.type_qualifier.allowed_types
-                                }
-                                if structuredproperty.type_qualifier
-                                else None
-                            ),
-                        ),
-                    )
-                    emitter.emit_mcp(mcp)
-
-                    logger.info(f"Created structured property {structuredproperty.urn}")
-
-    @classmethod
-    def from_datahub(cls, graph: DataHubGraph, urn: str) -> "StructuredProperties":
-        with StructuredPropertiesConfig.use_graph(graph):
-            structured_property: Optional[
-                StructuredPropertyDefinitionClass
-            ] = graph.get_aspect(urn, StructuredPropertyDefinitionClass)
-            if structured_property is None:
-                raise Exception(
-                    "StructuredPropertyDefinition aspect is None. Unable to create structured property."
-                )
-            return StructuredProperties(
-                urn=urn,
-                qualified_name=structured_property.qualifiedName,
-                display_name=structured_property.displayName,
-                type=structured_property.valueType,
-                description=structured_property.description,
-                entity_types=structured_property.entityTypes,
-                cardinality=structured_property.cardinality,
-                allowed_values=(
+    def from_yaml(file: str) -> List["StructuredProperties"]:
+        with open(file) as fp:
+            structuredproperties: List[dict] = yaml.safe_load(fp)
+
+        result: List[StructuredProperties] = []
+        for structuredproperty_raw in structuredproperties:
+            result.append(StructuredProperties.parse_obj(structuredproperty_raw))
+        return result
+
+    def generate_mcps(self) -> List[MetadataChangeProposalWrapper]:
+        mcp = MetadataChangeProposalWrapper(
+            entityUrn=self.urn,
+            aspect=StructuredPropertyDefinitionClass(
+                qualifiedName=self.fqn,
+                valueType=Urn.make_data_type_urn(self.type),
+                displayName=self.display_name,
+                description=self.description,
+                entityTypes=[
+                    Urn.make_entity_type_urn(entity_type)
+                    for entity_type in self.entity_types or []
+                ],
+                cardinality=self.cardinality,
+                immutable=self.immutable,
+                allowedValues=(
                     [
-                        AllowedValue(
-                            value=av.value,
-                            description=av.description,
-                        )
-                        for av in structured_property.allowedValues or []
+                        PropertyValueClass(value=v.value, description=v.description)
+                        for v in self.allowed_values
                     ]
-                    if structured_property.allowedValues is not None
+                    if self.allowed_values
                     else None
                 ),
-                type_qualifier=(
-                    {
-                        "allowed_types": structured_property.typeQualifier.get(
-                            "allowedTypes"
-                        )
-                    }
-                    if structured_property.typeQualifier
+                typeQualifier=(
+                    {"allowedTypes": self.type_qualifier.allowed_types}
+                    if self.type_qualifier
                     else None
                 ),
+            ),
+        )
+        return [mcp]
+
+    @staticmethod
+    def create(file: str, graph: DataHubGraph) -> None:
+        # TODO: Deprecate this method.
+        structuredproperties = StructuredProperties.from_yaml(file)
+        for structuredproperty in structuredproperties:
+            for mcp in structuredproperty.generate_mcps():
+                graph.emit_mcp(mcp)
+
+            logger.info(f"Created structured property {structuredproperty.urn}")
+
+    @classmethod
+    def from_datahub(cls, graph: DataHubGraph, urn: str) -> "StructuredProperties":
+        structured_property: Optional[
+            StructuredPropertyDefinitionClass
+        ] = graph.get_aspect(urn, StructuredPropertyDefinitionClass)
+        if structured_property is None:
+            raise Exception(
+                "StructuredPropertyDefinition aspect is None. Unable to create structured property."
             )
+        return StructuredProperties(
+            urn=urn,
+            qualified_name=structured_property.qualifiedName,
+            display_name=structured_property.displayName,
+            type=structured_property.valueType,
+            description=structured_property.description,
+            entity_types=structured_property.entityTypes,
+            cardinality=structured_property.cardinality,
+            allowed_values=(
+                [
+                    AllowedValue(
+                        value=av.value,
+                        description=av.description,
+                    )
+                    for av in structured_property.allowedValues or []
+                ]
+                if structured_property.allowedValues is not None
+                else None
+            ),
+            type_qualifier=(
+                {"allowed_types": structured_property.typeQualifier.get("allowedTypes")}
+                if structured_property.typeQualifier
+                else None
+            ),
+        )
 
     def to_yaml(
         self,

datahub/cli/cli_utils.py

@@ -327,6 +327,8 @@ def _ensure_valid_gms_url_acryl_cloud(url: str) -> str:
         url = f"{url}/gms"
     elif url.endswith("acryl.io/"):
         url = f"{url}gms"
+    if url.endswith("acryl.io/api/gms"):
+        url = url.replace("acryl.io/api/gms", "acryl.io/gms")
 
     return url
 

datahub/cli/delete_cli.py

@@ -1,4 +1,5 @@
 import logging
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
 from datetime import datetime
 from random import choices
@@ -214,14 +215,47 @@ def references(urn: str, dry_run: bool, force: bool) -> None:
 
 
 @delete.command()
-@click.option("--urn", required=True, type=str, help="the urn of the entity")
-def undo_by_filter(urn: str) -> None:
+@click.option("--urn", required=False, type=str, help="the urn of the entity")
+@click.option(
+    "-p",
+    "--platform",
+    required=False,
+    type=str,
+    help="Platform filter (e.g. snowflake)",
+)
+@click.option(
+    "-b",
+    "--batch-size",
+    required=False,
+    default=3000,
+    type=int,
+    help="Batch size when querying for entities to un-soft delete."
+    "Maximum 10000. Large batch sizes may cause timeouts.",
+)
+def undo_by_filter(
+    urn: Optional[str], platform: Optional[str], batch_size: int
+) -> None:
     """
-    Undo a soft deletion of an entity
+    Undo soft deletion by filters
     """
     graph = get_default_graph()
     logger.info(f"Using {graph}")
-    graph.set_soft_delete_status(urn=urn, delete=False)
+    if urn:
+        graph.set_soft_delete_status(urn=urn, delete=False)
+    else:
+        urns = list(
+            graph.get_urns_by_filter(
+                platform=platform,
+                query="*",
+                status=RemovedStatusFilter.ONLY_SOFT_DELETED,
+                batch_size=batch_size,
+            )
+        )
+        logger.info(f"Going to un-soft delete {len(urns)} urns")
+        urns_iter = progressbar.progressbar(urns, redirect_stdout=True)
+        for urn in urns_iter:
+            assert urn
+            graph.set_soft_delete_status(urn=urn, delete=False)
 
 
 @delete.command(no_args_is_help=True)
@@ -312,6 +346,9 @@ def undo_by_filter(urn: str) -> None:
     default=False,
     help="Only delete soft-deleted entities, for hard deletion",
 )
+@click.option(
+    "--workers", type=int, default=1, help="Num of workers to use for deletion."
+)
 @upgrade.check_upgrade
 @telemetry.with_telemetry()
 def by_filter(
@@ -329,6 +366,7 @@ def by_filter(
     batch_size: int,
     dry_run: bool,
     only_soft_deleted: bool,
+    workers: int = 1,
 ) -> None:
     """Delete metadata from datahub using a single urn or a combination of filters."""
 
@@ -349,16 +387,19 @@ def by_filter(
     # TODO: add some validation on entity_type
 
     if not force and not soft and not dry_run:
+        message = (
+            "Hard deletion will permanently delete data from DataHub and can be slow. "
+            "We generally recommend using soft deletes instead. "
+            "Do you want to continue?"
+        )
         if only_soft_deleted:
             click.confirm(
-                "This will permanently delete data from DataHub. Do you want to continue?",
+                message,
                 abort=True,
             )
         else:
             click.confirm(
-                "Hard deletion will permanently delete data from DataHub and can be slow. "
-                "We generally recommend using soft deletes instead. "
-                "Do you want to continue?",
+                message,
                 abort=True,
             )
 
@@ -429,26 +470,64 @@ def by_filter(
                 abort=True,
             )
 
-    urns_iter = urns
-    if not delete_by_urn and not dry_run:
-        urns_iter = progressbar.progressbar(urns, redirect_stdout=True)
+    _delete_urns_parallel(
+        graph=graph,
+        urns=urns,
+        aspect_name=aspect,
+        soft=soft,
+        dry_run=dry_run,
+        delete_by_urn=delete_by_urn,
+        start_time=start_time,
+        end_time=end_time,
+        workers=workers,
+    )
+
 
-    # Run the deletion.
+def _delete_urns_parallel(
+    graph: DataHubGraph,
+    urns: List[str],
+    delete_by_urn: bool,
+    start_time: Optional[datetime],
+    end_time: Optional[datetime],
+    aspect_name: Optional[str] = None,
+    soft: bool = True,
+    dry_run: bool = False,
+    workers: int = 1,
+) -> None:
     deletion_result = DeletionResult()
-    with PerfTimer() as timer:
-        for urn in urns_iter:
-            one_result = _delete_one_urn(
-                graph=graph,
-                urn=urn,
-                aspect_name=aspect,
-                soft=soft,
-                dry_run=dry_run,
-                start_time=start_time,
-                end_time=end_time,
+
+    def process_urn(urn):
+        return _delete_one_urn(
+            graph=graph,
+            urn=urn,
+            aspect_name=aspect_name,
+            soft=soft,
+            dry_run=dry_run,
+            start_time=start_time,
+            end_time=end_time,
+        )
+
+    with PerfTimer() as timer, ThreadPoolExecutor(max_workers=workers) as executor:
+        future_to_urn = {executor.submit(process_urn, urn): urn for urn in urns}
+
+        completed_futures = as_completed(future_to_urn)
+        if not delete_by_urn and not dry_run:
+            futures_iter = progressbar.progressbar(
+                as_completed(future_to_urn),
+                max_value=len(future_to_urn),
+                redirect_stdout=True,
             )
-            deletion_result.merge(one_result)
+        else:
+            futures_iter = completed_futures
+
+        for future in futures_iter:
+            try:
+                one_result = future.result()
+                deletion_result.merge(one_result)
+            except Exception as e:
+                urn = future_to_urn[future]
+                click.secho(f"Error processing URN {urn}: {e}", fg="red")
 
-    # Report out a summary of the deletion result.
     click.echo(
         deletion_result.format_message(
             dry_run=dry_run, soft=soft, time_sec=timer.elapsed_seconds()

datahub/cli/ingest_cli.py

@@ -27,6 +27,7 @@ from datahub.utilities.perf_timer import PerfTimer
 
 logger = logging.getLogger(__name__)
 
+INGEST_SRC_TABLE_COLUMNS = ["runId", "source", "startTime", "status", "URN"]
 RUNS_TABLE_COLUMNS = ["runId", "rows", "created at"]
 RUN_TABLE_COLUMNS = ["urn", "aspect name", "created at"]
 
@@ -437,6 +438,115 @@ def mcps(path: str) -> None:
     sys.exit(ret)
 
 
+@ingest.command()
+@click.argument("page_offset", type=int, default=0)
+@click.argument("page_size", type=int, default=100)
+@click.option("--urn", type=str, default=None, help="Filter by ingestion source URN.")
+@click.option(
+    "--source", type=str, default=None, help="Filter by ingestion source name."
+)
+@upgrade.check_upgrade
+@telemetry.with_telemetry()
+def list_source_runs(page_offset: int, page_size: int, urn: str, source: str) -> None:
+    """List ingestion source runs with their details, optionally filtered by URN or source."""
+
+    query = """
+    query listIngestionRuns($input: ListIngestionSourcesInput!) {
+      listIngestionSources(input: $input) {
+        ingestionSources {
+          urn
+          name
+          executions {
+            executionRequests {
+              id
+              result {
+                startTimeMs
+                status
+              }
+            }
+          }
+        }
+      }
+    }
+    """
+
+    # filter by urn and/or source using CONTAINS
+    filters = []
+    if urn:
+        filters.append({"field": "urn", "values": [urn], "condition": "CONTAIN"})
+    if source:
+        filters.append({"field": "name", "values": [source], "condition": "CONTAIN"})
+
+    variables = {
+        "input": {
+            "start": page_offset,
+            "count": page_size,
+            "filters": filters,
+        }
+    }
+
+    client = get_default_graph()
+    session = client._session
+    gms_host = client.config.server
+
+    url = f"{gms_host}/api/graphql"
+    try:
+        response = session.post(url, json={"query": query, "variables": variables})
+        response.raise_for_status()
+    except Exception as e:
+        click.echo(f"Error fetching data: {str(e)}")
+        return
+
+    try:
+        data = response.json()
+    except ValueError:
+        click.echo("Failed to parse JSON response from server.")
+        return
+
+    if not data:
+        click.echo("No response received from the server.")
+        return
+
+    # when urn or source filter does not match, exit gracefully
+    if (
+        not isinstance(data.get("data"), dict)
+        or "listIngestionSources" not in data["data"]
+    ):
+        click.echo("No matching ingestion sources found. Please check your filters.")
+        return
+
+    ingestion_sources = data["data"]["listIngestionSources"]["ingestionSources"]
+    if not ingestion_sources:
+        click.echo("No ingestion sources or executions found.")
+        return
+
+    rows = []
+    for ingestion_source in ingestion_sources:
+        urn = ingestion_source.get("urn", "N/A")
+        name = ingestion_source.get("name", "N/A")
+
+        executions = ingestion_source.get("executions", {}).get("executionRequests", [])
+        for execution in executions:
+            execution_id = execution.get("id", "N/A")
+            start_time = execution.get("result", {}).get("startTimeMs", "N/A")
+            start_time = (
+                datetime.fromtimestamp(start_time / 1000).strftime("%Y-%m-%d %H:%M:%S")
+                if start_time != "N/A"
+                else "N/A"
+            )
+            status = execution.get("result", {}).get("status", "N/A")
+
+            rows.append([execution_id, name, start_time, status, urn])
+
+    click.echo(
+        tabulate(
+            rows,
+            headers=INGEST_SRC_TABLE_COLUMNS,
+            tablefmt="grid",
+        )
+    )
+
+
 @ingest.command()
 @click.argument("page_offset", type=int, default=0)
 @click.argument("page_size", type=int, default=100)

datahub/cli/put_cli.py

@@ -105,7 +105,7 @@ def platform(
     """
 
     if name.startswith(f"urn:li:{DataPlatformUrn.ENTITY_TYPE}"):
-        platform_urn = DataPlatformUrn.create_from_string(name)
+        platform_urn = DataPlatformUrn.from_string(name)
         platform_name = platform_urn.get_entity_id_as_string()
     else:
         platform_name = name.lower()