acryl-datahub 1.0.0.2rc4__py3-none-any.whl → 1.0.0.3__py3-none-any.whl

{acryl_datahub-1.0.0.2rc4.dist-info → acryl_datahub-1.0.0.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (80.3.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

datahub/_version.py

@@ -1,6 +1,6 @@
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "1.0.0.2rc4"
+__version__ = "1.0.0.3"
 
 
 def is_dev_mode() -> bool:

datahub/api/circuit_breaker/operation_circuit_breaker.py

@@ -55,9 +55,9 @@ class OperationCircuitBreaker(AbstractCircuitBreaker):
             which is set as Airflow connection.
         :param partition: The partition to check the operation.
         :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/#operationtype
+            See valid types here: https://docs.datahub.com/docs/graphql/enums/#operationtype
         """
 
         start_time_millis: int = int(

datahub/api/entities/datacontract/datacontract.py

@@ -1,5 +1,5 @@
 import collections
-from typing import Iterable, List, Optional, Tuple
+from typing import Dict, Iterable, List, Optional, Tuple, Union
 
 from ruamel.yaml import YAML
 from typing_extensions import Literal
@@ -25,6 +25,8 @@ from datahub.metadata.schema_classes import (
     FreshnessContractClass,
     SchemaContractClass,
     StatusClass,
+    StructuredPropertiesClass,
+    StructuredPropertyValueAssignmentClass,
 )
 from datahub.utilities.urns.urn import guess_entity_type
 
@@ -47,8 +49,12 @@ class DataContract(v1_ConfigModel):
     entity: str = v1_Field(
         description="The entity urn that the Data Contract is associated with"
     )
-    # TODO: add support for properties
-    # properties: Optional[Dict[str, str]] = None
+    properties: Optional[Dict[str, Union[str, float, List[Union[str, float]]]]] = (
+        v1_Field(
+            default=None,
+            description="Structured properties associated with the data contract.",
+        )
+    )
 
     schema_field: Optional[SchemaAssertion] = v1_Field(default=None, alias="schema")
 
@@ -172,6 +178,30 @@ class DataContract(v1_ConfigModel):
         )
         yield from dq_assertion_mcps
 
+        # Construct the structured properties aspect if properties are defined
+        structured_properties_aspect: Optional[StructuredPropertiesClass] = None
+        if self.properties:
+            property_assignments: List[StructuredPropertyValueAssignmentClass] = []
+            for key, value in self.properties.items():
+                # Use f-string formatting for the property URN, like in dataset.py
+                prop_urn = f"urn:li:structuredProperty:{key}"
+                # Ensure value is a list for StructuredPropertyValueAssignmentClass
+                values_list = value if isinstance(value, list) else [value]
+                property_assignments.append(
+                    StructuredPropertyValueAssignmentClass(
+                        propertyUrn=prop_urn,
+                        values=[
+                            str(v) for v in values_list
+                        ],  # Ensure all values are strings
+                    )
+                )
+            if (
+                property_assignments
+            ):  # Only create aspect if there are valid assignments
+                structured_properties_aspect = StructuredPropertiesClass(
+                    properties=property_assignments
+                )
+
         # Now that we've generated the assertions, we can generate
         # the actual data contract.
         yield from MetadataChangeProposalWrapper.construct_many(
@@ -202,6 +232,8 @@ class DataContract(v1_ConfigModel):
                     if True
                     else None
                 ),
+                # Add structured properties aspect if defined
+                structured_properties_aspect,
             ],
         )
 

datahub/api/entities/datajob/dataflow.py

@@ -30,7 +30,7 @@ class DataFlow:
     """The DataHub representation of data-flow.
 
     Args:
-        urn (int): Unique identifier of the DataFlow in DataHub. For more detail refer https://datahubproject.io/docs/what/urn/.
+        urn (int): Unique identifier of the DataFlow in DataHub. For more detail refer https://docs.datahub.com/docs/what/urn/.
         id (str): Identifier of DataFlow in orchestrator.
         orchestrator (str): orchestrator. for example airflow.
         cluster (Optional[str]): [deprecated] Please use env.
@@ -40,8 +40,8 @@ class DataFlow:
         url (Optional[str]): URL pointing to DataFlow.
         tags (Set[str]): tags that need to be apply on DataFlow.
         owners (Set[str]): owners that need to be apply on DataFlow.
-        platform_instance (Optional[str]): The instance of the platform that all assets produced by this orchestrator belong to. For more detail refer https://datahubproject.io/docs/platform-instances/.
-        env (Optional[str]): The environment that all assets produced by this orchestrator belong to. For more detail and possible values refer https://datahubproject.io/docs/graphql/enums/#fabrictype.
+        platform_instance (Optional[str]): The instance of the platform that all assets produced by this orchestrator belong to. For more detail refer https://docs.datahub.com/docs/platform-instances/.
+        env (Optional[str]): The environment that all assets produced by this orchestrator belong to. For more detail and possible values refer https://docs.datahub.com/docs/graphql/enums/#fabrictype.
     """
 
     urn: DataFlowUrn = field(init=False)

datahub/api/entities/datajob/datajob.py

@@ -108,7 +108,9 @@ class DataJob:
         return [tags]
 
     def generate_mcp(
-        self, materialize_iolets: bool = True
+        self,
+        generate_lineage: bool = True,
+        materialize_iolets: bool = True,
     ) -> Iterable[MetadataChangeProposalWrapper]:
         env: Optional[str] = None
         if self.flow_urn.cluster.upper() in builder.ALL_ENV_TYPES:
@@ -152,9 +154,10 @@ class DataJob:
         )
         yield mcp
 
-        yield from self.generate_data_input_output_mcp(
-            materialize_iolets=materialize_iolets
-        )
+        if generate_lineage:
+            yield from self.generate_data_input_output_mcp(
+                materialize_iolets=materialize_iolets
+            )
 
         for owner in self.generate_ownership_aspect():
             mcp = MetadataChangeProposalWrapper(

datahub/api/entities/dataset/dataset.py

@@ -509,16 +509,14 @@ class Dataset(StrictModel):
     def generate_mcp(
         self,
     ) -> Iterable[Union[MetadataChangeProposalClass, MetadataChangeProposalWrapper]]:
-        mcp = MetadataChangeProposalWrapper(
-            entityUrn=self.urn,
-            aspect=DatasetPropertiesClass(
-                description=self.description,
-                name=self.name,
-                customProperties=self.properties,
-                externalUrl=self.external_url,
-            ),
-        )
-        yield mcp
+        patch_builder = self.patch_builder()
+
+        patch_builder.set_custom_properties(self.properties or {})
+        patch_builder.set_description(self.description)
+        patch_builder.set_display_name(self.name)
+        patch_builder.set_external_url(self.external_url)
+
+        yield from patch_builder.build()
 
         if self.schema_metadata:
             schema_fields = set()
@@ -981,7 +979,7 @@ class Dataset(StrictModel):
 
         def model_dump(self, **kwargs):
             """Custom model_dump method for Pydantic v2 to handle YAML serialization properly."""
-            exclude = kwargs.pop("exclude", set())
+            exclude = kwargs.pop("exclude", None) or set()
 
             # If id and name are identical, exclude name from the output
             if self.id == self.name and self.id is not None:

datahub/api/entities/forms/forms.py

@@ -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,
@@ -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/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/operation.py

@@ -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
@@ -103,9 +103,9 @@ 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.
         """
 

datahub/cli/check_cli.py

@@ -16,6 +16,7 @@ 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
@@ -259,7 +260,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,
@@ -472,7 +473,7 @@ WHERE
 @check.command()
 def server_config() -> None:
     """Print the server config."""
-    graph = get_default_graph()
+    graph = get_default_graph(ClientMode.CLI)
 
     server_config = graph.get_server_config()
 

datahub/cli/config_utils.py

@@ -17,8 +17,8 @@ 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"

datahub/cli/delete_cli.py

@@ -15,6 +15,7 @@ from datahub.cli import cli_utils
 from datahub.configuration.datetimes import ClickDatetime
 from datahub.emitter.aspect import ASPECT_MAP, TIMESERIES_ASPECT_MAP
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.ingestion.graph.filters import RemovedStatusFilter
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
@@ -48,7 +49,7 @@ def delete() -> None:
 
     See `datahub delete by-filter` for the list of available filters.
 
-    See https://datahubproject.io/docs/how/delete-metadata for more detailed docs.
+    See https://docs.datahub.com/docs/how/delete-metadata for more detailed docs.
     """
     pass
 
@@ -124,7 +125,7 @@ def by_registry(
     Delete all metadata written using the given registry id and version pair.
     """
 
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
 
     if soft and not dry_run:
         raise click.UsageError(
@@ -175,7 +176,7 @@ def references(urn: str, dry_run: bool, force: bool) -> None:
     Delete all references to an entity (but not the entity itself).
     """
 
-    graph = get_default_graph()
+    graph = get_default_graph(ClientMode.CLI)
     logger.info(f"Using graph: {graph}")
 
     references_count, related_aspects = graph.delete_references_to_urn(
@@ -238,7 +239,7 @@ def undo_by_filter(
     """
     Undo soft deletion by filters
     """
-    graph = get_default_graph()
+    graph = get_default_graph(ClientMode.CLI)
     logger.info(f"Using {graph}")
     if urn:
         graph.set_soft_delete_status(urn=urn, delete=False)
@@ -410,7 +411,7 @@ def by_filter(
                 abort=True,
             )
 
-    graph = get_default_graph()
+    graph = get_default_graph(ClientMode.CLI)
     logger.info(f"Using {graph}")
 
     # Determine which urns to delete.

datahub/cli/docker_cli.py

@@ -811,7 +811,7 @@ def quickstart(
         raise status.to_exception(
             header="Unable to run quickstart - the following issues were detected:",
             footer="If you think something went wrong, please file an issue at https://github.com/datahub-project/datahub/issues\n"
-            "or send a message in our Slack https://slack.datahubproject.io/\n"
+            "or send a message in our Slack https://datahub.com/slack/\n"
             f"Be sure to attach the logs from {log_file.name}",
         )
 
@@ -824,7 +824,7 @@ def quickstart(
         fg="green",
     )
     click.secho(
-        "Need support? Get in touch on Slack: https://slack.datahubproject.io/",
+        "Need support? Get in touch on Slack: https://datahub.com/slack/",
         fg="magenta",
     )
 

datahub/cli/exists_cli.py

@@ -6,6 +6,7 @@ import click
 from click_default_group import DefaultGroup
 
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 
@@ -36,4 +37,4 @@ def urn(ctx: Any, urn: Optional[str]) -> None:
             raise click.UsageError("Nothing for me to get. Maybe provide an urn?")
         urn = ctx.args[0]
         logger.debug(f"Using urn from args {urn}")
-    click.echo(json.dumps(get_default_graph().exists(urn)))
+    click.echo(json.dumps(get_default_graph(ClientMode.CLI).exists(urn)))

datahub/cli/get_cli.py

@@ -7,6 +7,7 @@ from click_default_group import DefaultGroup
 
 from datahub.cli.cli_utils import get_aspects_for_entity
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 
@@ -46,7 +47,7 @@ def urn(ctx: Any, urn: Optional[str], aspect: List[str], details: bool) -> None:
         urn = ctx.args[0]
         logger.debug(f"Using urn from args {urn}")
 
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
 
     if aspect:
         # If aspects are specified, we need to do the existence check first.

datahub/cli/iceberg_cli.py

@@ -13,6 +13,7 @@ import datahub.metadata.schema_classes
 from datahub.cli.cli_utils import post_entity
 from datahub.configuration.common import GraphError
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.schema_classes import SystemMetadataClass
 from datahub.telemetry import telemetry
 
@@ -178,7 +179,7 @@ def create(
     Create an iceberg warehouse.
     """
 
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
 
     urn = iceberg_data_platform_instance_urn(warehouse)
 
@@ -331,7 +332,7 @@ def update(
     Update iceberg warehouses. Can only update credentials, and role. Cannot update region
     """
 
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
 
     urn = iceberg_data_platform_instance_urn(warehouse)
 
@@ -407,7 +408,7 @@ def list() -> None:
     List iceberg warehouses
     """
 
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
 
     for warehouse in get_all_warehouses(client):
         click.echo(warehouse)
@@ -420,7 +421,7 @@ def list() -> None:
 @telemetry.with_telemetry()
 def get(warehouse: str) -> None:
     """Fetches the details of the specified iceberg warehouse"""
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
     urn = iceberg_data_platform_instance_urn(warehouse)
 
     if client.exists(urn):
@@ -455,7 +456,7 @@ def delete(warehouse: str, dry_run: bool, force: bool) -> None:
 
     urn = iceberg_data_platform_instance_urn(warehouse)
 
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
 
     if not client.exists(urn):
         raise click.ClickException(f"urn {urn} not found")

datahub/cli/ingest_cli.py

@@ -14,10 +14,11 @@ 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.cli.config_utils import CONDENSED_DATAHUB_CONFIG_PATH, load_client_config
 from datahub.configuration.common import GraphError
 from datahub.configuration.config_loader import load_config_file
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.ingestion.run.connection import ConnectionManager
 from datahub.ingestion.run.pipeline import Pipeline
 from datahub.telemetry import telemetry
@@ -269,7 +270,7 @@ def deploy(
     urn:li:dataHubIngestionSource:<name>
     """
 
-    datahub_graph = get_default_graph()
+    datahub_graph = get_default_graph(ClientMode.CLI)
 
     variables = deploy_source_vars(
         name=name,
@@ -360,6 +361,7 @@ def mcps(path: str) -> None:
     """
 
     click.echo("Starting ingestion...")
+    datahub_config = load_client_config()
     recipe: dict = {
         "source": {
             "type": "file",
@@ -367,6 +369,7 @@ def mcps(path: str) -> None:
                 "path": path,
             },
         },
+        "datahub_api": datahub_config,
     }
 
     pipeline = Pipeline.create(recipe, report_to=None)
@@ -422,7 +425,7 @@ def list_source_runs(page_offset: int, page_size: int, urn: str, source: str) ->
         }
     }
 
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
     session = client._session
     gms_host = client.config.server
 
@@ -508,7 +511,7 @@ def list_source_runs(page_offset: int, page_size: int, urn: str, source: str) ->
 def list_runs(page_offset: int, page_size: int, include_soft_deletes: bool) -> None:
     """List recent ingestion runs to datahub"""
 
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
     session = client._session
     gms_host = client.config.server
 
@@ -559,7 +562,7 @@ def show(
     run_id: str, start: int, count: int, include_soft_deletes: bool, show_aspect: bool
 ) -> None:
     """Describe a provided ingestion run to datahub"""
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
     session = client._session
     gms_host = client.config.server
 
@@ -609,7 +612,7 @@ def rollback(
     run_id: str, force: bool, dry_run: bool, safe: bool, report_dir: str
 ) -> None:
     """Rollback a provided ingestion run to datahub"""
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
 
     if not force and not dry_run:
         click.confirm(

datahub/cli/migrate.py

@@ -25,6 +25,7 @@ from datahub.emitter.mcp_builder import (
 )
 from datahub.emitter.rest_emitter import DatahubRestEmitter
 from datahub.ingestion.graph.client import (
+    ClientMode,
     DataHubGraph,
     RelatedEntity,
     get_default_graph,
@@ -147,7 +148,7 @@ def dataplatform2instance_func(
     migration_report = MigrationReport(run_id, dry_run, keep)
     system_metadata = SystemMetadataClass(runId=run_id)
 
-    graph = get_default_graph()
+    graph = get_default_graph(ClientMode.CLI)
 
     urns_to_migrate: List[str] = []
 
@@ -386,7 +387,7 @@ def migrate_containers(
 
 
 def get_containers_for_migration(env: str) -> List[Any]:
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
     containers_to_migrate = list(
         client.get_urns_by_filter(entity_types=["container"], env=env)
     )
@@ -445,7 +446,7 @@ def process_container_relationships(
     relationships: Iterable[RelatedEntity] = migration_utils.get_incoming_relationships(
         urn=src_urn
     )
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
     for relationship in relationships:
         log.debug(f"Incoming Relationship: {relationship}")
         target_urn: str = relationship.urn

datahub/cli/migration_utils.py

@@ -12,6 +12,7 @@ from datahub.ingestion.graph.client import (
     RelatedEntity,
     get_default_graph,
 )
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.schema_classes import (
     ChartInfoClass,
     ContainerClass,
@@ -243,7 +244,7 @@ def clone_aspect(
     run_id: str = str(uuid.uuid4()),
     dry_run: bool = False,
 ) -> Iterable[MetadataChangeProposalWrapper]:
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
     aspect_map = cli_utils.get_aspects_for_entity(
         client._session,
         client.config.server,
@@ -274,7 +275,7 @@ def clone_aspect(
 
 
 def get_incoming_relationships(urn: str) -> Iterable[RelatedEntity]:
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
     yield from client.get_related_entities(
         entity_urn=urn,
         relationship_types=[
@@ -290,7 +291,7 @@ def get_incoming_relationships(urn: str) -> Iterable[RelatedEntity]:
 
 
 def get_outgoing_relationships(urn: str) -> Iterable[RelatedEntity]:
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
     yield from client.get_related_entities(
         entity_urn=urn,
         relationship_types=[