1# --------------------------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License. See License.txt in the project root for license information.
4# --------------------------------------------------------------------------------------------
5
6# pylint: disable=line-too-long
7# pylint: disable=too-many-statements
8
9from azure.cli.core.commands import CliCommandType
10
11from azure.cli.command_modules.cosmosdb._client_factory import (
12    cf_db_accounts,
13    cf_db_private_endpoint_connections,
14    cf_db_private_link_resources,
15    cf_sql_resources,
16    cf_mongo_db_resources,
17    cf_cassandra_resources,
18    cf_gremlin_resources,
19    cf_table_resources,
20    cf_restorable_database_accounts,
21    cf_restorable_sql_databases,
22    cf_restorable_sql_containers,
23    cf_restorable_sql_resources,
24    cf_restorable_mongodb_databases,
25    cf_restorable_mongodb_collections,
26    cf_restorable_mongodb_resources
27)
28
29from azure.cli.command_modules.cosmosdb._format import (
30    database_output,
31    list_database_output,
32    collection_output,
33    list_collection_output,
34    list_connection_strings_output
35)
36
37from azure.cli.command_modules.cosmosdb._transformers import (
38    transform_network_rule_list_output,
39    transform_db_account_json_output,
40    transform_db_account_list_output
41)
42
43from ._validators import (
44    validate_private_endpoint_connection_id
45)
46
47DATABASE_DEPRECATION_INFO = 'cosmosdb sql database, cosmosdb mongodb database, cosmosdb cassandra keyspace or cosmosdb gremlin database'
48
49COLLECTION_DEPRECATON_INFO = 'cosmosdb sql container, cosmosdb mongodb collection, cosmosdb cassandra table, cosmosdb gremlin graph or cosmosdb table'
50
51
52def load_command_table(self, _):
53
54    cosmosdb_sdk = CliCommandType(
55        operations_tmpl='azure.mgmt.cosmosdb.operations#DatabaseAccountsOperations.{}',
56        client_factory=cf_db_accounts)
57
58    cosmosdb_private_endpoint_connections_sdk = CliCommandType(
59        operations_tmpl='azure.mgmt.cosmosdb.operations#PrivateEndpointConnectionsOperations.{}',
60        client_factory=cf_db_private_endpoint_connections)
61
62    cosmosdb_private_link_resources_sdk = CliCommandType(
63        operations_tmpl='azure.mgmt.cosmosdb.operations#PrivateLinkResourcesOperations.{}',
64        client_factory=cf_db_private_link_resources)
65
66    cosmosdb_sql_sdk = CliCommandType(
67        operations_tmpl='azure.mgmt.cosmosdb.operations#SqlResourcesOperations.{}',
68        client_factory=cf_sql_resources)
69
70    cosmosdb_mongo_sdk = CliCommandType(
71        operations_tmpl='azure.mgmt.cosmosdb.operations#MongoDBResourcesOperations.{}',
72        client_factory=cf_mongo_db_resources)
73
74    cosmosdb_cassandra_sdk = CliCommandType(
75        operations_tmpl='azure.mgmt.cosmosdb.operations#CassandraResourcesOperations.{}',
76        client_factory=cf_cassandra_resources)
77
78    cosmosdb_gremlin_sdk = CliCommandType(
79        operations_tmpl='azure.mgmt.cosmosdb.operations#GremlinResourcesOperations.{}',
80        client_factory=cf_gremlin_resources)
81
82    cosmosdb_table_sdk = CliCommandType(
83        operations_tmpl='azure.mgmt.cosmosdb.operations#TableResourcesOperations.{}',
84        client_factory=cf_table_resources)
85
86    cosmosdb_restorable_database_accounts_sdk = CliCommandType(
87        operations_tmpl='azure.mgmt.cosmosdb.operations#RestorableDatabaseAccountsOperations.{}',
88        client_factory=cf_restorable_database_accounts)
89
90    cosmosdb_sql_restorable_database_sdk = CliCommandType(
91        operations_tmpl='azure.mgmt.cosmosdb.operations#RestorableSqlDatabasesOperations.{}',
92        client_factory=cf_restorable_sql_databases)
93
94    cosmosdb_sql_restorable_container_sdk = CliCommandType(
95        operations_tmpl='azure.mgmt.cosmosdb.operations#RestorableSqlContainersOperations.{}',
96        client_factory=cf_restorable_sql_containers)
97
98    cosmosdb_sql_restorable_resources_sdk = CliCommandType(
99        operations_tmpl='azure.mgmt.cosmosdb.operations#RestorableSqlResourcesOperations.{}',
100        client_factory=cf_restorable_sql_resources)
101
102    cosmosdb_mongodb_restorable_database_sdk = CliCommandType(
103        operations_tmpl='azure.mgmt.cosmosdb.operations#RestorableMongodbDatabasesOperations.{}',
104        client_factory=cf_restorable_mongodb_databases)
105
106    cosmosdb_mongodb_restorable_collection_sdk = CliCommandType(
107        operations_tmpl='azure.mgmt.cosmosdb.operations#RestorableMongodbCollectionsOperations.{}',
108        client_factory=cf_restorable_mongodb_collections)
109
110    cosmosdb_mongodb_restorable_resources_sdk = CliCommandType(
111        operations_tmpl='azure.mgmt.cosmosdb.operations#RestorableMongodbResourcesOperations.{}',
112        client_factory=cf_restorable_mongodb_resources)
113
114    with self.command_group('cosmosdb', cosmosdb_sdk, client_factory=cf_db_accounts) as g:
115        g.show_command('show', 'get', transform=transform_db_account_json_output)
116        g.command('list-keys', 'list_keys', deprecate_info=g.deprecate(redirect='cosmosdb keys list', hide=True))
117        g.command('list-read-only-keys', 'list_read_only_keys', deprecate_info=g.deprecate(redirect='cosmosdb keys list --type read-only-keys', hide=True))
118        g.command('list-connection-strings', 'list_connection_strings', table_transformer=list_connection_strings_output, deprecate_info=g.deprecate(redirect='cosmosdb keys list --type connection-strings', hide=True))
119        g.custom_command('regenerate-key', 'cli_cosmosdb_regenerate_key', deprecate_info=g.deprecate(redirect='cosmosdb keys regenerate', hide=True))
120        g.command('check-name-exists', 'check_name_exists')
121        g.command('delete', 'begin_delete', confirmation=True)
122        g.command('failover-priority-change', 'begin_failover_priority_change')
123        g.custom_command('create', 'cli_cosmosdb_create', transform=transform_db_account_json_output)
124        g.custom_command('update', 'cli_cosmosdb_update', transform=transform_db_account_json_output)
125        g.custom_command('list', 'cli_cosmosdb_list', transform=transform_db_account_list_output)
126        g.custom_command('restore', 'cli_cosmosdb_restore', transform=transform_db_account_json_output)
127
128    with self.command_group('cosmosdb private-endpoint-connection',
129                            cosmosdb_private_endpoint_connections_sdk,
130                            client_factory=cf_db_private_endpoint_connections) as g:
131        g.custom_command('approve', 'approve_private_endpoint_connection',
132                         validator=validate_private_endpoint_connection_id)
133        g.custom_command('reject', 'reject_private_endpoint_connection',
134                         validator=validate_private_endpoint_connection_id)
135        g.command('delete', 'begin_delete', validator=validate_private_endpoint_connection_id)
136        g.show_command('show', 'get', validator=validate_private_endpoint_connection_id)
137
138    with self.command_group('cosmosdb private-link-resource',
139                            cosmosdb_private_link_resources_sdk,
140                            client_factory=cf_db_private_link_resources) as g:
141        from azure.cli.core.commands.transform import gen_dict_to_list_transform
142        g.show_command('list', 'list_by_database_account', transform=gen_dict_to_list_transform(key='values'))
143
144    # SQL api
145    with self.command_group('cosmosdb sql'):
146        pass
147    with self.command_group('cosmosdb sql database', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
148        g.custom_command('create', 'cli_cosmosdb_sql_database_create')
149        g.custom_command('exists', 'cli_cosmosdb_sql_database_exists')
150        g.command('list', 'list_sql_databases')
151        g.show_command('show', 'get_sql_database')
152        g.command('delete', 'begin_delete_sql_database', confirmation=True)
153
154    with self.command_group('cosmosdb sql container', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
155        g.custom_command('create', 'cli_cosmosdb_sql_container_create')
156        g.custom_command('update', 'cli_cosmosdb_sql_container_update')
157        g.custom_command('exists', 'cli_cosmosdb_sql_container_exists')
158        g.command('list', 'list_sql_containers')
159        g.show_command('show', 'get_sql_container')
160        g.command('delete', 'begin_delete_sql_container', confirmation=True)
161
162    with self.command_group('cosmosdb sql stored-procedure', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
163        g.custom_command('create', 'cli_cosmosdb_sql_stored_procedure_create_update')
164        g.custom_command('update', 'cli_cosmosdb_sql_stored_procedure_create_update')
165        g.command('list', 'list_sql_stored_procedures')
166        g.show_command('show', 'get_sql_stored_procedure')
167        g.command('delete', 'begin_delete_sql_stored_procedure', confirmation=True)
168
169    with self.command_group('cosmosdb sql trigger', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
170        g.custom_command('create', 'cli_cosmosdb_sql_trigger_create')
171        g.custom_command('update', 'cli_cosmosdb_sql_trigger_update')
172        g.command('list', 'list_sql_triggers')
173        g.show_command('show', 'get_sql_trigger')
174        g.command('delete', 'begin_delete_sql_trigger', confirmation=True)
175
176    with self.command_group('cosmosdb sql user-defined-function', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
177        g.custom_command('create', 'cli_cosmosdb_sql_user_defined_function_create_update')
178        g.custom_command('update', 'cli_cosmosdb_sql_user_defined_function_create_update')
179        g.command('list', 'list_sql_user_defined_functions')
180        g.show_command('show', 'get_sql_user_defined_function')
181        g.command('delete', 'begin_delete_sql_user_defined_function', confirmation=True)
182
183    # MongoDB api
184    with self.command_group('cosmosdb mongodb'):
185        pass
186    with self.command_group('cosmosdb mongodb database', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
187        g.custom_command('create', 'cli_cosmosdb_mongodb_database_create')
188        g.custom_command('exists', 'cli_cosmosdb_mongodb_database_exists')
189        g.command('list', 'list_mongo_db_databases')
190        g.show_command('show', 'get_mongo_db_database')
191        g.command('delete', 'begin_delete_mongo_db_database', confirmation=True)
192
193    with self.command_group('cosmosdb mongodb collection', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
194        g.custom_command('create', 'cli_cosmosdb_mongodb_collection_create')
195        g.custom_command('update', 'cli_cosmosdb_mongodb_collection_update')
196        g.custom_command('exists', 'cli_cosmosdb_mongodb_collection_exists')
197        g.command('list', 'list_mongo_db_collections')
198        g.show_command('show', 'get_mongo_db_collection')
199        g.command('delete', 'begin_delete_mongo_db_collection', confirmation=True)
200
201    # Cassandra api
202    with self.command_group('cosmosdb cassandra'):
203        pass
204    with self.command_group('cosmosdb cassandra keyspace', cosmosdb_cassandra_sdk, client_factory=cf_cassandra_resources) as g:
205        g.custom_command('create', 'cli_cosmosdb_cassandra_keyspace_create')
206        g.custom_command('exists', 'cli_cosmosdb_cassandra_keyspace_exists')
207        g.command('list', 'list_cassandra_keyspaces')
208        g.show_command('show', 'get_cassandra_keyspace')
209        g.command('delete', 'begin_delete_cassandra_keyspace', confirmation=True)
210
211    with self.command_group('cosmosdb cassandra table', cosmosdb_cassandra_sdk, client_factory=cf_cassandra_resources) as g:
212        g.custom_command('create', 'cli_cosmosdb_cassandra_table_create')
213        g.custom_command('update', 'cli_cosmosdb_cassandra_table_update')
214        g.custom_command('exists', 'cli_cosmosdb_cassandra_table_exists')
215        g.command('list', 'list_cassandra_tables')
216        g.show_command('show', 'get_cassandra_table')
217        g.command('delete', 'begin_delete_cassandra_table', confirmation=True)
218
219    # Gremlin api
220    with self.command_group('cosmosdb gremlin'):
221        pass
222    with self.command_group('cosmosdb gremlin database', cosmosdb_gremlin_sdk, client_factory=cf_gremlin_resources) as g:
223        g.custom_command('create', 'cli_cosmosdb_gremlin_database_create')
224        g.custom_command('exists', 'cli_cosmosdb_gremlin_database_exists')
225        g.command('list', 'list_gremlin_databases')
226        g.show_command('show', 'get_gremlin_database')
227        g.command('delete', 'begin_delete_gremlin_database', confirmation=True)
228
229    with self.command_group('cosmosdb gremlin graph', cosmosdb_gremlin_sdk, client_factory=cf_gremlin_resources) as g:
230        g.custom_command('create', 'cli_cosmosdb_gremlin_graph_create')
231        g.custom_command('update', 'cli_cosmosdb_gremlin_graph_update')
232        g.custom_command('exists', 'cli_cosmosdb_gremlin_graph_exists')
233        g.command('list', 'list_gremlin_graphs')
234        g.show_command('show', 'get_gremlin_graph')
235        g.command('delete', 'begin_delete_gremlin_graph', confirmation=True)
236
237    # Table api
238    with self.command_group('cosmosdb table', cosmosdb_table_sdk, client_factory=cf_table_resources) as g:
239        g.custom_command('create', 'cli_cosmosdb_table_create')
240        g.custom_command('exists', 'cli_cosmosdb_table_exists')
241        g.command('list', 'list_tables')
242        g.show_command('show', 'get_table')
243        g.command('delete', 'begin_delete_table', confirmation=True)
244
245    # Offer throughput
246    with self.command_group('cosmosdb sql database throughput', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
247        g.show_command('show', 'get_sql_database_throughput')
248        g.custom_command('migrate', 'cli_cosmosdb_sql_database_throughput_migrate')
249        g.custom_command('update', 'cli_cosmosdb_sql_database_throughput_update')
250
251    with self.command_group('cosmosdb sql container throughput', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
252        g.show_command('show', 'get_sql_container_throughput')
253        g.custom_command('update', 'cli_cosmosdb_sql_container_throughput_update')
254        g.custom_command('migrate', 'cli_cosmosdb_sql_container_throughput_migrate')
255
256    with self.command_group('cosmosdb mongodb database throughput', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
257        g.show_command('show', 'get_mongo_db_database_throughput')
258        g.custom_command('update', 'cli_cosmosdb_mongodb_database_throughput_update')
259        g.custom_command('migrate', 'cli_cosmosdb_mongodb_database_throughput_migrate')
260
261    with self.command_group('cosmosdb mongodb collection throughput', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
262        g.show_command('show', 'get_mongo_db_collection_throughput')
263        g.custom_command('update', 'cli_cosmosdb_mongodb_collection_throughput_update')
264        g.custom_command('migrate', 'cli_cosmosdb_mongodb_collection_throughput_migrate')
265
266    with self.command_group('cosmosdb cassandra keyspace throughput', cosmosdb_cassandra_sdk, client_factory=cf_cassandra_resources) as g:
267        g.show_command('show', 'get_cassandra_keyspace_throughput')
268        g.custom_command('update', 'cli_cosmosdb_cassandra_keyspace_throughput_update')
269        g.custom_command('migrate', 'cli_cosmosdb_cassandra_keyspace_throughput_migrate')
270
271    with self.command_group('cosmosdb cassandra table throughput', cosmosdb_cassandra_sdk, client_factory=cf_cassandra_resources) as g:
272        g.show_command('show', 'get_cassandra_table_throughput')
273        g.custom_command('update', 'cli_cosmosdb_cassandra_table_throughput_update')
274        g.custom_command('migrate', 'cli_cosmosdb_cassandra_table_throughput_migrate')
275
276    with self.command_group('cosmosdb gremlin database throughput', cosmosdb_gremlin_sdk, client_factory=cf_gremlin_resources) as g:
277        g.show_command('show', 'get_gremlin_database_throughput')
278        g.custom_command('update', 'cli_cosmosdb_gremlin_database_throughput_update')
279        g.custom_command('migrate', 'cli_cosmosdb_gremlin_database_throughput_migrate')
280
281    with self.command_group('cosmosdb gremlin graph throughput', cosmosdb_gremlin_sdk, client_factory=cf_gremlin_resources) as g:
282        g.show_command('show', 'get_gremlin_graph_throughput')
283        g.custom_command('update', 'cli_cosmosdb_gremlin_graph_throughput_update')
284        g.custom_command('migrate', 'cli_cosmosdb_gremlin_graph_throughput_migrate')
285
286    with self.command_group('cosmosdb table throughput', cosmosdb_table_sdk, client_factory=cf_table_resources) as g:
287        g.show_command('show', 'get_table_throughput')
288        g.custom_command('update', 'cli_cosmosdb_table_throughput_update')
289        g.custom_command('migrate', 'cli_cosmosdb_table_throughput_migrate')
290
291    with self.command_group('cosmosdb identity', client_factory=cf_db_accounts, is_preview=True) as g:
292        g.custom_show_command('show', 'cli_cosmosdb_identity_show')
293        g.custom_command('assign', 'cli_cosmosdb_identity_assign')
294        g.custom_command('remove', 'cli_cosmosdb_identity_remove')
295
296    # virtual network rules
297    with self.command_group('cosmosdb network-rule', None, client_factory=cf_db_accounts) as g:
298        g.custom_command('list', 'cli_cosmosdb_network_rule_list', transform=transform_network_rule_list_output)
299        g.custom_command('add', 'cli_cosmosdb_network_rule_add', transform=transform_db_account_json_output)
300        g.custom_command('remove', 'cli_cosmosdb_network_rule_remove', transform=transform_db_account_json_output)
301
302    # key operations
303    with self.command_group('cosmosdb keys', cosmosdb_sdk, client_factory=cf_db_accounts) as g:
304        g.custom_command('list', 'cli_cosmosdb_keys', table_transformer=list_connection_strings_output)
305        g.custom_command('regenerate', 'cli_cosmosdb_regenerate_key')
306
307    # # database operations
308    with self.command_group('cosmosdb database', deprecate_info=self.deprecate(redirect=DATABASE_DEPRECATION_INFO, hide=True)) as g:
309        g.cosmosdb_custom('show', 'cli_cosmosdb_database_show', table_transformer=database_output)
310        g.cosmosdb_custom('list', 'cli_cosmosdb_database_list', table_transformer=list_database_output)
311        g.cosmosdb_custom('exists', 'cli_cosmosdb_database_exists')
312        g.cosmosdb_custom('create', 'cli_cosmosdb_database_create', table_transformer=database_output)
313        g.cosmosdb_custom('delete', 'cli_cosmosdb_database_delete', confirmation=True)
314
315    # collection operations
316    with self.command_group('cosmosdb collection', deprecate_info=self.deprecate(redirect=COLLECTION_DEPRECATON_INFO, hide=True)) as g:
317        g.cosmosdb_custom('show', 'cli_cosmosdb_collection_show', table_transformer=collection_output)
318        g.cosmosdb_custom('list', 'cli_cosmosdb_collection_list', table_transformer=list_collection_output)
319        g.cosmosdb_custom('exists', 'cli_cosmosdb_collection_exists')
320        g.cosmosdb_custom('create', 'cli_cosmosdb_collection_create', table_transformer=collection_output)
321        g.cosmosdb_custom('delete', 'cli_cosmosdb_collection_delete', confirmation=True)
322        g.cosmosdb_custom('update', 'cli_cosmosdb_collection_update')
323
324    # SQL role definition operations
325    with self.command_group('cosmosdb sql role definition', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
326        g.custom_command('create', 'cli_cosmosdb_sql_role_definition_create', supports_no_wait=True)
327        g.custom_command('update', 'cli_cosmosdb_sql_role_definition_update', supports_no_wait=True)
328        g.custom_command('exists', 'cli_cosmosdb_sql_role_definition_exists')
329        g.command('list', 'list_sql_role_definitions')
330        g.show_command('show', 'get_sql_role_definition')
331        g.command('delete', 'begin_delete_sql_role_definition', confirmation=True, supports_no_wait=True)
332        g.wait_command('wait', 'get_sql_role_definition')
333
334    # SQL role assignment operations
335    with self.command_group('cosmosdb sql role assignment', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
336        g.custom_command('create', 'cli_cosmosdb_sql_role_assignment_create', supports_no_wait=True)
337        g.custom_command('update', 'cli_cosmosdb_sql_role_assignment_update', supports_no_wait=True)
338        g.custom_command('exists', 'cli_cosmosdb_sql_role_assignment_exists')
339        g.command('list', 'list_sql_role_assignments')
340        g.show_command('show', 'get_sql_role_assignment')
341        g.command('delete', 'begin_delete_sql_role_assignment', confirmation=True, supports_no_wait=True)
342        g.wait_command('wait', 'get_sql_role_assignment')
343
344    with self.command_group('cosmosdb restorable-database-account', cosmosdb_restorable_database_accounts_sdk, client_factory=cf_restorable_database_accounts) as g:
345        g.show_command('show', 'get_by_location')
346        g.custom_command('list', 'cli_cosmosdb_restorable_database_account_list')
347
348    with self.command_group('cosmosdb sql restorable-database', cosmosdb_sql_restorable_database_sdk, client_factory=cf_restorable_sql_databases) as g:
349        g.command('list', 'list')
350
351    with self.command_group('cosmosdb sql restorable-container', cosmosdb_sql_restorable_container_sdk, client_factory=cf_restorable_sql_containers) as g:
352        g.command('list', 'list')
353
354    with self.command_group('cosmosdb sql restorable-resource', cosmosdb_sql_restorable_resources_sdk, client_factory=cf_restorable_sql_resources) as g:
355        g.command('list', 'list')
356
357    with self.command_group('cosmosdb mongodb restorable-database', cosmosdb_mongodb_restorable_database_sdk, client_factory=cf_restorable_mongodb_databases) as g:
358        g.command('list', 'list')
359
360    with self.command_group('cosmosdb mongodb restorable-collection', cosmosdb_mongodb_restorable_collection_sdk, client_factory=cf_restorable_mongodb_collections) as g:
361        g.command('list', 'list')
362
363    with self.command_group('cosmosdb mongodb restorable-resource', cosmosdb_mongodb_restorable_resources_sdk, client_factory=cf_restorable_mongodb_resources) as g:
364        g.command('list', 'list')
365
366    # Retrieve backup info
367    with self.command_group('cosmosdb sql', cosmosdb_sql_sdk, client_factory=cf_sql_resources, is_preview=True) as g:
368        g.custom_command('retrieve-latest-backup-time', 'cli_retrieve_latest_backup_time')
369