From 037850736db98f67ada2bcbd7de22c399d89fafd Mon Sep 17 00:00:00 2001 From: Parth Bansal Date: Wed, 30 Oct 2024 22:46:10 +0100 Subject: [PATCH] update --- .codegen.json | 17 +- .codegen/__init__.py.tmpl | 194 -------------- .codegen/error_mapping.py.tmpl | 20 -- .codegen/error_overrides.py.tmpl | 20 -- .codegen/example.py.tmpl | 112 -------- .codegen/lib.tmpl | 12 - .codegen/service.py.tmpl | 423 ------------------------------- 7 files changed, 2 insertions(+), 796 deletions(-) delete mode 100644 .codegen/__init__.py.tmpl delete mode 100644 .codegen/error_mapping.py.tmpl delete mode 100644 .codegen/error_overrides.py.tmpl delete mode 100644 .codegen/example.py.tmpl delete mode 100644 .codegen/lib.tmpl delete mode 100644 .codegen/service.py.tmpl diff --git a/.codegen.json b/.codegen.json index a1886bd80..3a880d1a9 100644 --- a/.codegen.json +++ b/.codegen.json @@ -1,20 +1,6 @@ { - "formatter": "yapf -pri $FILENAMES && autoflake -i $FILENAMES && isort $FILENAMES", + "mode": "py_v0", "changelog_config": ".codegen/changelog_config.yml", - "template_libraries": [ - ".codegen/lib.tmpl" - ], - "packages": { - ".codegen/service.py.tmpl": "databricks/sdk/service/{{.Name}}.py" - }, - "batch": { - ".codegen/__init__.py.tmpl": "databricks/sdk/__init__.py", - ".codegen/error_mapping.py.tmpl": "databricks/sdk/errors/platform.py", - ".codegen/error_overrides.py.tmpl": "databricks/sdk/errors/overrides.py" - }, - "samples": { - ".codegen/example.py.tmpl": "examples/{{if .IsAccount}}account{{else}}workspace{{end}}/{{.Service.SnakeName}}/{{.Method.SnakeName}}_{{.SnakeName}}.py" - }, "version": { "databricks/sdk/version.py": "__version__ = '$VERSION'" }, @@ -28,6 +14,7 @@ "pip install '.[dev]'" ], "post_generate": [ + "make fmt", "pytest -m 'not integration' --cov=databricks --cov-report html tests", "pip install .", "python docs/gen-client-docs.py" diff --git a/.codegen/__init__.py.tmpl b/.codegen/__init__.py.tmpl deleted file mode 100644 index d54e9dfff..000000000 --- a/.codegen/__init__.py.tmpl +++ /dev/null @@ -1,194 +0,0 @@ -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}', ...)" diff --git a/.codegen/error_mapping.py.tmpl b/.codegen/error_mapping.py.tmpl deleted file mode 100644 index b3cc8cea6..000000000 --- a/.codegen/error_mapping.py.tmpl +++ /dev/null @@ -1,20 +0,0 @@ -# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. - -from .base import DatabricksError - -{{range .ExceptionTypes}} -class {{.PascalName}}({{if .Inherit -}} - {{.Inherit.PascalName}} - {{- else -}} - DatabricksError - {{- end -}}): - """{{.Comment " " 100 | trimSuffix "\"" }}""" -{{end}} - -STATUS_CODE_MAPPING = { {{range .ErrorStatusCodeMapping}} - {{.StatusCode}}: {{.PascalName}},{{- end}} -} - -ERROR_CODE_MAPPING = { {{range .ErrorCodeMapping}} - '{{.ErrorCode}}': {{.PascalName}},{{- end}} -} diff --git a/.codegen/error_overrides.py.tmpl b/.codegen/error_overrides.py.tmpl deleted file mode 100644 index adcfea555..000000000 --- a/.codegen/error_overrides.py.tmpl +++ /dev/null @@ -1,20 +0,0 @@ -# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. - -from .base import _ErrorOverride -from .platform import * -import re - - -_ALL_OVERRIDES = [ - {{ range .ErrorOverrides -}} - _ErrorOverride( - debug_name="{{.Name}}", - path_regex=re.compile(r'{{.PathRegex}}'), - verb="{{.Verb}}", - status_code_matcher=re.compile(r'{{replaceAll "'" "\\'" .StatusCodeMatcher}}'), - error_code_matcher=re.compile(r'{{replaceAll "'" "\\'" .ErrorCodeMatcher}}'), - message_matcher=re.compile(r'{{replaceAll "'" "\\'" .MessageMatcher}}'), - custom_error={{.OverrideErrorCode.PascalName}}, - ), -{{- end }} -] diff --git a/.codegen/example.py.tmpl b/.codegen/example.py.tmpl deleted file mode 100644 index dba71d9bf..000000000 --- a/.codegen/example.py.tmpl +++ /dev/null @@ -1,112 +0,0 @@ -from databricks.sdk import {{if .IsAccount}}AccountClient{{else}}WorkspaceClient{{end}} -from databricks.sdk.service import _internal{{range .Suite.ServiceToPackage}}, {{.}}{{end}} -import time, base64, os - -{{$example := .}} -{{if .IsAccount}}a{{else}}w{{end}} = {{if .IsAccount}}Account{{else}}Workspace{{end}}Client() - -{{range .Init}} -{{.SnakeName}} = {{template "expr" .Value}} -{{end}} - -{{range .Calls}} -{{if .Service -}} - {{template "svc-call" .}} -{{- else -}} - {{with .Assign}}{{.SnakeName}} = {{end}}{{template "expr" .}} -{{- end}} -{{end}} - -{{with .Cleanup}} -# cleanup -{{range . -}} - {{template "svc-call" .}} -{{end}} -{{end}} - -{{define "svc-call" -}} - {{with .Assign}}{{.SnakeName}} = {{end}}{{if .IsAccount}}a{{else}}w{{end}}.{{.Service.SnakeName}}.{{.Original.SnakeName}}{{if eq .Original.SnakeName "import"}}_{{end}}({{template "method-args" .}}) - {{- if .IsWait}}.result(){{end}} -{{- end}} - -{{define "expr" -}} -{{- if eq .Type "binary" -}} - {{template "expr" .Left}} {{.Op}} {{template "expr" .Right}} -{{- else if eq .Type "index" -}} - {{template "expr" .Left}}[{{template "expr" .Right}}] -{{- else if eq .Type "boolean" -}} - {{if .Value}}True{{else}}False{{end}} -{{- else if eq .Type "heredoc" -}} -"""{{.Value}}""" -{{- else if eq .Type "literal" -}} - {{.Value}} -{{- else if eq .Type "lookup" -}} - {{template "expr" .X}}.{{.Field.SnakeName}} -{{- else if eq .Type "enum" -}} - {{.Package}}.{{.Entity.PascalName}}.{{.ConstantName}} -{{- else if eq .Type "variable" -}} - {{if eq .SnakeName "true"}}True - {{- else if eq .SnakeName "false"}}False - {{else}}{{.SnakeName}}{{end}} -{{- else if eq .Type "entity" -}} - {{.Package}}.{{.PascalName}}({{template "kwargs" .FieldValues}}) -{{- else if eq .Type "call" -}} - {{template "call" .}} -{{- else if eq .Type "map" -}} - { {{range .Pairs}}{{template "expr" .Key}}: {{template "expr" .Value}},{{end}} } -{{- else if eq .Type "array" -}} - [ {{range $i, $x := .Values}}{{if $i}}, {{end}}{{template "expr" .}}{{end}} ] -{{- else -}} - /* UNKNOWN: {{.Type}} */ -{{- end -}} -{{- end}} - -{{define "method-args" -}} - {{with .Request -}} - {{template "kwargs" .}} - {{- else -}} - {{template "args" .}} - {{- end}} -{{- end}} - -{{define "kwargs" -}} - {{range $i, $x := . -}} - {{if $i}}, {{end}}{{.SnakeName}}={{template "expr" .Value}} - {{- end}} -{{- end}} - -{{define "args" -}} - {{range $i, $x := .Args -}} - {{if $i}}, {{end}}{{template "expr" .}} - {{- end}} -{{- end}} - -{{define "call" -}} -{{- if eq .PascalName "GetEnvOrSkipTest" -}} -os.environ[{{template "args" .}}] -{{- else if eq .PascalName "Dir" -}} -os.path.dirname({{template "args" .}}) -{{- else if eq .PascalName "Sprintf" -}} -{{range $i, $x := .Args}}{{if eq $i 0}}{{template "expr" .}} % ({{else}} {{if gt $i 1}}, {{end}} {{template "expr" .}} {{end}}{{end}}) -{{- else if eq .PascalName "MustParseInt64" -}} -{{template "args" .}} -{{- else if eq .PascalName "RandomEmail" -}} -f'sdk-{time.time_ns()}@example.com' -{{- else if eq .PascalName "RandomName" -}} -f'sdk-{time.time_ns()}' -{{- else if eq .PascalName "RandomHex" -}} -hex(time.time_ns())[2:] -{{- else if eq .PascalName "EncodeToString" -}} -base64.b64encode({{template "args" .}}.encode()).decode() -{{- else if eq .PascalName "CanonicalHostName" -}} -w.config.host -{{- else if eq .PascalName "SharedRunningCluster" -}} -w.clusters.ensure_cluster_is_running(os.environ["DATABRICKS_CLUSTER_ID"]) and os.environ["DATABRICKS_CLUSTER_ID"] -{{- else if eq .PascalName "DltNotebook" -}} -"CREATE LIVE TABLE dlt_sample AS SELECT 1" -{{- else if eq .PascalName "MyNotebookPath" -}} -f'/Users/{w.current_user.me().user_name}/sdk-{time.time_ns()}' -{{- else -}} -{{.SnakeName}}({{range $i, $x := .Args}}{{if $i}}, {{end}}{{template "expr" .}}{{end}}) -{{- end -}} -{{- end}} diff --git a/.codegen/lib.tmpl b/.codegen/lib.tmpl deleted file mode 100644 index 50233ca08..000000000 --- a/.codegen/lib.tmpl +++ /dev/null @@ -1,12 +0,0 @@ -{{ define "safe-name" -}} - {{/* https://docs.python.org/3/reference/lexical_analysis.html#keywords */}} - {{- $keywords := list "False" "await" "else" "import" "pass" "None" "break" "except" "in" "raise" - "True" "class" "finally" "is" "return" "and" "continue" "for" "lambda" "try" - "as" "def" "from" "nonlocal" "while" "assert" "del" "global" "not" "with" - "async" "elif" "if" "or" "yield" -}} - {{.}}{{ if in $keywords . }}_{{ end }} -{{- end}} - -{{ define "safe-snake-name" -}} - {{ template "safe-name" .SnakeName }} -{{- end}} diff --git a/.codegen/service.py.tmpl b/.codegen/service.py.tmpl deleted file mode 100644 index a062e4d10..000000000 --- a/.codegen/service.py.tmpl +++ /dev/null @@ -1,423 +0,0 @@ -# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. - -from __future__ import annotations -from dataclasses import dataclass -from datetime import timedelta -from enum import Enum -from typing import Dict, List, Any, Iterator, Type, Callable, Optional, BinaryIO -import time -import random -import logging -import requests - -from ..data_plane import DataPlaneService -from ..errors import OperationTimeout, OperationFailed -from ._internal import _enum, _from_dict, _repeated_dict, _repeated_enum, Wait, _escape_multi_segment_path_parameter -from ..oauth import Token - -_LOG = logging.getLogger('databricks.sdk') - -{{range .ImportedEntities}} -from databricks.sdk.service import {{.Package.Name}}{{end}} - -# all definitions in this file are in alphabetical order -{{range .Types}} -{{if or .Fields .IsEmpty -}}{{if not .IsRequest}}@dataclass -class {{.PascalName}}{{if eq "List" .PascalName}}Request{{end}}:{{if .Description}} - """{{.Comment " " 100}}""" - {{end}} - {{- range .RequiredFields}} - {{template "safe-snake-name" .}}: {{template "type" .Entity}}{{if .Description}} - """{{.Comment " " 100 | trimSuffix "\""}}"""{{end}} - {{end}} - {{- range .NonRequiredFields}} - {{template "safe-snake-name" .}}: Optional[{{template "type" .Entity}}] = None{{if .Description}} - """{{.Comment " " 100 | trimSuffix "\""}}"""{{end}} - {{end}} - {{if or .IsEmpty .HasJsonField .HasHeaderField .HasByteStreamField -}} - def as_dict(self) -> dict: - """Serializes the {{.PascalName}}{{if eq "List" .PascalName}}Request{{end}} into a dictionary suitable for use as a JSON request body.""" - body = {} - {{range .Fields}}if self.{{template "safe-snake-name" .}}{{with .Entity.IsPrimitive}} is not None{{end}}: body['{{.Name}}'] = {{template "as_request_type" .}} - {{end -}} - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> {{.PascalName}}{{if eq "List" .PascalName}}Request{{end}}: - """Deserializes the {{.PascalName}}{{if eq "List" .PascalName}}Request{{end}} from a dictionary.""" - return cls({{range $i, $f := .Fields}}{{if $i}}, {{end}}{{template "safe-snake-name" $f}}={{template "from_dict_type" $f}}{{end}}) - {{end}} -{{end}} -{{else if .ArrayValue}}type {{.PascalName}} []{{template "type" .ArrayValue}} -{{else if .MapValue}}{{.PascalName}} = {{template "type" .}} -{{else if .Enum}}class {{.PascalName}}(Enum): - {{if .Description}}"""{{.Comment " " 100 | trimSuffix "\"" }}"""{{end}} - {{range .Enum }} - {{.ConstantName}} = '{{.Content}}'{{end}}{{end}} -{{end}} -{{- define "from_dict_type" -}} - {{- if not .Entity }}None - {{- else if .Entity.ArrayValue }} - {{- if (or .Entity.ArrayValue.IsObject .Entity.ArrayValue.IsExternal) }}_repeated_dict(d, '{{.Name}}', {{template "type" .Entity.ArrayValue}}) - {{- else if .Entity.ArrayValue.Enum }}_repeated_enum(d, '{{.Name}}', {{template "type" .Entity.ArrayValue}}) - {{- else}}d.get('{{.Name}}', None){{- end -}} - {{- else if or .Entity.IsObject .Entity.IsExternal .Entity.IsEmpty }}_from_dict(d, '{{.Name}}', {{template "type" .Entity}}) - {{- else if .Entity.Enum }}_enum(d, '{{.Name}}', {{template "type" .Entity}}) - {{- else if and .IsHeader (or .Entity.IsInt64 .Entity.IsInt) }} int(d.get('{{.Name}}', None)) - {{- else}}d.get('{{.Name}}', None){{- end -}} -{{- end -}} -{{- define "as_request_type" -}} - {{- if not .Entity }}None # ERROR: No Type - {{- /* This should be done recursively, but recursion in text templates is not supported. */ -}} - {{- else if .Entity.ArrayValue }}[{{if or .Entity.ArrayValue.IsObject .Entity.ArrayValue.IsExternal}}v.as_dict(){{ else if .Entity.ArrayValue.Enum }}v.value{{else}}v{{end}} for v in self.{{template "safe-snake-name" .}}] - {{- else if or .Entity.IsObject .Entity.IsExternal .Entity.IsEmpty }}self.{{template "safe-snake-name" .}}.as_dict() - {{- else if .Entity.Enum }}self.{{template "safe-snake-name" .}}.value - {{- else}}self.{{template "safe-snake-name" .}}{{- end -}} -{{- end -}} -{{- define "type" -}} - {{- if not . }}any # ERROR: No Type - {{- else if .IsExternal }}{{.Package.Name}}.{{.PascalName}} - {{- else if .ArrayValue }}List[{{template "type" .ArrayValue}}] - {{- else if .MapValue }}Dict[str,{{template "type" .MapValue}}] - {{- else if .IsObject }}{{.PascalName}}{{if eq "List" .PascalName}}Request{{end}} - {{- else if .Enum }}{{.PascalName}} - {{- else if .IsString}}str - {{- else if .IsAny}}Any - {{- else if .IsEmpty}}{{.PascalName}} - {{- else if .IsBool}}bool - {{- else if .IsInt64}}int - {{- else if .IsFloat64}}float - {{- else if .IsInt}}int - {{- else if .IsByteStream}}BinaryIO - {{- else}}any /* MISSING TYPE */ - {{- end -}} -{{- end -}} - -{{- define "type-doc" -}} - {{- if .IsExternal }}:class:`{{.PascalName}}` - {{- else if .IsEmpty}}:class:`{{template "type" .}}` - {{- else if .ArrayValue }}List[{{template "type-doc" .ArrayValue}}] - {{- else if .MapValue }}Dict[str,{{template "type-doc" .MapValue}}] - {{- else if .IsObject }}:class:`{{.PascalName}}{{if eq "List" .PascalName}}Request{{end}}` - {{- else if .Enum }}:class:`{{.PascalName}}` - {{- else}}{{template "type" . }} - {{- end -}} -{{- end -}} - -{{range .Services}} -class {{.PascalName}}API:{{if .Description}} - """{{.Comment " " 110}}""" - {{end}} - def __init__(self, api_client{{if .IsDataPlane}}, control_plane{{end}}): - self._api = api_client - {{if .IsDataPlane -}} - self._control_plane = control_plane - self._data_plane_service = DataPlaneService() - {{end -}} - {{range .Subservices}} - self._{{.SnakeName}} = {{.PascalName}}API(self._api){{end}} - - {{range .Subservices}} - @property - def {{.SnakeName}}(self) -> {{.PascalName}}API: - {{if .Description}}"""{{.Summary}}"""{{end}} - return self._{{.SnakeName}} - {{end}} - - {{range .Waits}} - def {{template "safe-snake-name" .}}(self{{range .Binding}}, {{template "safe-snake-name" .PollField}}: {{template "type" .PollField.Entity}}{{end}}, - timeout=timedelta(minutes={{.Timeout}}), callback: Optional[Callable[[{{.Poll.Response.PascalName}}], None]] = None) -> {{.Poll.Response.PascalName}}: - deadline = time.time() + timeout.total_seconds() - target_states = ({{range .Success}}{{.Entity.PascalName}}.{{.ConstantName}}, {{end}}){{if .Failure}} - failure_states = ({{range .Failure}}{{.Entity.PascalName}}.{{.ConstantName}}, {{end}}){{end}} - status_message = 'polling...' - attempt = 1 - while time.time() < deadline: - poll = self.{{template "safe-snake-name" .Poll}}({{range $i, $b := .Binding}}{{if $i}}, {{end}}{{template "safe-snake-name" .PollField}}={{template "safe-snake-name" .PollField}}{{- end}}) - status = poll{{range .StatusPath}}.{{template "safe-snake-name" .}}{{end}} - {{if .ComplexMessagePath -}} - status_message = f'current status: {status}' - if poll.{{template "safe-snake-name" .MessagePathHead}}: - status_message = poll{{range .MessagePath}}.{{template "safe-snake-name" .}}{{end}} - {{- else if .MessagePath -}} - status_message = poll{{range .MessagePath}}.{{template "safe-snake-name" .}}{{end}} - {{- else -}} - status_message = f'current status: {status}' - {{- end}} - if status in target_states: - return poll - if callback: - callback(poll) - {{if .Failure -}} - if status in failure_states: - msg = f'failed to reach {{range $i, $e := .Success}}{{if $i}} or {{end}}{{$e.Content}}{{end}}, got {status}: {status_message}' - raise OperationFailed(msg) - {{end}}prefix = f"{{range $i, $b := .Binding}}{{if $i}}, {{end -}} - {{template "safe-snake-name" .PollField}}={{"{"}}{{template "safe-snake-name" .PollField}}{{"}"}} - {{- end}}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f'{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)') - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f'timed out after {timeout}: {status_message}') - {{end}} - - {{range .Methods}} - def {{template "safe-snake-name" .}}({{ template "method-parameters" . }}){{template "method-return-type" .}}: - {{if .Description}}"""{{.Comment " " 110 | trimSuffix "\"" }} - {{with .Request}}{{range .RequiredFields}} - :param {{template "safe-snake-name" .}}: {{template "type-doc" .Entity}}{{if .Description}} - {{.Comment " " 110 | trimSuffix "\"" }}{{end}} - {{- end}}{{range .NonRequiredFields}} - :param {{template "safe-snake-name" .}}: {{template "type-doc" .Entity}} (optional){{if .Description}} - {{.Comment " " 110 | trimSuffix "\"" }}{{end}} - {{- end}} - {{end}} - {{if and .Wait (and (not .IsCrudRead) (not (eq .SnakeName "get_run"))) -}} - :returns: - Long-running operation waiter for {{template "type-doc" .Wait.Poll.Response}}. - See :method:{{template "safe-snake-name" .Wait}} for more details. - {{- else if not .Response.IsEmpty }}:returns: {{if .Response.ArrayValue -}} - Iterator over {{template "type-doc" .Response.ArrayValue}} - {{- else if .Pagination -}} - Iterator over {{template "type-doc" .Pagination.Entity}} - {{- else -}} - {{template "type-doc" .Response}} - {{- end}}{{end}} - """{{end}} - {{if .Request -}} - {{template "method-serialize" .}} - {{- end}} - {{- if .Service.IsDataPlane}} - {{template "data-plane" .}} - {{- end}} - {{template "method-headers" . }} - {{if .Response.HasHeaderField -}} - {{template "method-response-headers" . }} - {{- end}} - {{template "method-call" .}} - - {{if and .Wait (and (not .IsCrudRead) (not (eq .SnakeName "get_run"))) }} - def {{.SnakeName}}_and_wait({{ template "method-parameters" . }}, - timeout=timedelta(minutes={{.Wait.Timeout}})) -> {{.Wait.Poll.Response.PascalName}}: - return self.{{template "safe-snake-name" .}}({{range $i, $x := .Request.Fields}}{{if $i}}, {{end}}{{template "safe-snake-name" .}}={{template "safe-snake-name" .}}{{end}}).result(timeout=timeout) - {{end}} - {{end -}} -{{- end}} - -{{define "data-plane" -}} - def info_getter(): - response = self._control_plane.{{.Service.DataPlaneInfoMethod.SnakeName}}( - {{- range .Service.DataPlaneInfoMethod.Request.Fields }} - {{.SnakeName}} = {{.SnakeName}}, - {{- end}} - ) - if response.{{(index .DataPlaneInfoFields 0).SnakeName}} is None: - raise Exception("Resource does not support direct Data Plane access") - return response{{range .DataPlaneInfoFields}}.{{.SnakeName}}{{end}} - - get_params = [{{- range .Service.DataPlaneInfoMethod.Request.Fields }}{{.SnakeName}},{{end}}] - data_plane_details = self._data_plane_service.get_data_plane_details('{{.SnakeName}}', get_params, info_getter, self._api.get_oauth_token) - token = data_plane_details.token - - def auth(r: requests.PreparedRequest) -> requests.PreparedRequest: - authorization = f"{token.token_type} {token.access_token}" - r.headers["Authorization"] = authorization - return r -{{- end}} - -{{define "method-parameters" -}} - self{{if .Request}} - {{- if .Request.MapValue }}, contents: {{template "type" .Request }}{{ end }} - {{range .Request.RequiredFields}}, {{template "safe-snake-name" .}}: {{template "type" .Entity}}{{end}} - {{if .Request.NonRequiredFields}}, * - {{range .Request.NonRequiredFields}}, {{template "safe-snake-name" .}}: Optional[{{template "type" .Entity}}] = None{{end}} - {{- end}} - {{- end}} -{{- end}} - -{{define "method-serialize" -}} - {{if and .Request.HasJsonField .RequestBodyField -}} - body = {{template "safe-snake-name" .RequestBodyField}} - {{- else -}} - {{if or .Request.HasJsonField .Request.HasQueryField -}} - {{if .Request.HasJsonField}}body = {}{{end}}{{if .Request.HasQueryField}} - query = {}{{end}} - {{- range .Request.Fields}}{{ if and (not .IsPath) (not .IsHeader) }} - {{- if .IsQuery }} - if {{template "safe-snake-name" .}} is not None: query['{{.Name}}'] = {{template "method-param-bind" .}}{{end}} - {{- if .IsJson }} - if {{template "safe-snake-name" .}} is not None: body['{{.Name}}'] = {{template "method-param-bind" .}}{{end}} - {{- end}} - {{- end}} - {{- end}} - {{- end}} -{{- end}} - -{{ define "method-headers" -}} - headers = { - {{- range $k, $v := .FixedRequestHeaders}}'{{ $k }}': '{{ $v }}',{{ end -}} - } -{{- end }} - -{{ define "method-response-headers" -}} - response_headers = [ - {{- range $h := .ResponseHeaders}}'{{ $h.Name }}',{{ end -}} - ] -{{- end }} - -{{- define "method-param-bind" -}} - {{- if not .Entity }}None # ERROR: No Type - {{- else if .Entity.ArrayValue }}[ - {{- if or .Entity.ArrayValue.IsObject .Entity.ArrayValue.IsExternal -}}v.as_dict() - {{- else if .Entity.ArrayValue.Enum -}}v.value - {{- else}}v{{end}} for v in {{template "safe-snake-name" .}}] - {{- else if .Entity.IsObject }}{{template "safe-snake-name" .}}.as_dict() - {{- else if .Entity.Enum }}{{template "safe-snake-name" .}}.value - {{- else}}{{template "safe-snake-name" .}}{{- end -}} -{{- end -}} - -{{define "method-call" -}} - {{if .Pagination -}}{{template "method-call-paginated" .}} - {{- else if and .Wait (and (not .IsCrudRead) (not (eq .SnakeName "get_run"))) -}}{{template "method-call-retried" .}} - {{- else}}{{template "method-call-default" .}}{{end}} -{{- end}} - -{{define "method-call-retried" -}} - {{if .Response}}op_response = {{end}}{{template "method-do" .}} - return Wait(self.{{template "safe-snake-name" .Wait}} - {{if .Response}}, response = {{.Response.PascalName}}.from_dict(op_response){{end}} - {{range .Wait.Binding}}, {{template "safe-snake-name" .PollField}}={{if .IsResponseBind}}op_response['{{.Bind.Name}}']{{else}}{{template "safe-snake-name" .Bind}}{{end}} - {{- end}}) -{{- end}} - -{{define "method-call-paginated" -}} - {{if .Pagination.MultiRequest}} - {{if .NeedsOffsetDedupe -}} - # deduplicate items that may have been added during iteration - seen = set() - {{- end}}{{if and .Pagination.Offset (not (eq .Path "/api/2.1/clusters/events")) }} - query['{{.Pagination.Offset.Name}}'] = - {{- if eq .Pagination.Increment 1 -}} - 1 - {{- else if contains .Path "/scim/v2/" -}} - 1 - {{- else -}} - 0 - {{- end}}{{end}}{{if and .Pagination.Limit (contains .Path "/scim/v2/")}} - if "{{.Pagination.Limit.Name}}" not in query: query['{{.Pagination.Limit.Name}}'] = 100 - {{- end}} - while True: - json = {{template "method-do" .}} - if '{{.Pagination.Results.Name}}' in json: - for v in json['{{.Pagination.Results.Name}}']: - {{if .NeedsOffsetDedupe -}} - i = v['{{.IdentifierField.Name}}'] - if i in seen: - continue - seen.add(i) - {{end -}} - yield {{.Pagination.Entity.PascalName}}.from_dict(v) - {{ if .Pagination.Token -}} - if '{{.Pagination.Token.Bind.Name}}' not in json or not json['{{.Pagination.Token.Bind.Name}}']: - return - {{if or (eq "GET" .Verb) (eq "HEAD" .Verb)}}query{{else}}body{{end}}['{{.Pagination.Token.PollField.Name}}'] = json['{{.Pagination.Token.Bind.Name}}'] - {{- else if eq .Path "/api/2.1/clusters/events" -}} - if 'next_page' not in json or not json['next_page']: - return - body = json['next_page'] - {{- else -}} - if '{{.Pagination.Results.Name}}' not in json or not json['{{.Pagination.Results.Name}}']: - return - {{ if eq .Pagination.Increment 1 -}} - query['{{.Pagination.Offset.Name}}'] += 1 - {{- else -}} - query['{{.Pagination.Offset.Name}}'] += len(json['{{.Pagination.Results.Name}}']) - {{- end}} - {{- end}} - {{else -}} - json = {{template "method-do" .}} - parsed = {{.Response.PascalName}}.from_dict(json).{{template "safe-snake-name" .Pagination.Results}} - return parsed if parsed is not None else [] - {{end}} -{{- end}} - -{{define "method-call-default" -}} - {{if not .Response.IsEmpty -}} - res = {{end}}{{template "method-do" .}} - {{if not .Response.IsEmpty -}} - {{- if .Response.ArrayValue -}} - return [{{.Response.ArrayValue.PascalName}}.from_dict(v) for v in res] - {{- else if .Response.MapValue -}} - return res - {{- else -}} - return {{template "type" .Response}}.from_dict(res) - {{- end}} - {{- end}} -{{- end}} - -{{define "method-do" -}} - self._api.do('{{.Verb}}', - {{- if .Service.IsDataPlane -}} - url=data_plane_details.endpoint_url - {{- else -}} - {{ template "path" . }} - {{- end -}} - {{if .Request}} - {{- if .Request.HasQueryField}}, query=query{{end}} - {{- if .Request.MapValue}}, body=contents - {{- else if .Request.HasJsonField}}, body=body{{end}} - {{end}} - , headers=headers - {{if .Response.HasHeaderField -}} - , response_headers=response_headers - {{- end}} - {{- if and .IsRequestByteStream .RequestBodyField }}, data={{template "safe-snake-name" .RequestBodyField}}{{ end }} - {{- if .Service.IsDataPlane -}} - ,auth=auth - {{- end -}} - {{- if .IsResponseByteStream }}, raw=True{{ end }}) -{{- end}} - -{{- define "path" -}} -{{- if .PathParts -}} - f'{{range .PathParts -}} - {{- .Prefix -}} - {{- if .Field -}} - {{- "{" -}} - {{- if .Field.IsPathMultiSegment -}}_escape_multi_segment_path_parameter({{ template "path-parameter" . }}) - {{- else -}}{{ template "path-parameter" . }} - {{- end -}} - {{- "}" -}} - {{- else if .IsAccountId}} - {{- "{" -}} - self._api.account_id - {{- "}" -}} - {{- end -}} - {{- end }}' -{{- else -}} - '{{.Path}}' -{{- end -}} -{{- end -}} - -{{- define "path-parameter" -}} - {{template "safe-snake-name" .Field}}{{with .Field.Entity.Enum}}.value{{end}} -{{- end -}} - -{{define "method-return-type" -}} - {{if and .Wait (and (not .IsCrudRead) (not (eq .SnakeName "get_run"))) }} -> Wait[{{.Wait.Poll.Response.PascalName}}] - {{- else if not .Response.IsEmpty }} -> {{if .Response.ArrayValue -}} - Iterator[{{template "type" .Response.ArrayValue}}] - {{- else if .Pagination -}} - Iterator[{{template "type" .Pagination.Entity}}] - {{- else -}} - {{- if .Response.IsExternal -}} - {{.Response.Package.Name}}.{{.Response.PascalName}} - {{- else -}} - {{.Response.PascalName}} - {{- end -}} - {{- end}}{{end}} -{{- end}}