acryl-datahub 1.0.0.3rc9__py3-none-any.whl → 1.0.0.3rc11__py3-none-any.whl

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=[

datahub/cli/put_cli.py

@@ -8,6 +8,7 @@ from datahub.cli.cli_utils import post_entity
 from datahub.configuration.config_loader import load_config_file
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.schema_classes import (
     DataPlatformInfoClass as DataPlatformInfo,
     PlatformTypeClass,
@@ -53,7 +54,7 @@ def aspect(urn: str, aspect: str, aspect_data: str, run_id: Optional[str]) -> No
         aspect_data, allow_stdin=True, resolve_env_vars=False, process_directives=False
     )
 
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
 
     system_metadata: Union[None, SystemMetadataClass] = None
     if run_id:
@@ -118,7 +119,7 @@ def platform(
         displayName=display_name or platform_name,
         logoUrl=logo,
     )
-    datahub_graph = get_default_graph()
+    datahub_graph = get_default_graph(ClientMode.CLI)
     mcp = MetadataChangeProposalWrapper(
         entityUrn=str(platform_urn),
         aspect=data_platform_info,

datahub/cli/specific/assertions_cli.py

@@ -15,6 +15,7 @@ from datahub.api.entities.assertion.compiler_interface import (
 from datahub.emitter.mce_builder import make_assertion_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.integrations.assertion.registry import ASSERTION_PLATFORMS
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
@@ -39,7 +40,7 @@ def upsert(file: str) -> None:
 
     assertions_spec: AssertionsConfigSpec = AssertionsConfigSpec.from_yaml(file)
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         for assertion_spec in assertions_spec.assertions:
             try:
                 mcp = MetadataChangeProposalWrapper(

datahub/cli/specific/datacontract_cli.py

@@ -6,6 +6,7 @@ from click_default_group import DefaultGroup
 
 from datahub.api.entities.datacontract.datacontract import DataContract
 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
 
@@ -28,7 +29,7 @@ def upsert(file: str) -> None:
     data_contract: DataContract = DataContract.from_yaml(file)
     urn = data_contract.urn
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if not graph.exists(data_contract.entity):
             raise ValueError(
                 f"Cannot define a data contract for non-existent entity {data_contract.entity}"
@@ -72,7 +73,7 @@ def delete(urn: Optional[str], file: Optional[str], hard: bool) -> None:
         data_contract = DataContract.from_yaml(file)
         urn = data_contract.urn
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if not graph.exists(urn):
             raise ValueError(f"Data Contract {urn} does not exist")
 

datahub/cli/specific/dataproduct_cli.py

@@ -20,6 +20,7 @@ from datahub.emitter.mce_builder import (
     validate_ownership_type,
 )
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.schema_classes import OwnerClass, OwnershipTypeClass
 from datahub.specific.dataproduct import DataProductPatchBuilder
 from datahub.telemetry import telemetry
@@ -81,7 +82,7 @@ def mutate(file: Path, validate_assets: bool, external_url: str, upsert: bool) -
 
     config_dict = load_file(pathlib.Path(file))
     id = config_dict.get("id") if isinstance(config_dict, dict) else None
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         data_product: DataProduct = DataProduct.from_yaml(file, graph)
         external_url_override = (
             external_url
@@ -162,7 +163,7 @@ def upsert(file: Path, validate_assets: bool, external_url: str) -> None:
 def diff(file: Path, update: bool) -> None:
     """Diff a Data Product file with its twin in DataHub"""
 
-    with get_default_graph() as emitter:
+    with get_default_graph(ClientMode.CLI) as emitter:
         id: Optional[str] = None
         try:
             data_product_local: DataProduct = DataProduct.from_yaml(file, emitter)
@@ -216,7 +217,7 @@ def delete(urn: str, file: Path, hard: bool) -> None:
         raise click.Abort()
 
     graph: DataHubGraph
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         data_product_urn = (
             urn if urn.startswith("urn:li:dataProduct") else f"urn:li:dataProduct:{urn}"
         )
@@ -248,7 +249,7 @@ def get(urn: str, to_file: str) -> None:
     if not urn.startswith("urn:li:dataProduct:"):
         urn = f"urn:li:dataProduct:{urn}"
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if graph.exists(urn):
             dataproduct: DataProduct = DataProduct.from_datahub(graph=graph, id=urn)
             click.secho(
@@ -306,7 +307,7 @@ def set_description(urn: str, description: str, md_file: Path) -> None:
 
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.set_description(description)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "set description")
         for mcp in dataproduct_patcher.build():
             graph.emit(mcp)
@@ -342,7 +343,7 @@ def add_owner(urn: str, owner: str, owner_type: str) -> None:
             owner=_get_owner_urn(owner), type=owner_type, typeUrn=owner_type_urn
         )
     )
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "add owners")
         for mcp in dataproduct_patcher.build():
             graph.emit(mcp)
@@ -360,7 +361,7 @@ def remove_owner(urn: str, owner_urn: str) -> None:
         urn = f"urn:li:dataProduct:{urn}"
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.remove_owner(owner=_get_owner_urn(owner_urn))
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "remove owners")
         for mcp in dataproduct_patcher.build():
             click.echo(json.dumps(mcp.to_obj()))
@@ -382,7 +383,7 @@ def add_asset(urn: str, asset: str, validate_assets: bool) -> None:
         urn = f"urn:li:dataProduct:{urn}"
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.add_asset(asset)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "add assets")
         if validate_assets:
             _abort_if_non_existent_urn(
@@ -409,7 +410,7 @@ def remove_asset(urn: str, asset: str, validate_assets: bool) -> None:
         urn = f"urn:li:dataProduct:{urn}"
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.remove_asset(asset)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "remove assets")
         if validate_assets:
             _abort_if_non_existent_urn(

datahub/cli/specific/dataset_cli.py

@@ -12,6 +12,7 @@ from click_default_group import DefaultGroup
 from datahub.api.entities.dataset.dataset import Dataset, DatasetRetrievalConfig
 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.com.linkedin.pegasus2avro.common import Siblings
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
@@ -54,7 +55,7 @@ def get(urn: str, to_file: str) -> None:
     if not urn.startswith("urn:li:dataset:"):
         urn = f"urn:li:dataset:{urn}"
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if graph.exists(urn):
             dataset: Dataset = Dataset.from_datahub(graph=graph, urn=urn)
             click.secho(
@@ -82,7 +83,7 @@ def add_sibling(urn: str, sibling_urns: Tuple[str]) -> None:
     all_urns.add(urn)
     for sibling_urn in sibling_urns:
         all_urns.add(sibling_urn)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         for _urn in all_urns:
             _emit_sibling(graph, urn, _urn, all_urns)
 
@@ -181,7 +182,7 @@ def sync(file: str, to_datahub: bool, dry_run: bool) -> None:
     dry_run_prefix = "[dry-run]: " if dry_run else ""  # prefix to use in messages
 
     failures: List[str] = []
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         datasets = Dataset.from_yaml(file)
         for dataset in datasets:
             assert (