diff --git a/.codegen/__init__.py.tmpl b/.codegen/__init__.py.tmpl index c52b8ed65..a8e730af0 100644 --- a/.codegen/__init__.py.tmpl +++ b/.codegen/__init__.py.tmpl @@ -17,6 +17,22 @@ from databricks.sdk.service.{{.Package.Name}} import {{.PascalName}}API{{end}} {{- getOrDefault $mixins $genApi $genApi -}} {{- end -}} +def _make_dbutils(config: client.Config): + # We try to directly check if we are in runtime, instead of + # trying to import from databricks.sdk.runtime. This is to prevent + # remote dbutils from being created without the config, which is both + # expensive (will need to check all credential providers) and can + # throw errors (when no env vars are set). + try: + from dbruntime import UserNamespaceInitializer + except ImportError: + return dbutils.RemoteDbUtils(config) + + # We are in runtime, so we can use the runtime dbutils + from databricks.sdk.runtime import dbutils as runtime_dbutils + return runtime_dbutils + + class WorkspaceClient: def __init__(self, *{{range $args}}, {{.}}: str = None{{end}}, debug_truncate_bytes: int = None, @@ -33,7 +49,7 @@ class WorkspaceClient: product=product, product_version=product_version) self.config = config.copy() - self.dbutils = dbutils.RemoteDbUtils(self.config) + self.dbutils = _make_dbutils(self.config) self.api_client = client.ApiClient(self.config) self.files = FilesMixin(self.api_client) {{- range .Services}}{{if not .IsAccounts}} diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index f364ad2ff..c1c3d189c 100755 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -59,6 +59,22 @@ SecretsAPI, WorkspaceAPI) +def _make_dbutils(config: client.Config): + # We try to directly check if we are in runtime, instead of + # trying to import from databricks.sdk.runtime. This is to prevent + # remote dbutils from being created without the config, which is both + # expensive (will need to check all credential providers) and can + # throw errors (when no env vars are set). + try: + from dbruntime import UserNamespaceInitializer + except ImportError: + return dbutils.RemoteDbUtils(config) + + # We are in runtime, so we can use the runtime dbutils + from databricks.sdk.runtime import dbutils as runtime_dbutils + return runtime_dbutils + + class WorkspaceClient: def __init__(self, @@ -108,7 +124,7 @@ def __init__(self, product=product, product_version=product_version) self.config = config.copy() - self.dbutils = dbutils.RemoteDbUtils(self.config) + self.dbutils = _make_dbutils(self.config) self.api_client = client.ApiClient(self.config) self.files = FilesMixin(self.api_client) self.account_access_control_proxy = AccountAccessControlProxyAPI(self.api_client)