From 7e8f03eec45c2daf7d0f9ef59448475d04371d0a Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Tue, 6 Feb 2024 17:34:46 +0100 Subject: [PATCH 1/6] Do not terminate listing for token-based pagination resources on empty response --- .codegen/service.py.tmpl | 19 ++-- .gitattributes | 2 +- databricks/sdk/service/catalog.py | 56 +++++------- databricks/sdk/service/compute.py | 14 ++- databricks/sdk/service/iam.py | 90 +++++++++---------- databricks/sdk/service/jobs.py | 14 ++- databricks/sdk/service/ml.py | 63 ++++++------- databricks/sdk/service/oauth2.py | 7 +- databricks/sdk/service/pipelines.py | 14 ++- databricks/sdk/service/settings.py | 14 ++- databricks/sdk/service/sharing.py | 7 +- databricks/sdk/service/sql.py | 37 ++++---- databricks/sdk/service/vectorsearch.py | 14 ++- databricks/sdk/service/workspace.py | 7 +- examples/r/wait_catalog_workspace_bindings.py | 5 ++ 15 files changed, 164 insertions(+), 199 deletions(-) create mode 100755 examples/r/wait_catalog_workspace_bindings.py diff --git a/.codegen/service.py.tmpl b/.codegen/service.py.tmpl index dd1abfea0..ad0009d7f 100644 --- a/.codegen/service.py.tmpl +++ b/.codegen/service.py.tmpl @@ -270,16 +270,15 @@ class {{.Name}}API:{{if .Description}} {{- end}} while True: json = {{template "method-do" .}} - if '{{.Pagination.Results.Name}}' not in json or not json['{{.Pagination.Results.Name}}']: - return - 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.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 eq .Path "/api/2.0/clusters/events" -}} if 'next_page' not in json or not json['next_page']: return diff --git a/.gitattributes b/.gitattributes index d36f22eeb..209ba7dec 100755 --- a/.gitattributes +++ b/.gitattributes @@ -190,6 +190,7 @@ examples/queries/create_queries.py linguist-generated=true examples/queries/get_queries.py linguist-generated=true examples/queries/update_queries.py linguist-generated=true examples/query_history/list_sql_query_history.py linguist-generated=true +examples/r/wait_catalog_workspace_bindings.py linguist-generated=true examples/recipients/create_recipients.py linguist-generated=true examples/recipients/get_recipients.py linguist-generated=true examples/recipients/list_recipients.py linguist-generated=true @@ -285,7 +286,6 @@ examples/workspace/list_workspace_integration.py linguist-generated=true examples/workspace_assignment/list_workspace_assignment_on_aws.py linguist-generated=true examples/workspace_assignment/update_workspace_assignment_on_aws.py linguist-generated=true examples/workspace_bindings/get_catalog_workspace_bindings.py linguist-generated=true -examples/workspace_bindings/update_catalog_workspace_bindings.py linguist-generated=true examples/workspace_conf/get_status_repos.py linguist-generated=true examples/workspaces/create_workspaces.py linguist-generated=true examples/workspaces/get_workspaces.py linguist-generated=true diff --git a/databricks/sdk/service/catalog.py b/databricks/sdk/service/catalog.py index c1c068aff..5c240c62e 100755 --- a/databricks/sdk/service/catalog.py +++ b/databricks/sdk/service/catalog.py @@ -5804,10 +5804,9 @@ def list(self, '/api/2.1/unity-catalog/external-locations', query=query, headers=headers) - if 'external_locations' not in json or not json['external_locations']: - return - for v in json['external_locations']: - yield ExternalLocationInfo.from_dict(v) + if 'external_locations' in json: + for v in json['external_locations']: + yield ExternalLocationInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -5988,10 +5987,9 @@ def list(self, while True: json = self._api.do('GET', '/api/2.1/unity-catalog/functions', query=query, headers=headers) - if 'functions' not in json or not json['functions']: - return - for v in json['functions']: - yield FunctionInfo.from_dict(v) + if 'functions' in json: + for v in json['functions']: + yield FunctionInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -6718,10 +6716,9 @@ def list(self, f'/api/2.1/unity-catalog/models/{full_name}/versions', query=query, headers=headers) - if 'model_versions' not in json or not json['model_versions']: - return - for v in json['model_versions']: - yield ModelVersionInfo.from_dict(v) + if 'model_versions' in json: + for v in json['model_versions']: + yield ModelVersionInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -6928,10 +6925,9 @@ def list(self, while True: json = self._api.do('GET', '/api/2.1/unity-catalog/models', query=query, headers=headers) - if 'registered_models' not in json or not json['registered_models']: - return - for v in json['registered_models']: - yield RegisteredModelInfo.from_dict(v) + if 'registered_models' in json: + for v in json['registered_models']: + yield RegisteredModelInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -7112,10 +7108,9 @@ def list(self, while True: json = self._api.do('GET', '/api/2.1/unity-catalog/schemas', query=query, headers=headers) - if 'schemas' not in json or not json['schemas']: - return - for v in json['schemas']: - yield SchemaInfo.from_dict(v) + if 'schemas' in json: + for v in json['schemas']: + yield SchemaInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -7306,10 +7301,9 @@ def list(self, '/api/2.1/unity-catalog/storage-credentials', query=query, headers=headers) - if 'storage_credentials' not in json or not json['storage_credentials']: - return - for v in json['storage_credentials']: - yield StorageCredentialInfo.from_dict(v) + if 'storage_credentials' in json: + for v in json['storage_credentials']: + yield StorageCredentialInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -7709,10 +7703,9 @@ def list(self, while True: json = self._api.do('GET', '/api/2.1/unity-catalog/tables', query=query, headers=headers) - if 'tables' not in json or not json['tables']: - return - for v in json['tables']: - yield TableInfo.from_dict(v) + if 'tables' in json: + for v in json['tables']: + yield TableInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -7765,10 +7758,9 @@ def list_summaries(self, while True: json = self._api.do('GET', '/api/2.1/unity-catalog/table-summaries', query=query, headers=headers) - if 'tables' not in json or not json['tables']: - return - for v in json['tables']: - yield TableSummary.from_dict(v) + if 'tables' in json: + for v in json['tables']: + yield TableSummary.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] diff --git a/databricks/sdk/service/compute.py b/databricks/sdk/service/compute.py index 816f0db3a..efa4c54b8 100755 --- a/databricks/sdk/service/compute.py +++ b/databricks/sdk/service/compute.py @@ -6405,10 +6405,9 @@ def events(self, while True: json = self._api.do('POST', '/api/2.0/clusters/events', body=body, headers=headers) - if 'events' not in json or not json['events']: - return - for v in json['events']: - yield ClusterEvent.from_dict(v) + if 'events' in json: + for v in json['events']: + yield ClusterEvent.from_dict(v) if 'next_page' not in json or not json['next_page']: return body = json['next_page'] @@ -7733,10 +7732,9 @@ def list(self, while True: json = self._api.do('GET', '/api/2.0/policy-families', query=query, headers=headers) - if 'policy_families' not in json or not json['policy_families']: - return - for v in json['policy_families']: - yield PolicyFamily.from_dict(v) + if 'policy_families' in json: + for v in json['policy_families']: + yield PolicyFamily.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] diff --git a/databricks/sdk/service/iam.py b/databricks/sdk/service/iam.py index 5a4131f19..34cf118d7 100755 --- a/databricks/sdk/service/iam.py +++ b/databricks/sdk/service/iam.py @@ -1465,14 +1465,13 @@ def list(self, f'/api/2.0/accounts/{self._api.account_id}/scim/v2/Groups', query=query, headers=headers) - if 'Resources' not in json or not json['Resources']: - return - for v in json['Resources']: - i = v['id'] - if i in seen: - continue - seen.add(i) - yield Group.from_dict(v) + if 'Resources' in json: + for v in json['Resources']: + i = v['id'] + if i in seen: + continue + seen.add(i) + yield Group.from_dict(v) query['startIndex'] += len(json['Resources']) def patch(self, @@ -1705,14 +1704,13 @@ def list(self, f'/api/2.0/accounts/{self._api.account_id}/scim/v2/ServicePrincipals', query=query, headers=headers) - if 'Resources' not in json or not json['Resources']: - return - for v in json['Resources']: - i = v['id'] - if i in seen: - continue - seen.add(i) - yield ServicePrincipal.from_dict(v) + if 'Resources' in json: + for v in json['Resources']: + i = v['id'] + if i in seen: + continue + seen.add(i) + yield ServicePrincipal.from_dict(v) query['startIndex'] += len(json['Resources']) def patch(self, @@ -2005,14 +2003,13 @@ def list(self, f'/api/2.0/accounts/{self._api.account_id}/scim/v2/Users', query=query, headers=headers) - if 'Resources' not in json or not json['Resources']: - return - for v in json['Resources']: - i = v['id'] - if i in seen: - continue - seen.add(i) - yield User.from_dict(v) + if 'Resources' in json: + for v in json['Resources']: + i = v['id'] + if i in seen: + continue + seen.add(i) + yield User.from_dict(v) query['startIndex'] += len(json['Resources']) def patch(self, @@ -2267,14 +2264,13 @@ def list(self, if "count" not in query: query['count'] = 100 while True: json = self._api.do('GET', '/api/2.0/preview/scim/v2/Groups', query=query, headers=headers) - if 'Resources' not in json or not json['Resources']: - return - for v in json['Resources']: - i = v['id'] - if i in seen: - continue - seen.add(i) - yield Group.from_dict(v) + if 'Resources' in json: + for v in json['Resources']: + i = v['id'] + if i in seen: + continue + seen.add(i) + yield Group.from_dict(v) query['startIndex'] += len(json['Resources']) def patch(self, @@ -2646,14 +2642,13 @@ def list(self, '/api/2.0/preview/scim/v2/ServicePrincipals', query=query, headers=headers) - if 'Resources' not in json or not json['Resources']: - return - for v in json['Resources']: - i = v['id'] - if i in seen: - continue - seen.add(i) - yield ServicePrincipal.from_dict(v) + if 'Resources' in json: + for v in json['Resources']: + i = v['id'] + if i in seen: + continue + seen.add(i) + yield ServicePrincipal.from_dict(v) query['startIndex'] += len(json['Resources']) def patch(self, @@ -2955,14 +2950,13 @@ def list(self, if "count" not in query: query['count'] = 100 while True: json = self._api.do('GET', '/api/2.0/preview/scim/v2/Users', query=query, headers=headers) - if 'Resources' not in json or not json['Resources']: - return - for v in json['Resources']: - i = v['id'] - if i in seen: - continue - seen.add(i) - yield User.from_dict(v) + if 'Resources' in json: + for v in json['Resources']: + i = v['id'] + if i in seen: + continue + seen.add(i) + yield User.from_dict(v) query['startIndex'] += len(json['Resources']) def patch(self, diff --git a/databricks/sdk/service/jobs.py b/databricks/sdk/service/jobs.py index e7fef2de3..c9453ae04 100755 --- a/databricks/sdk/service/jobs.py +++ b/databricks/sdk/service/jobs.py @@ -4945,10 +4945,9 @@ def list(self, while True: json = self._api.do('GET', '/api/2.1/jobs/list', query=query, headers=headers) - if 'jobs' not in json or not json['jobs']: - return - for v in json['jobs']: - yield BaseJob.from_dict(v) + if 'jobs' in json: + for v in json['jobs']: + yield BaseJob.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -5017,10 +5016,9 @@ def list_runs(self, while True: json = self._api.do('GET', '/api/2.1/jobs/runs/list', query=query, headers=headers) - if 'runs' not in json or not json['runs']: - return - for v in json['runs']: - yield BaseRun.from_dict(v) + if 'runs' in json: + for v in json['runs']: + yield BaseRun.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] diff --git a/databricks/sdk/service/ml.py b/databricks/sdk/service/ml.py index ff2da8655..8194f5436 100755 --- a/databricks/sdk/service/ml.py +++ b/databricks/sdk/service/ml.py @@ -3710,10 +3710,9 @@ def get_history(self, while True: json = self._api.do('GET', '/api/2.0/mlflow/metrics/get-history', query=query, headers=headers) - if 'metrics' not in json or not json['metrics']: - return - for v in json['metrics']: - yield Metric.from_dict(v) + if 'metrics' in json: + for v in json['metrics']: + yield Metric.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -3807,10 +3806,9 @@ def list_artifacts(self, while True: json = self._api.do('GET', '/api/2.0/mlflow/artifacts/list', query=query, headers=headers) - if 'files' not in json or not json['files']: - return - for v in json['files']: - yield FileInfo.from_dict(v) + if 'files' in json: + for v in json['files']: + yield FileInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -3844,10 +3842,9 @@ def list_experiments(self, while True: json = self._api.do('GET', '/api/2.0/mlflow/experiments/list', query=query, headers=headers) - if 'experiments' not in json or not json['experiments']: - return - for v in json['experiments']: - yield Experiment.from_dict(v) + if 'experiments' in json: + for v in json['experiments']: + yield Experiment.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -4125,10 +4122,9 @@ def search_experiments(self, while True: json = self._api.do('POST', '/api/2.0/mlflow/experiments/search', body=body, headers=headers) - if 'experiments' not in json or not json['experiments']: - return - for v in json['experiments']: - yield Experiment.from_dict(v) + if 'experiments' in json: + for v in json['experiments']: + yield Experiment.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return body['page_token'] = json['next_page_token'] @@ -4186,10 +4182,9 @@ def search_runs(self, while True: json = self._api.do('POST', '/api/2.0/mlflow/runs/search', body=body, headers=headers) - if 'runs' not in json or not json['runs']: - return - for v in json['runs']: - yield Run.from_dict(v) + if 'runs' in json: + for v in json['runs']: + yield Run.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return body['page_token'] = json['next_page_token'] @@ -4903,10 +4898,9 @@ def list_models(self, while True: json = self._api.do('GET', '/api/2.0/mlflow/registered-models/list', query=query, headers=headers) - if 'registered_models' not in json or not json['registered_models']: - return - for v in json['registered_models']: - yield Model.from_dict(v) + if 'registered_models' in json: + for v in json['registered_models']: + yield Model.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -4963,10 +4957,9 @@ def list_webhooks(self, while True: json = self._api.do('GET', '/api/2.0/mlflow/registry-webhooks/list', query=query, headers=headers) - if 'webhooks' not in json or not json['webhooks']: - return - for v in json['webhooks']: - yield RegistryWebhook.from_dict(v) + if 'webhooks' in json: + for v in json['webhooks']: + yield RegistryWebhook.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -5062,10 +5055,9 @@ def search_model_versions(self, while True: json = self._api.do('GET', '/api/2.0/mlflow/model-versions/search', query=query, headers=headers) - if 'model_versions' not in json or not json['model_versions']: - return - for v in json['model_versions']: - yield ModelVersion.from_dict(v) + if 'model_versions' in json: + for v in json['model_versions']: + yield ModelVersion.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -5108,10 +5100,9 @@ def search_models(self, '/api/2.0/mlflow/registered-models/search', query=query, headers=headers) - if 'registered_models' not in json or not json['registered_models']: - return - for v in json['registered_models']: - yield Model.from_dict(v) + if 'registered_models' in json: + for v in json['registered_models']: + yield Model.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] diff --git a/databricks/sdk/service/oauth2.py b/databricks/sdk/service/oauth2.py index 9a48d9082..bf0fc4436 100755 --- a/databricks/sdk/service/oauth2.py +++ b/databricks/sdk/service/oauth2.py @@ -629,10 +629,9 @@ def list(self, f'/api/2.0/accounts/{self._api.account_id}/oauth2/published-apps/', query=query, headers=headers) - if 'apps' not in json or not json['apps']: - return - for v in json['apps']: - yield PublishedAppOutput.from_dict(v) + if 'apps' in json: + for v in json['apps']: + yield PublishedAppOutput.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] diff --git a/databricks/sdk/service/pipelines.py b/databricks/sdk/service/pipelines.py index 6ef8efa73..b27b5cbe0 100755 --- a/databricks/sdk/service/pipelines.py +++ b/databricks/sdk/service/pipelines.py @@ -1891,10 +1891,9 @@ def list_pipeline_events(self, f'/api/2.0/pipelines/{pipeline_id}/events', query=query, headers=headers) - if 'events' not in json or not json['events']: - return - for v in json['events']: - yield PipelineEvent.from_dict(v) + if 'events' in json: + for v in json['events']: + yield PipelineEvent.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -1940,10 +1939,9 @@ def list_pipelines(self, while True: json = self._api.do('GET', '/api/2.0/pipelines', query=query, headers=headers) - if 'statuses' not in json or not json['statuses']: - return - for v in json['statuses']: - yield PipelineStateInfo.from_dict(v) + if 'statuses' in json: + for v in json['statuses']: + yield PipelineStateInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] diff --git a/databricks/sdk/service/settings.py b/databricks/sdk/service/settings.py index c13adb756..e542c250f 100755 --- a/databricks/sdk/service/settings.py +++ b/databricks/sdk/service/settings.py @@ -2075,10 +2075,9 @@ def list_network_connectivity_configurations(self, f'/api/2.0/accounts/{self._api.account_id}/network-connectivity-configs', query=query, headers=headers) - if 'items' not in json or not json['items']: - return - for v in json['items']: - yield NetworkConnectivityConfiguration.from_dict(v) + if 'items' in json: + for v in json['items']: + yield NetworkConnectivityConfiguration.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -2110,10 +2109,9 @@ def list_private_endpoint_rules( f'/api/2.0/accounts/{self._api.account_id}/network-connectivity-configs/{network_connectivity_config_id}/private-endpoint-rules', query=query, headers=headers) - if 'items' not in json or not json['items']: - return - for v in json['items']: - yield NccAzurePrivateEndpointRule.from_dict(v) + if 'items' in json: + for v in json['items']: + yield NccAzurePrivateEndpointRule.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] diff --git a/databricks/sdk/service/sharing.py b/databricks/sdk/service/sharing.py index 37753bdcb..ef4f98684 100755 --- a/databricks/sdk/service/sharing.py +++ b/databricks/sdk/service/sharing.py @@ -1644,10 +1644,9 @@ def list(self, while True: json = self._api.do('GET', '/api/2.1/unity-catalog/clean-rooms', query=query, headers=headers) - if 'clean_rooms' not in json or not json['clean_rooms']: - return - for v in json['clean_rooms']: - yield CleanRoomInfo.from_dict(v) + if 'clean_rooms' in json: + for v in json['clean_rooms']: + yield CleanRoomInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] diff --git a/databricks/sdk/service/sql.py b/databricks/sdk/service/sql.py index 7eacd64e3..d73648292 100755 --- a/databricks/sdk/service/sql.py +++ b/databricks/sdk/service/sql.py @@ -4073,14 +4073,13 @@ def list(self, query['page'] = 1 while True: json = self._api.do('GET', '/api/2.0/preview/sql/dashboards', query=query, headers=headers) - if 'results' not in json or not json['results']: - return - for v in json['results']: - i = v['id'] - if i in seen: - continue - seen.add(i) - yield Dashboard.from_dict(v) + if 'results' in json: + for v in json['results']: + i = v['id'] + if i in seen: + continue + seen.add(i) + yield Dashboard.from_dict(v) query['page'] += 1 def restore(self, dashboard_id: str): @@ -4387,14 +4386,13 @@ def list(self, query['page'] = 1 while True: json = self._api.do('GET', '/api/2.0/preview/sql/queries', query=query, headers=headers) - if 'results' not in json or not json['results']: - return - for v in json['results']: - i = v['id'] - if i in seen: - continue - seen.add(i) - yield Query.from_dict(v) + if 'results' in json: + for v in json['results']: + i = v['id'] + if i in seen: + continue + seen.add(i) + yield Query.from_dict(v) query['page'] += 1 def restore(self, query_id: str): @@ -4499,10 +4497,9 @@ def list(self, while True: json = self._api.do('GET', '/api/2.0/sql/history/queries', query=query, headers=headers) - if 'res' not in json or not json['res']: - return - for v in json['res']: - yield QueryInfo.from_dict(v) + if 'res' in json: + for v in json['res']: + yield QueryInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] diff --git a/databricks/sdk/service/vectorsearch.py b/databricks/sdk/service/vectorsearch.py index 95ad717cb..52aaf172a 100755 --- a/databricks/sdk/service/vectorsearch.py +++ b/databricks/sdk/service/vectorsearch.py @@ -978,10 +978,9 @@ def list_endpoints(self, *, page_token: Optional[str] = None) -> Iterator[Endpoi while True: json = self._api.do('GET', '/api/2.0/vector-search/endpoints', query=query, headers=headers) - if 'endpoints' not in json or not json['endpoints']: - return - for v in json['endpoints']: - yield EndpointInfo.from_dict(v) + if 'endpoints' in json: + for v in json['endpoints']: + yield EndpointInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] @@ -1117,10 +1116,9 @@ def list_indexes(self, while True: json = self._api.do('GET', '/api/2.0/vector-search/indexes', query=query, headers=headers) - if 'vector_indexes' not in json or not json['vector_indexes']: - return - for v in json['vector_indexes']: - yield MiniVectorIndex.from_dict(v) + if 'vector_indexes' in json: + for v in json['vector_indexes']: + yield MiniVectorIndex.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['page_token'] = json['next_page_token'] diff --git a/databricks/sdk/service/workspace.py b/databricks/sdk/service/workspace.py index cf5dacc3a..142fff57d 100755 --- a/databricks/sdk/service/workspace.py +++ b/databricks/sdk/service/workspace.py @@ -1516,10 +1516,9 @@ def list(self, while True: json = self._api.do('GET', '/api/2.0/repos', query=query, headers=headers) - if 'repos' not in json or not json['repos']: - return - for v in json['repos']: - yield RepoInfo.from_dict(v) + if 'repos' in json: + for v in json['repos']: + yield RepoInfo.from_dict(v) if 'next_page_token' not in json or not json['next_page_token']: return query['next_page_token'] = json['next_page_token'] diff --git a/examples/r/wait_catalog_workspace_bindings.py b/examples/r/wait_catalog_workspace_bindings.py new file mode 100755 index 000000000..1352ed169 --- /dev/null +++ b/examples/r/wait_catalog_workspace_bindings.py @@ -0,0 +1,5 @@ +from databricks.sdk import WorkspaceClient + +w = WorkspaceClient() + +w.r.wait(update_function) From d35a6f5c5a62e698c9d5e270ac4b2029aea38eda Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Thu, 8 Feb 2024 10:44:17 +0100 Subject: [PATCH 2/6] fix --- .codegen/service.py.tmpl | 14 +++++++++----- databricks/sdk/clock.py | 0 databricks/sdk/service/iam.py | 12 ++++++++++++ databricks/sdk/service/sql.py | 4 ++++ tests/clock.py | 0 tests/test_core.py | 13 +++++++------ 6 files changed, 32 insertions(+), 11 deletions(-) create mode 100644 databricks/sdk/clock.py create mode 100644 tests/clock.py diff --git a/.codegen/service.py.tmpl b/.codegen/service.py.tmpl index ad0009d7f..052116424 100644 --- a/.codegen/service.py.tmpl +++ b/.codegen/service.py.tmpl @@ -279,19 +279,23 @@ class {{.Name}}API:{{if .Description}} seen.add(i) {{end -}} yield {{.Pagination.Entity.PascalName}}.from_dict(v) - {{if eq .Path "/api/2.0/clusters/events" -}} + {{ if .Pagination.Token -}} + if '{{.Pagination.Token.Bind.Name}}' not in json or not json['{{.Pagination.Token.Bind.Name}}']: + return + {{if eq "GET" .Verb}}query{{else}}body{{end}}['{{.Pagination.Token.PollField.Name}}'] = json['{{.Pagination.Token.Bind.Name}}'] + {{- else if eq .Path "/api/2.0/clusters/events" -}} if 'next_page' not in json or not json['next_page']: return body = json['next_page'] - {{- else if .Pagination.Token -}} - if '{{.Pagination.Token.Bind.Name}}' not in json or not json['{{.Pagination.Token.Bind.Name}}']: + {{- else -}} + if '{{.Pagination.Results.Name}}' not in json or not json['{{.Pagination.Results.Name}}']: return - {{if eq "GET" .Verb}}query{{else}}body{{end}}['{{.Pagination.Token.PollField.Name}}'] = json['{{.Pagination.Token.Bind.Name}}'] - {{- else if eq .Pagination.Increment 1 -}} + {{ 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}} diff --git a/databricks/sdk/clock.py b/databricks/sdk/clock.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/service/iam.py b/databricks/sdk/service/iam.py index 34cf118d7..18d16d3fb 100755 --- a/databricks/sdk/service/iam.py +++ b/databricks/sdk/service/iam.py @@ -1472,6 +1472,8 @@ def list(self, continue seen.add(i) yield Group.from_dict(v) + if 'Resources' not in json or not json['Resources']: + return query['startIndex'] += len(json['Resources']) def patch(self, @@ -1711,6 +1713,8 @@ def list(self, continue seen.add(i) yield ServicePrincipal.from_dict(v) + if 'Resources' not in json or not json['Resources']: + return query['startIndex'] += len(json['Resources']) def patch(self, @@ -2010,6 +2014,8 @@ def list(self, continue seen.add(i) yield User.from_dict(v) + if 'Resources' not in json or not json['Resources']: + return query['startIndex'] += len(json['Resources']) def patch(self, @@ -2271,6 +2277,8 @@ def list(self, continue seen.add(i) yield Group.from_dict(v) + if 'Resources' not in json or not json['Resources']: + return query['startIndex'] += len(json['Resources']) def patch(self, @@ -2649,6 +2657,8 @@ def list(self, continue seen.add(i) yield ServicePrincipal.from_dict(v) + if 'Resources' not in json or not json['Resources']: + return query['startIndex'] += len(json['Resources']) def patch(self, @@ -2957,6 +2967,8 @@ def list(self, continue seen.add(i) yield User.from_dict(v) + if 'Resources' not in json or not json['Resources']: + return query['startIndex'] += len(json['Resources']) def patch(self, diff --git a/databricks/sdk/service/sql.py b/databricks/sdk/service/sql.py index d73648292..b7dc9a297 100755 --- a/databricks/sdk/service/sql.py +++ b/databricks/sdk/service/sql.py @@ -4080,6 +4080,8 @@ def list(self, continue seen.add(i) yield Dashboard.from_dict(v) + if 'results' not in json or not json['results']: + return query['page'] += 1 def restore(self, dashboard_id: str): @@ -4393,6 +4395,8 @@ def list(self, continue seen.add(i) yield Query.from_dict(v) + if 'results' not in json or not json['results']: + return query['page'] += 1 def restore(self, query_id: str): diff --git a/tests/clock.py b/tests/clock.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_core.py b/tests/test_core.py index ca2eaac31..4935a18f6 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -27,6 +27,7 @@ from databricks.sdk.version import __version__ from .conftest import noop_credentials +from .clock import FakeClock def test_parse_dsn(): @@ -422,7 +423,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_')) + api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock()) res = api_client.do('GET', '/foo') assert 'foo' in res @@ -445,7 +446,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_')) + api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock()) res = api_client.do('GET', '/foo') assert 'foo' in res @@ -462,7 +463,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_', retry_timeout_seconds=1)) + api_client = ApiClient(Config(host=host, token='_', retry_timeout_seconds=1), clock=FakeClock()) with pytest.raises(TimeoutError): api_client.do('GET', '/foo') @@ -484,7 +485,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_')) + api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock()) res = api_client.do('GET', '/foo') assert 'foo' in res @@ -502,7 +503,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_')) + api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock()) with pytest.raises(DatabricksError): api_client.do('GET', '/foo') @@ -520,7 +521,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_')) + api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock()) res = api_client.do('GET', '/foo') assert 'foo' in res From e8f82240c78fccff5370a2e259de330abc2898b9 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Thu, 8 Feb 2024 10:45:50 +0100 Subject: [PATCH 3/6] everything --- databricks/sdk/clock.py | 50 +++++++++++++++++++++++++++++++++++++++ databricks/sdk/core.py | 10 ++++++-- databricks/sdk/retries.py | 14 +++++++---- tests/clock.py | 14 +++++++++++ 4 files changed, 81 insertions(+), 7 deletions(-) diff --git a/databricks/sdk/clock.py b/databricks/sdk/clock.py index e69de29bb..f881a0b69 100644 --- a/databricks/sdk/clock.py +++ b/databricks/sdk/clock.py @@ -0,0 +1,50 @@ +import abc +import time + + +class Clock(metaclass=abc.ABCMeta): + @abc.abstractmethod + def time(self) -> float: + """ + Return the current time in seconds since the Epoch. + Fractions of a second may be present if the system clock provides them. + + :return: The current time in seconds since the Epoch. + """ + pass + + @abc.abstractmethod + def sleep(self, seconds: float) -> None: + """ + Return the current time in seconds since the Epoch. + Fractions of a second may be present if the system clock provides them. + + :param seconds: The duration to sleep in seconds. + :return: + """ + pass + + +class RealClock(Clock): + """ + A real clock that uses the ``time`` module to get the current time and sleep. + """ + + def time(self) -> float: + """ + Return the current time in seconds since the Epoch. + Fractions of a second may be present if the system clock provides them. + + :return: The current time in seconds since the Epoch. + """ + return time.time() + + def sleep(self, seconds: float) -> None: + """ + Return the current time in seconds since the Epoch. + Fractions of a second may be present if the system clock provides them. + + :param seconds: The duration to sleep in seconds. + :return: + """ + time.sleep(seconds) diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py index 2b7442708..2e0c8e0c9 100644 --- a/databricks/sdk/core.py +++ b/databricks/sdk/core.py @@ -12,6 +12,7 @@ from .credentials_provider import * from .errors import DatabricksError, error_mapper from .retries import retried +from .clock import Clock, RealClock __all__ = ['Config', 'DatabricksError'] @@ -22,12 +23,16 @@ class ApiClient: _cfg: Config _RETRY_AFTER_DEFAULT: int = 1 - def __init__(self, cfg: Config = None): + def __init__(self, cfg: Config = None, clock: Clock=None): if cfg is None: cfg = Config() + if clock is None: + clock = RealClock() + self._cfg = cfg + self._clock = clock # See https://github.com/databricks/databricks-sdk-go/blob/main/client/client.go#L34-L35 self._debug_truncate_bytes = cfg.debug_truncate_bytes if cfg.debug_truncate_bytes else 96 self._retry_timeout_seconds = cfg.retry_timeout_seconds if cfg.retry_timeout_seconds else 300 @@ -123,7 +128,8 @@ def do(self, headers = {} headers['User-Agent'] = self._user_agent_base retryable = retried(timeout=timedelta(seconds=self._retry_timeout_seconds), - is_retryable=self._is_retryable) + is_retryable=self._is_retryable, + clock=self._clock) return retryable(self._perform)(method, path, query=query, diff --git a/databricks/sdk/retries.py b/databricks/sdk/retries.py index a91467c4a..1ac8046de 100644 --- a/databricks/sdk/retries.py +++ b/databricks/sdk/retries.py @@ -1,30 +1,34 @@ import functools import logging -import time from datetime import timedelta from random import random from typing import Callable, Optional, Sequence, Type +from .clock import Clock, RealClock + logger = logging.getLogger(__name__) def retried(*, on: Sequence[Type[BaseException]] = None, is_retryable: Callable[[BaseException], Optional[str]] = None, - timeout=timedelta(minutes=20)): + timeout=timedelta(minutes=20), + clock: Clock=None): has_allowlist = on is not None has_callback = is_retryable is not None if not (has_allowlist or has_callback) or (has_allowlist and has_callback): raise SyntaxError('either on=[Exception] or callback=lambda x: .. is required') + if clock is None: + clock = RealClock() def decorator(func): @functools.wraps(func) def wrapper(*args, **kwargs): - deadline = time.time() + timeout.total_seconds() + deadline = clock.time() + timeout.total_seconds() attempt = 1 last_err = None - while time.time() < deadline: + while clock.time() < deadline: try: return func(*args, **kwargs) except Exception as err: @@ -50,7 +54,7 @@ def wrapper(*args, **kwargs): raise err logger.debug(f'Retrying: {retry_reason} (sleeping ~{sleep}s)') - time.sleep(sleep + random()) + clock.sleep(sleep + random()) attempt += 1 raise TimeoutError(f'Timed out after {timeout}') from last_err diff --git a/tests/clock.py b/tests/clock.py index e69de29bb..25ef969b4 100644 --- a/tests/clock.py +++ b/tests/clock.py @@ -0,0 +1,14 @@ +from databricks.sdk.clock import Clock + +class FakeClock(Clock): + """ + A simple clock that can be used to mock time in tests. + """ + def __init__(self, start_time: float = 0.0): + self._start_time = start_time + + def time(self) -> float: + return self._start_time + + def sleep(self, seconds: float) -> None: + self._start_time += seconds From a0a0df227933c719cdabc55985228bc8a12bf8fe Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Thu, 8 Feb 2024 10:46:54 +0100 Subject: [PATCH 4/6] extra change --- .gitattributes | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.gitattributes b/.gitattributes index 209ba7dec..d36f22eeb 100755 --- a/.gitattributes +++ b/.gitattributes @@ -190,7 +190,6 @@ examples/queries/create_queries.py linguist-generated=true examples/queries/get_queries.py linguist-generated=true examples/queries/update_queries.py linguist-generated=true examples/query_history/list_sql_query_history.py linguist-generated=true -examples/r/wait_catalog_workspace_bindings.py linguist-generated=true examples/recipients/create_recipients.py linguist-generated=true examples/recipients/get_recipients.py linguist-generated=true examples/recipients/list_recipients.py linguist-generated=true @@ -286,6 +285,7 @@ examples/workspace/list_workspace_integration.py linguist-generated=true examples/workspace_assignment/list_workspace_assignment_on_aws.py linguist-generated=true examples/workspace_assignment/update_workspace_assignment_on_aws.py linguist-generated=true examples/workspace_bindings/get_catalog_workspace_bindings.py linguist-generated=true +examples/workspace_bindings/update_catalog_workspace_bindings.py linguist-generated=true examples/workspace_conf/get_status_repos.py linguist-generated=true examples/workspaces/create_workspaces.py linguist-generated=true examples/workspaces/get_workspaces.py linguist-generated=true From d710eb3f20ec6812b1249f3cb2c09e2791ef15a6 Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Thu, 8 Feb 2024 10:51:56 +0100 Subject: [PATCH 5/6] remove clock changes --- databricks/sdk/clock.py | 50 --------------------------------------- databricks/sdk/core.py | 10 ++------ databricks/sdk/retries.py | 14 ++++------- tests/clock.py | 14 ----------- tests/test_core.py | 13 +++++----- 5 files changed, 13 insertions(+), 88 deletions(-) delete mode 100644 databricks/sdk/clock.py delete mode 100644 tests/clock.py diff --git a/databricks/sdk/clock.py b/databricks/sdk/clock.py deleted file mode 100644 index f881a0b69..000000000 --- a/databricks/sdk/clock.py +++ /dev/null @@ -1,50 +0,0 @@ -import abc -import time - - -class Clock(metaclass=abc.ABCMeta): - @abc.abstractmethod - def time(self) -> float: - """ - Return the current time in seconds since the Epoch. - Fractions of a second may be present if the system clock provides them. - - :return: The current time in seconds since the Epoch. - """ - pass - - @abc.abstractmethod - def sleep(self, seconds: float) -> None: - """ - Return the current time in seconds since the Epoch. - Fractions of a second may be present if the system clock provides them. - - :param seconds: The duration to sleep in seconds. - :return: - """ - pass - - -class RealClock(Clock): - """ - A real clock that uses the ``time`` module to get the current time and sleep. - """ - - def time(self) -> float: - """ - Return the current time in seconds since the Epoch. - Fractions of a second may be present if the system clock provides them. - - :return: The current time in seconds since the Epoch. - """ - return time.time() - - def sleep(self, seconds: float) -> None: - """ - Return the current time in seconds since the Epoch. - Fractions of a second may be present if the system clock provides them. - - :param seconds: The duration to sleep in seconds. - :return: - """ - time.sleep(seconds) diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py index 2e0c8e0c9..2b7442708 100644 --- a/databricks/sdk/core.py +++ b/databricks/sdk/core.py @@ -12,7 +12,6 @@ from .credentials_provider import * from .errors import DatabricksError, error_mapper from .retries import retried -from .clock import Clock, RealClock __all__ = ['Config', 'DatabricksError'] @@ -23,16 +22,12 @@ class ApiClient: _cfg: Config _RETRY_AFTER_DEFAULT: int = 1 - def __init__(self, cfg: Config = None, clock: Clock=None): + def __init__(self, cfg: Config = None): if cfg is None: cfg = Config() - if clock is None: - clock = RealClock() - self._cfg = cfg - self._clock = clock # See https://github.com/databricks/databricks-sdk-go/blob/main/client/client.go#L34-L35 self._debug_truncate_bytes = cfg.debug_truncate_bytes if cfg.debug_truncate_bytes else 96 self._retry_timeout_seconds = cfg.retry_timeout_seconds if cfg.retry_timeout_seconds else 300 @@ -128,8 +123,7 @@ def do(self, headers = {} headers['User-Agent'] = self._user_agent_base retryable = retried(timeout=timedelta(seconds=self._retry_timeout_seconds), - is_retryable=self._is_retryable, - clock=self._clock) + is_retryable=self._is_retryable) return retryable(self._perform)(method, path, query=query, diff --git a/databricks/sdk/retries.py b/databricks/sdk/retries.py index 1ac8046de..a91467c4a 100644 --- a/databricks/sdk/retries.py +++ b/databricks/sdk/retries.py @@ -1,34 +1,30 @@ import functools import logging +import time from datetime import timedelta from random import random from typing import Callable, Optional, Sequence, Type -from .clock import Clock, RealClock - logger = logging.getLogger(__name__) def retried(*, on: Sequence[Type[BaseException]] = None, is_retryable: Callable[[BaseException], Optional[str]] = None, - timeout=timedelta(minutes=20), - clock: Clock=None): + timeout=timedelta(minutes=20)): has_allowlist = on is not None has_callback = is_retryable is not None if not (has_allowlist or has_callback) or (has_allowlist and has_callback): raise SyntaxError('either on=[Exception] or callback=lambda x: .. is required') - if clock is None: - clock = RealClock() def decorator(func): @functools.wraps(func) def wrapper(*args, **kwargs): - deadline = clock.time() + timeout.total_seconds() + deadline = time.time() + timeout.total_seconds() attempt = 1 last_err = None - while clock.time() < deadline: + while time.time() < deadline: try: return func(*args, **kwargs) except Exception as err: @@ -54,7 +50,7 @@ def wrapper(*args, **kwargs): raise err logger.debug(f'Retrying: {retry_reason} (sleeping ~{sleep}s)') - clock.sleep(sleep + random()) + time.sleep(sleep + random()) attempt += 1 raise TimeoutError(f'Timed out after {timeout}') from last_err diff --git a/tests/clock.py b/tests/clock.py deleted file mode 100644 index 25ef969b4..000000000 --- a/tests/clock.py +++ /dev/null @@ -1,14 +0,0 @@ -from databricks.sdk.clock import Clock - -class FakeClock(Clock): - """ - A simple clock that can be used to mock time in tests. - """ - def __init__(self, start_time: float = 0.0): - self._start_time = start_time - - def time(self) -> float: - return self._start_time - - def sleep(self, seconds: float) -> None: - self._start_time += seconds diff --git a/tests/test_core.py b/tests/test_core.py index 4935a18f6..ca2eaac31 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -27,7 +27,6 @@ from databricks.sdk.version import __version__ from .conftest import noop_credentials -from .clock import FakeClock def test_parse_dsn(): @@ -423,7 +422,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock()) + api_client = ApiClient(Config(host=host, token='_')) res = api_client.do('GET', '/foo') assert 'foo' in res @@ -446,7 +445,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock()) + api_client = ApiClient(Config(host=host, token='_')) res = api_client.do('GET', '/foo') assert 'foo' in res @@ -463,7 +462,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_', retry_timeout_seconds=1), clock=FakeClock()) + api_client = ApiClient(Config(host=host, token='_', retry_timeout_seconds=1)) with pytest.raises(TimeoutError): api_client.do('GET', '/foo') @@ -485,7 +484,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock()) + api_client = ApiClient(Config(host=host, token='_')) res = api_client.do('GET', '/foo') assert 'foo' in res @@ -503,7 +502,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock()) + api_client = ApiClient(Config(host=host, token='_')) with pytest.raises(DatabricksError): api_client.do('GET', '/foo') @@ -521,7 +520,7 @@ def inner(h: BaseHTTPRequestHandler): requests.append(h.requestline) with http_fixture_server(inner) as host: - api_client = ApiClient(Config(host=host, token='_'), clock=FakeClock()) + api_client = ApiClient(Config(host=host, token='_')) res = api_client.do('GET', '/foo') assert 'foo' in res From 3685fdee63c91d90571d1c2f473c2595fdc29ecc Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Thu, 8 Feb 2024 11:36:51 +0100 Subject: [PATCH 6/6] remove errant example --- examples/r/wait_catalog_workspace_bindings.py | 5 ----- 1 file changed, 5 deletions(-) delete mode 100755 examples/r/wait_catalog_workspace_bindings.py diff --git a/examples/r/wait_catalog_workspace_bindings.py b/examples/r/wait_catalog_workspace_bindings.py deleted file mode 100755 index 1352ed169..000000000 --- a/examples/r/wait_catalog_workspace_bindings.py +++ /dev/null @@ -1,5 +0,0 @@ -from databricks.sdk import WorkspaceClient - -w = WorkspaceClient() - -w.r.wait(update_function)