diff --git a/.gitignore b/.gitignore index d5e32c90..aaa3c355 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,4 @@ venv .DS_Store .coverage .tox/ +.vscode/ diff --git a/databricks_cli/accounts/__init__.py b/databricks_cli/accounts/__init__.py new file mode 100644 index 00000000..b0c9feac --- /dev/null +++ b/databricks_cli/accounts/__init__.py @@ -0,0 +1,22 @@ +# Databricks CLI +# Copyright 2017 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/databricks_cli/accounts/api.py b/databricks_cli/accounts/api.py new file mode 100644 index 00000000..e6273487 --- /dev/null +++ b/databricks_cli/accounts/api.py @@ -0,0 +1,130 @@ +"""Implement Databricks Accounts API, interfacing with the AccountsService.""" +# Databricks CLI +# Copyright 2017 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from databricks_cli.sdk import AccountsService + + +class AccountsApi(object): + """Implement the databricks '2.0/accounts' API Interface.""" + + def __init__(self, api_client): + self.client = AccountsService(api_client) + + def create_credentials(self, account_id, json): + """Create a new credentials object for the AWS IAM role reference.""" + if account_id is None: + raise TypeError('Expected account_id') + endpoint = "/accounts/%s/credentials" % (account_id) + return self.client.client.perform_query('POST', endpoint, data=json) + + def get_credentials(self, account_id, credentials_id): + """Get the credentials object for the given credentials id.""" + return self.client.get_credentials(account_id, credentials_id) + + def list_credentials(self, account_id): + """Get all credentials objects for the given account.""" + return self.client.list_credentials(account_id) + + def delete_credentials(self, account_id, credentials_id): + """Delete the credentials object for the given credentials id.""" + return self.client.delete_credentials(account_id, credentials_id) + + def create_storage_config(self, account_id, json): + """Create a new storage config object for the AWS bucket reference.""" + if account_id is None: + raise TypeError('Expected account_id') + endpoint = "/accounts/%s/storage-configurations" % (account_id) + return self.client.client.perform_query('POST', endpoint, data=json) + + def get_storage_config(self, account_id, storage_config_id): + """Get the storage config object for the given storage config id.""" + return self.client.get_storage_config(account_id, storage_config_id) + + def list_storage_configs(self, account_id): + """Get all storage config objects for the given account.""" + return self.client.list_storage_configs(account_id) + + def delete_storage_config(self, account_id, storage_config_id): + """Delete the storage config object for the given storage config id.""" + return self.client.delete_storage_config(account_id, storage_config_id) + + def create_network(self, account_id, json): + """Create a new network object for the AWS network infrastructure reference.""" + if account_id is None: + raise TypeError('Expected account_id') + endpoint = "/accounts/%s/networks" % (account_id) + return self.client.client.perform_query('POST', endpoint, data=json) + + def get_network(self, account_id, network_id): + """Get the network object for the given network id.""" + return self.client.get_network(account_id, network_id) + + def list_networks(self, account_id): + """Get all network objects for the given account.""" + return self.client.list_networks(account_id) + + def delete_network(self, account_id, network_id): + """Delete the network object for the given network id.""" + return self.client.delete_network(account_id, network_id) + + def create_customer_managed_key(self, account_id, json): + """Create a new customer managed key object for the AWS KMS key reference.""" + if account_id is None: + raise TypeError('Expected account_id') + endpoint = "/accounts/%s/customer-managed-keys" % (account_id) + return self.client.client.perform_query('POST', endpoint, data=json) + + def get_customer_managed_key(self, account_id, customer_managed_key_id): + """Get the customer managed key object for the given customer managed key id.""" + return self.client.get_customer_managed_key(account_id, customer_managed_key_id) + + def list_customer_managed_keys(self, account_id): + """Get all customer managed key objects for the given account.""" + return self.client.list_customer_managed_keys(account_id) + + def create_workspace(self, account_id, json): + """Create a new workspace with the required references.""" + if account_id is None: + raise TypeError('Expected account_id') + endpoint = "/accounts/%s/workspaces" % (account_id) + return self.client.client.perform_query('POST', endpoint, data=json) + + def get_workspace(self, account_id, workspace_id): + """Get the workspace details for the given workspace id.""" + return self.client.get_workspace(account_id, workspace_id) + + def list_workspaces(self, account_id): + """Get all workspaces for the given account.""" + return self.client.list_workspaces(account_id) + + def delete_workspace(self, account_id, workspace_id): + """Delete the workspace for the given workspace id.""" + return self.client.delete_workspace(account_id, workspace_id) + + def list_customer_managed_key_hist_by_workspace(self, account_id, workspace_id): + """Get the history of customer managed key objects for the given workspace id.""" + return self.client.list_customer_managed_key_hist_by_workspace(account_id, workspace_id) + + def list_customer_managed_key_hist_by_account(self, account_id): + """Get the history of customer managed key objects for the given account.""" + return self.client.list_customer_managed_key_hist_by_account(account_id) diff --git a/databricks_cli/accounts/cli.py b/databricks_cli/accounts/cli.py new file mode 100644 index 00000000..a9832227 --- /dev/null +++ b/databricks_cli/accounts/cli.py @@ -0,0 +1,412 @@ +"""Provide the API methods for Databricks Accounts REST endpoint.""" +# Databricks CLI +# Copyright 2017 Databricks, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"), except +# that the use of services to which certain application programming +# interfaces (each, an "API") connect requires that the user first obtain +# a license for the use of the APIs from Databricks, Inc. ("Databricks"), +# by creating an account at www.databricks.com and agreeing to either (a) +# the Community Edition Terms of Service, (b) the Databricks Terms of +# Service, or (c) another written agreement between Licensee and Databricks +# for the use of the APIs. +# +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import click + +from databricks_cli.click_types import AccountIdClickType, CredentialsIdClickType, \ + StorageConfigIdClickType, NetworkIdClickType, CustomerManagedKeyIdClickType, \ + WorkspaceIdClickType, JsonClickType +from databricks_cli.accounts.api import AccountsApi +from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, pretty_format, json_cli_base +from databricks_cli.configure.config import provide_api_client, profile_option, debug_option +from databricks_cli.version import print_version_callback, version + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Create a credentials object for the AWS IAM role reference.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option('--json-file', default=None, type=click.Path(), + help='File containing JSON request to POST to \ + /api/2.0/accounts/{account_id}/credentials.') +@click.option('--json', default=None, type=JsonClickType(), + help=JsonClickType.help('/api/2.0/accounts/{account_id}/credentials')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_credentials_cli(api_client, account_id, json_file, json): + """Create a credentials object for the AWS IAM role reference.""" + json_cli_base(json_file, json, + lambda json: AccountsApi(api_client).create_credentials(account_id, json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get the credentials object for the given credentials id.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option("--credentials-id", required=True, type=CredentialsIdClickType(), + help=CredentialsIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_credentials_cli(api_client, account_id, credentials_id): + """Get the credentials object for the given credentials id.""" + content = AccountsApi(api_client).get_credentials(account_id, credentials_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get all credentials objects for the given account.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_credentials_cli(api_client, account_id): + """Get all credentials objects for the given account.""" + content = AccountsApi(api_client).list_credentials(account_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Delete the credentials object for the given credentials id.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option("--credentials-id", required=True, type=CredentialsIdClickType(), + help=CredentialsIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_credentials_cli(api_client, account_id, credentials_id): + """Delete the credentials object for the given credentials id.""" + content = AccountsApi(api_client).delete_credentials(account_id, credentials_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Create a storage config object for the AWS bucket reference.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option('--json-file', default=None, type=click.Path(), + help='File containing JSON request to POST to \ + /api/2.0/accounts/{account_id}/storage-configurations.') +@click.option('--json', default=None, type=JsonClickType(), + help=JsonClickType.help('/api/2.0/accounts/{account_id}/storage-configurations')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_storage_config_cli(api_client, account_id, json_file, json): + """Create a storage config object for the AWS bucket reference.""" + json_cli_base(json_file, json, + lambda json: AccountsApi(api_client).create_storage_config(account_id, json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get the storage config object for the given storage config id.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option("--storage-config-id", required=True, type=StorageConfigIdClickType(), + help=StorageConfigIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_storage_config_cli(api_client, account_id, storage_config_id): + """Get the storage config object for the given storage config id.""" + content = AccountsApi(api_client).get_storage_config(account_id, storage_config_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get all storage config objects for the given account.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_storage_configs_cli(api_client, account_id): + """Get all storage config objects for the given account.""" + content = AccountsApi(api_client).list_storage_configs(account_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Delete the storage config object for the given storage config id.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option("--storage-config-id", required=True, type=StorageConfigIdClickType(), + help=StorageConfigIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_storage_config_cli(api_client, account_id, storage_config_id): + """Delete the storage config object for the given storage config id.""" + content = AccountsApi(api_client).delete_storage_config(account_id, storage_config_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Create a network object for the AWS network infrastructure reference.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option('--json-file', default=None, type=click.Path(), + help='File containing JSON request to POST to \ + /api/2.0/accounts/{account_id}/networks.') +@click.option('--json', default=None, type=JsonClickType(), + help=JsonClickType.help('/api/2.0/accounts/{account_id}/networks')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_network_cli(api_client, account_id, json_file, json): + """Create a network object for the AWS network infrastructure reference.""" + json_cli_base(json_file, json, + lambda json: AccountsApi(api_client).create_network(account_id, json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get the network object for the given network id.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option("--network-id", required=True, type=NetworkIdClickType(), + help=NetworkIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_network_cli(api_client, account_id, network_id): + """Get the network object for the given network id.""" + content = AccountsApi(api_client).get_network(account_id, network_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get all network objects for the given account.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_networks_cli(api_client, account_id): + """Get all network objects for the given account.""" + content = AccountsApi(api_client).list_networks(account_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Delete the network object for the given network id.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option("--network-id", required=True, type=NetworkIdClickType(), + help=NetworkIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_network_cli(api_client, account_id, network_id): + """Delete the network object for the given network id.""" + content = AccountsApi(api_client).delete_network(account_id, network_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Create a customer managed key object for the AWS KMS key reference.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option('--json-file', default=None, type=click.Path(), + help='File containing JSON request to POST to \ + /api/2.0/accounts/{account_id}/customer-managed-keys.') +@click.option('--json', default=None, type=JsonClickType(), + help=JsonClickType.help('/api/2.0/accounts/{account_id}/customer-managed-keys')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_customer_managed_key_cli(api_client, account_id, json_file, json): + """Create a customer managed key object for the AWS KMS key reference.""" + json_cli_base(json_file, json, + lambda json: AccountsApi(api_client). + create_customer_managed_key(account_id, json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get customer managed key object for the given customer managed key id.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option("--customer-managed-key-id", required=True, type=CustomerManagedKeyIdClickType(), + help=CustomerManagedKeyIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_customer_managed_key_cli(api_client, account_id, customer_managed_key_id): + """Get the customer managed key object for the given customer managed key id.""" + content = AccountsApi(api_client).get_customer_managed_key(account_id, customer_managed_key_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get all customer managed key objects for the given account.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_customer_managed_keys_cli(api_client, account_id): + """Get all customer managed key objects for the given account.""" + content = AccountsApi(api_client).list_customer_managed_keys(account_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Create a workspace with the required references.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option('--json-file', default=None, type=click.Path(), + help='File containing JSON request to POST to \ + /api/2.0/accounts/{account_id}/workspaces.') +@click.option('--json', default=None, type=JsonClickType(), + help=JsonClickType.help('/api/2.0/accounts/{account_id}/workspaces')) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def create_workspace_cli(api_client, account_id, json_file, json): + """Create a workspace with the required references.""" + json_cli_base(json_file, json, + lambda json: AccountsApi(api_client).create_workspace(account_id, json)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get the workspace details for the given workspace id.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option("--workspace-id", required=True, type=WorkspaceIdClickType(), + help=WorkspaceIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def get_workspace_cli(api_client, account_id, workspace_id): + """Get the workspace details for the given workspace id.""" + content = AccountsApi(api_client).get_workspace(account_id, workspace_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get all workspaces for the given account.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_workspaces_cli(api_client, account_id): + """Get all workspaces for the given account.""" + content = AccountsApi(api_client).list_workspaces(account_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Delete the workspace for the given workspace id.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option("--workspace-id", required=True, type=WorkspaceIdClickType(), + help=WorkspaceIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def delete_workspace_cli(api_client, account_id, workspace_id): + """Delete the workspace for the given workspace id.""" + content = AccountsApi(api_client).delete_workspace(account_id, workspace_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get history of customer managed key objects for the given workspace id.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@click.option("--workspace-id", required=True, type=WorkspaceIdClickType(), + help=WorkspaceIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_customer_managed_key_hist_by_workspace_cli(api_client, account_id, workspace_id): + """Get the history of customer managed key objects for the given workspace id.""" + content = AccountsApi(api_client). \ + list_customer_managed_key_hist_by_workspace(account_id, workspace_id) + click.echo(pretty_format(content)) + + +@click.command(context_settings=CONTEXT_SETTINGS, + short_help="Get the history of customer managed key objects for the given account.") +@click.option('--account-id', required=True, type=AccountIdClickType(), + help=AccountIdClickType.help) +@debug_option +@profile_option +@eat_exceptions +@provide_api_client +def list_customer_managed_key_hist_by_account_cli(api_client, account_id): + """Get the history of customer managed key objects for the given account.""" + content = AccountsApi(api_client).list_customer_managed_key_hist_by_account(account_id) + click.echo(pretty_format(content)) + + +@click.group(context_settings=CONTEXT_SETTINGS, + short_help='Utility to interact with Databricks accounts.') +@click.option('--version', '-v', is_flag=True, callback=print_version_callback, + expose_value=False, is_eager=True, help=version) +@debug_option +@profile_option +@eat_exceptions +def accounts_group(): + """ + Utility to interact with Databricks accounts. + """ + + +accounts_group.add_command(create_credentials_cli, name="create-credentials") +accounts_group.add_command(get_credentials_cli, name="get-credentials") +accounts_group.add_command(list_credentials_cli, name="list-credentials") +accounts_group.add_command(delete_credentials_cli, name="delete-credentials") +accounts_group.add_command(create_storage_config_cli, name="create-storage-config") +accounts_group.add_command(get_storage_config_cli, name="get-storage-config") +accounts_group.add_command(list_storage_configs_cli, name="list-storage-config") +accounts_group.add_command(delete_storage_config_cli, name="delete-storage-config") +accounts_group.add_command(create_network_cli, name="create-network") +accounts_group.add_command(get_network_cli, name="get-network") +accounts_group.add_command(list_networks_cli, name="list-network") +accounts_group.add_command(delete_network_cli, name="delete-network") +accounts_group.add_command(create_customer_managed_key_cli, name="create-cust-managed-key") +accounts_group.add_command(get_customer_managed_key_cli, name="get-cust-managed-key") +accounts_group.add_command(list_customer_managed_keys_cli, name="list-cust-managed-key") +accounts_group.add_command(create_workspace_cli, name="create-workspace") +accounts_group.add_command(get_workspace_cli, name="get-workspace") +accounts_group.add_command(list_workspaces_cli, name="list-workspace") +accounts_group.add_command(delete_workspace_cli, name="delete-workspace") +accounts_group.add_command(list_customer_managed_key_hist_by_workspace_cli, + name="list-cust-managed-key-hist-by-ws") +accounts_group.add_command(list_customer_managed_key_hist_by_account_cli, + name="list-cust-managed-key-hist-by-acc") diff --git a/databricks_cli/cli.py b/databricks_cli/cli.py index b05a703b..c1addde6 100644 --- a/databricks_cli/cli.py +++ b/databricks_cli/cli.py @@ -37,7 +37,11 @@ from databricks_cli.stack.cli import stack_group from databricks_cli.groups.cli import groups_group from databricks_cli.instance_pools.cli import instance_pools_group +<<<<<<< HEAD +from databricks_cli.accounts.cli import accounts_group +======= from databricks_cli.pipelines.cli import pipelines_group +>>>>>>> upstream/master @click.group(context_settings=CONTEXT_SETTINGS) @@ -60,6 +64,7 @@ def cli(): cli.add_command(stack_group, name='stack') cli.add_command(groups_group, name='groups') cli.add_command(instance_pools_group, name="instance-pools") +cli.add_command(accounts_group, name="accounts") cli.add_command(pipelines_group, name='pipelines') if __name__ == "__main__": diff --git a/databricks_cli/click_types.py b/databricks_cli/click_types.py index 18852792..d1e8087f 100644 --- a/databricks_cli/click_types.py +++ b/databricks_cli/click_types.py @@ -90,6 +90,36 @@ class SecretPrincipalClickType(ParamType): help = 'The name of the principal.' +class AccountIdClickType(ParamType): + name = 'ACCOUNT_ID' + help = 'Databricks Account Id' + + +class CredentialsIdClickType(ParamType): + name = 'CREDENTIALS_ID' + help = 'Databricks Workspace Credentials Id' + + +class StorageConfigIdClickType(ParamType): + name = 'STORAGE_CONFIG_ID' + help = 'Databricks Workspace Storage Configuration Id' + + +class NetworkIdClickType(ParamType): + name = 'NETWORK_ID' + help = 'Databricks Workspace Network Id' + + +class CustomerManagedKeyIdClickType(ParamType): + name = 'CUSTOMER_MANAGED_KEY_ID' + help = 'Databricks Workspace Customer Managed Key Id' + + +class WorkspaceIdClickType(ParamType): + name = 'WORKSPACE_ID' + help = 'Databricks Workspace Id' + + class PipelineSpecClickType(ParamType): name = 'SPEC' help = 'The path to the pipelines deployment spec file' diff --git a/databricks_cli/sdk/api_client.py b/databricks_cli/sdk/api_client.py index d5870a6a..f504b543 100644 --- a/databricks_cli/sdk/api_client.py +++ b/databricks_cli/sdk/api_client.py @@ -93,6 +93,10 @@ def __init__(self, user=None, password=None, host=None, token=None, self.default_headers.update(user_agent) self.verify = verify + """Adding for Accounts API support""" + self.restful_methods = ['REST-GET', 'DELETE'] + self.accounts_url = "%s://%s/api/%s" % ("https", "accounts.cloud.databricks.com", apiVersion) + def close(self): """Close the client""" pass @@ -100,6 +104,13 @@ def close(self): # helper functions starting here def perform_query(self, method, path, data = {}, headers = None): + """Adding for Accounts API support""" + path_specific_url = "" + if path.startswith("/accounts"): + path_specific_url = self.accounts_url + else: + path_specific_url = self.url + """set up connection and perform query""" if headers is None: headers = self.default_headers @@ -110,12 +121,18 @@ def perform_query(self, method, path, data = {}, headers = None): with warnings.catch_warnings(): warnings.simplefilter("ignore", exceptions.InsecureRequestWarning) - if method == 'GET': + if method in self.restful_methods: + restful_method = method + if method == 'REST-GET': + restful_method = 'GET' + resp = self.session.request(restful_method, path_specific_url + path, verify = self.verify, + headers = headers) + elif method == 'GET': translated_data = {k: _translate_boolean_to_query_param(data[k]) for k in data} - resp = self.session.request(method, self.url + path, params = translated_data, + resp = self.session.request(method, path_specific_url + path, params = translated_data, verify = self.verify, headers = headers) else: - resp = self.session.request(method, self.url + path, data = json.dumps(data), + resp = self.session.request(method, path_specific_url + path, data = json.dumps(data), verify = self.verify, headers = headers) try: resp.raise_for_status() diff --git a/databricks_cli/sdk/service.py b/databricks_cli/sdk/service.py index 532fab01..73a22105 100644 --- a/databricks_cli/sdk/service.py +++ b/databricks_cli/sdk/service.py @@ -796,6 +796,107 @@ def list_instance_pools(self, headers=None): return self.client.perform_query('GET', '/instance-pools/list', data=_data, headers=headers) +class AccountsService(object): + def __init__(self, client): + self.client = client + + def get_credentials(self, account_id, credentials_id, headers=None): + if account_id is None or credentials_id is None: + raise TypeError('Expected account_id and credentials_id') + endpoint = "/accounts/%s/credentials/%s" % (account_id, credentials_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + def list_credentials(self, account_id, headers=None): + if account_id is None: + raise TypeError('Expected account_id') + endpoint = "/accounts/%s/credentials" % (account_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + def delete_credentials(self, account_id, credentials_id, headers=None): + if account_id is None or credentials_id is None: + raise TypeError('Expected account_id and credentials_id') + endpoint = "/accounts/%s/credentials/%s" % (account_id, credentials_id) + return self.client.perform_query('DELETE', endpoint, data={}, headers=headers) + + def get_storage_config(self, account_id, storage_config_id, headers=None): + if account_id is None or storage_config_id is None: + raise TypeError('Expected account_id and storage_config_id') + endpoint = "/accounts/%s/storage-configurations/%s" % (account_id, storage_config_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + def list_storage_configs(self, account_id, headers=None): + if account_id is None: + raise TypeError('Expected account_id') + endpoint = "/accounts/%s/storage-configurations" % (account_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + def delete_storage_config(self, account_id, storage_config_id, headers=None): + if account_id is None or storage_config_id is None: + raise TypeError('Expected account_id and storage_config_id') + endpoint = "/accounts/%s/storage-configurations/%s" % (account_id, storage_config_id) + return self.client.perform_query('DELETE', endpoint, data={}, headers=headers) + + def get_network(self, account_id, network_id, headers=None): + if account_id is None or network_id is None: + raise TypeError('Expected account_id and network_id') + endpoint = "/accounts/%s/networks/%s" % (account_id, network_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + def list_networks(self, account_id, headers=None): + if account_id is None: + raise TypeError('Expected account_id') + endpoint = "/accounts/%s/networks" % (account_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + def delete_network(self, account_id, network_id, headers=None): + if account_id is None or network_id is None: + raise TypeError('Expected account_id and network_id') + endpoint = "/accounts/%s/networks/%s" % (account_id, network_id) + return self.client.perform_query('DELETE', endpoint, data={}, headers=headers) + + def get_customer_managed_key(self, account_id, customer_managed_key_id, headers=None): + if account_id is None or customer_managed_key_id is None: + raise TypeError('Expected account_id and customer_managed_key_id') + endpoint = "/accounts/%s/customer-managed-keys/%s" % (account_id, customer_managed_key_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + def list_customer_managed_keys(self, account_id, headers=None): + if account_id is None: + raise TypeError('Expected account_id') + endpoint = "/accounts/%s/customer-managed-keys" % (account_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + def get_workspace(self, account_id, workspace_id, headers=None): + if account_id is None or workspace_id is None: + raise TypeError('Expected account_id and workspace_id') + endpoint = "/accounts/%s/workspaces/%s" % (account_id, workspace_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + def list_workspaces(self, account_id, headers=None): + if account_id is None: + raise TypeError('Expected account_id') + endpoint = "/accounts/%s/workspaces" % (account_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + def delete_workspace(self, account_id, workspace_id, headers=None): + if account_id is None or workspace_id is None: + raise TypeError('Expected account_id and workspace_id') + endpoint = "/accounts/%s/workspaces/%s" % (account_id, workspace_id) + return self.client.perform_query('DELETE', endpoint, data={}, headers=headers) + + def list_customer_managed_key_hist_by_workspace(self, account_id, workspace_id, headers=None): + if account_id is None or workspace_id is None: + raise TypeError('Expected account_id and workspace_id') + endpoint = "/accounts/%s/workspaces/%s/customer-managed-key-history" % (account_id, workspace_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + def list_customer_managed_key_hist_by_account(self, account_id, headers=None): + if account_id is None: + raise TypeError('Expected account_id') + endpoint = "/accounts/%s/customer-managed-key-history" % (account_id) + return self.client.perform_query('REST-GET', endpoint, data={}, headers=headers) + + class DeltaPipelinesService(object): def __init__(self, client): self.client = client