__init__.py.tmpl

import databricks.sdk.core as client
import databricks.sdk.dbutils as dbutils
from databricks.sdk.credentials_provider import CredentialsStrategy

from databricks.sdk.mixins.files import DbfsExt
from databricks.sdk.mixins.compute import ClustersExt
from databricks.sdk.mixins.workspace import WorkspaceExt
from databricks.sdk.mixins.open_ai_client import ServingEndpointsExt
{{- range .Services}}
from databricks.sdk.service.{{.Package.Name}} import {{.PascalName}}API{{end}}
from databricks.sdk.service.provisioning import Workspace
from databricks.sdk import azure
from typing import Optional

{{$args := list "host" "account_id" "username" "password" "client_id" "client_secret"
  "token" "profile" "config_file" "azure_workspace_resource_id" "azure_client_secret"
  "azure_client_id" "azure_tenant_id" "azure_environment" "auth_type" "cluster_id"
  "google_credentials" "google_service_account" }}

{{- define "api" -}}
  {{- $mixins := dict "ClustersAPI" "ClustersExt" "DbfsAPI" "DbfsExt" "WorkspaceAPI" "WorkspaceExt" "ServingEndpointsAPI" "ServingEndpointsExt" -}}
  {{- $genApi := concat .PascalName "API" -}}
  {{- 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:
    """
    The WorkspaceClient is a client for the workspace-level Databricks REST API.
    """
    def __init__(self, *{{range $args}}, {{.}}: Optional[str] = None{{end}},
                 debug_truncate_bytes: Optional[int] = None,
                 debug_headers: Optional[bool] = None,
                 product="unknown",
                 product_version="0.0.0",
                 credentials_strategy: Optional[CredentialsStrategy] = None,
                 credentials_provider: Optional[CredentialsStrategy] = None,
                 config: Optional[client.Config] = None):
        if not config:
          config = client.Config({{range $args}}{{.}}={{.}}, {{end}}
            credentials_strategy=credentials_strategy,
            credentials_provider=credentials_provider,
            debug_truncate_bytes=debug_truncate_bytes,
            debug_headers=debug_headers,
            product=product,
            product_version=product_version)
        self._config = config.copy()
        self._dbutils = _make_dbutils(self._config)
        self._api_client = client.ApiClient(self._config)

        {{- range .Services}}{{if and (not .IsAccounts) (not .HasParent) .HasDataPlaneAPI (not .IsDataPlane)}}
        {{.SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}

        {{- range .Services}}
        {{- if and (not .IsAccounts) (not .HasParent)}}
        {{- if .IsDataPlane}}
        self._{{.SnakeName}} = {{template "api" .}}(self._api_client, {{.ControlPlaneService.SnakeName}})
        {{- else if .HasDataPlaneAPI}}
        self._{{.SnakeName}} = {{.SnakeName}}
        {{- else}}
        self._{{.SnakeName}} = {{template "api" .}}(self._api_client)
        {{- end -}}
        {{- end -}}
        {{end}}

    @property
    def config(self) -> client.Config:
        return self._config

    @property
    def api_client(self) -> client.ApiClient:
        return self._api_client

    @property
    def dbutils(self) -> dbutils.RemoteDbUtils:
        return self._dbutils

    {{- range .Services}}{{if and (not .IsAccounts) (not .HasParent)}}
    @property
    def {{.SnakeName}}(self) -> {{template "api" .}}:
        {{if .Description}}"""{{.Summary}}"""{{end}}
        return self._{{.SnakeName}}
    {{end -}}{{end}}

    def get_workspace_id(self) -> int:
        """Get the workspace ID of the workspace that this client is connected to."""
        response = self._api_client.do("GET",
                                       "/api/2.0/preview/scim/v2/Me",
                                       response_headers=['X-Databricks-Org-Id'])
        return int(response["X-Databricks-Org-Id"])

    def __repr__(self):
        return f"WorkspaceClient(host='{self._config.host}', auth_type='{self._config.auth_type}', ...)"

class AccountClient:
    """
    The AccountClient is a client for the account-level Databricks REST API.
    """

    def __init__(self, *{{range $args}}, {{.}}: Optional[str] = None{{end}},
                 debug_truncate_bytes: Optional[int] = None,
                 debug_headers: Optional[bool] = None,
                 product="unknown",
                 product_version="0.0.0",
                 credentials_strategy: Optional[CredentialsStrategy] = None,
                 credentials_provider: Optional[CredentialsStrategy] = None,
                 config: Optional[client.Config] = None):
        if not config:
          config = client.Config({{range $args}}{{.}}={{.}}, {{end}}
            credentials_strategy=credentials_strategy,
            credentials_provider=credentials_provider,
            debug_truncate_bytes=debug_truncate_bytes,
            debug_headers=debug_headers,
            product=product,
            product_version=product_version)
        self._config = config.copy()
        self._api_client = client.ApiClient(self._config)

        {{- range .Services}}{{if and .IsAccounts (not .HasParent) .HasDataPlaneAPI (not .IsDataPlane)}}
        {{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client){{end -}}{{end}}

        {{- range .Services}}
        {{- if and .IsAccounts (not .HasParent)}}
        {{- if .IsDataPlane}}
        self._{{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client, {{.ControlPlaneService.SnakeName}})
        {{- else if .HasDataPlaneAPI}}
        self._{{(.TrimPrefix "account").SnakeName}} = {{(.TrimPrefix "account").SnakeName}}
        {{- else}}
        self._{{(.TrimPrefix "account").SnakeName}} = {{template "api" .}}(self._api_client)
        {{- end -}}
        {{- end -}}
        {{end}}

    @property
    def config(self) -> client.Config:
        return self._config

    @property
    def api_client(self) -> client.ApiClient:
        return self._api_client

    {{- range .Services}}{{if and .IsAccounts (not .HasParent)}}
    @property
    def {{(.TrimPrefix "account").SnakeName}}(self) -> {{template "api" .}}:{{if .Description}}
        """{{.Summary}}"""{{end}}
        return self._{{(.TrimPrefix "account").SnakeName}}
    {{end -}}{{end}}

    def get_workspace_client(self, workspace: Workspace) -> WorkspaceClient:
        """Constructs a ``WorkspaceClient`` for the given workspace.

        Returns a ``WorkspaceClient`` that is configured to use the same
        credentials as this ``AccountClient``. The underlying config is
        copied from this ``AccountClient``, but the ``host`` and
        ``azure_workspace_resource_id`` are overridden to match the
        given workspace, and the ``account_id`` field is cleared.

        Usage:

        .. code-block::

            wss = list(a.workspaces.list())
            if len(wss) == 0:
                pytest.skip("no workspaces")
            w = a.get_workspace_client(wss[0])
            assert w.current_user.me().active

        :param workspace: The workspace to construct a client for.
        :return: A ``WorkspaceClient`` for the given workspace.
        """
        config = self._config.deep_copy()
        config.host = config.environment.deployment_url(workspace.deployment_name)
        config.azure_workspace_resource_id = azure.get_azure_resource_id(workspace)
        config.account_id = None
        config.init_auth()
        return WorkspaceClient(config=config)

    def __repr__(self):
        return f"AccountClient(account_id='{self._config.account_id}', auth_type='{self._config.auth_type}', ...)"