[Cosmos DB] Add Sql stored procedure, udf and trigger cmdlets

src/azure-cli/azure/cli/command_modules/cosmosdb/_params.py

@@ -115,6 +115,32 @@ def load_arguments(self, _):
         c.argument('conflict_resolution_policy', options_list=['--conflict-resolution-policy', '-c'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Conflict Resolution Policy, you can enter it as a string or as a file, e.g., --conflict-resolution-policy @policy-file.json or ' + SQL_GREMLIN_CONFLICT_RESOLUTION_POLICY_EXAMPLE)
         c.argument('throughput', help='The throughput of SQL container (RU/s). Default value is 400')
 
+# SQL stored procedure
+    with self.argument_context('cosmosdb sql stored_procedure') as c:
+        c.argument('account_name', options_list=['--account-name', '-a'], help="Cosmosdb account name", id_part=None)
+        c.argument('database_name', options_list=['--database-name', '-d'], help="Database name")
+        c.argument('container_name', options_list=['--container-name', '-c'], help="Container name")
+        c.argument('stored_procedure_name', options_list=['--name', '-n'], help="StoredProcedure name")
+        c.argument('stored_procedure_body', options_list=['--body', '-b'], help="StoredProcedure body")
+
+# SQL trigger
+    with self.argument_context('cosmosdb sql trigger') as c:
+        c.argument('account_name', options_list=['--account-name', '-a'], help="Cosmosdb account name", id_part=None)
+        c.argument('database_name', options_list=['--database-name', '-d'], help="Database name")
+        c.argument('container_name', options_list=['--container', '-c'], help="Container name")
+        c.argument('trigger_name', options_list=['--name', '-n'], help="Trigger name")
+        c.argument('trigger_body', options_list=['--body', '-b'], help="Trigger body")
+        c.argument('trigger_type', options_list=['--type', '-t'], help="Trigger type, values can be: Pre or Post")
+        c.argument('trigger_operation', options_list=['--operation'], help="The operation of the trigger. Values can be: All, Create, Update, Delete, Replace")
+
+# SQL user defined function
+    with self.argument_context('cosmosdb sql user_defined_function') as c:
+        c.argument('account_name', options_list=['--account-name', '-a'], help="Cosmosdb account name", id_part=None)
+        c.argument('database_name', options_list=['--database-name', '-d'], help="Database name")
+        c.argument('container_name', options_list=['--container', '-c'], help="Container name")
+        c.argument('user_defined_function_name', options_list=['--name', '-n'], help="UserDefinedFunction name")
+        c.argument('user_defined_function_body', options_list=['--body', '-b'], help="UserDefinedFunction body")
+
 # MongoDB
     with self.argument_context('cosmosdb mongodb database') as c:
         c.argument('account_name', options_list=['--account-name', '-a'], help="Cosmosdb account name", id_part=None)

src/azure-cli/azure/cli/command_modules/cosmosdb/commands.py

@@ -77,6 +77,27 @@ def load_command_table(self, _):
         g.command('list', 'list_sql_containers')
         g.command('show', 'get_sql_container')
         g.command('delete', 'delete_sql_container')
+
+    with self.command_group('cosmosdb sql stored_procedure', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
+        g.custom_command('create', 'cli_cosmosdb_sql_stored_procedure_create_update')
+        g.custom_command('update', 'cli_cosmosdb_sql_stored_procedure_create_update')
+        g.command('list', 'list_sql_stored_procedures')
+        g.command('show', 'get_sql_stored_procedure')
+        g.command('delete', 'delete_sql_stored_procedure')
+
+    with self.command_group('cosmosdb sql trigger', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
+        g.custom_command('create', 'cli_cosmosdb_sql_trigger_create')
+        g.custom_command('update', 'cli_cosmosdb_sql_trigger_update')
+        g.command('list', 'list_sql_triggers')
+        g.command('show', 'get_sql_trigger')
+        g.command('delete', 'delete_sql_trigger')
+
+    with self.command_group('cosmosdb sql user_defined_function', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
+        g.custom_command('create', 'cli_cosmosdb_sql_user_defined_function_create_update')
+        g.custom_command('update', 'cli_cosmosdb_sql_user_defined_function_create_update')
+        g.command('list', 'list_sql_user_defined_functions')
+        g.command('show', 'get_sql_user_defined_function')
+        g.command('delete', 'delete_sql_user_defined_function')
 
     # MongoDB api
     with self.command_group('cosmosdb mongodb', is_preview=True):

src/azure-cli/azure/cli/command_modules/cosmosdb/custom.py

@@ -21,6 +21,12 @@
     SqlContainerResource,
     SqlContainerCreateUpdateParameters,
     ContainerPartitionKey,
+    SqlStoredProcedureResource,
+    SqlStoredProcedureCreateUpdateParameters,
+    SqlTriggerResource,
+    SqlTriggerCreateUpdateParameters,
+    SqlUserDefinedFunctionResource,
+    SqlUserDefinedFunctionCreateUpdateParameters,
     TableResource,
     TableCreateUpdateParameters,
     MongoDBDatabaseResource,
@@ -282,7 +288,6 @@ def cli_cosmosdb_sql_container_create(client,
                                               container_name,
                                               sql_container_create_update_resource)
 
-
 def cli_cosmosdb_sql_container_update(client,
                                       resource_group_name,
                                       account_name,
@@ -319,6 +324,139 @@ def cli_cosmosdb_sql_container_update(client,
                                               container_name,
                                               sql_container_create_update_resource)
 
+def cli_cosmosdb_sql_stored_procedure_create_update(client,
+                                      resource_group_name,
+                                      account_name,
+                                      database_name,
+                                      container_name,
+                                      stored_procedure_name,
+                                      stored_procedure_body):
+
+    """Creates or Updates an Azure Cosmos DB SQL stored procedure """
+    sql_stored_procedure_resource = SqlStoredProcedureResource(id=stored_procedure_name)
+    sql_stored_procedure_resource.body = stored_procedure_body
+
+    sql_stored_procedure_create_update_resource = SqlStoredProcedureCreateUpdateParameters(
+        resource=sql_stored_procedure_resource,
+        options={})
+
+    return client.create_update_sql_stored_procedure(resource_group_name,
+                                              account_name,
+                                              database_name,
+                                              container_name,
+                                              stored_procedure_name,
+                                              sql_stored_procedure_create_update_resource)
+
+def cli_cosmosdb_sql_user_defined_function_create_update(client,
+                                      resource_group_name,
+                                      account_name,
+                                      database_name,
+                                      container_name,
+                                      user_defined_function_name,
+                                      user_defined_function_body):
+
+    """Creates or Updates an Azure Cosmos DB SQL user defined function """
+    sql_user_defined_function_resource = SqlUserDefinedFunctionResource(id=user_defined_function_name)
+    sql_user_defined_function_resource.body = user_defined_function_body
+
+    sql_user_defined_function_create_update_resource = SqlUserDefinedFunctionCreateUpdateParameters(
+        resource=sql_user_defined_function_resource,
+        options={})
+
+    return client.create_update_sql_user_defined_function(resource_group_name,
+                                              account_name,
+                                              database_name,
+                                              container_name,
+                                              user_defined_function_name,
+                                              sql_user_defined_function_create_update_resource)
+
+def _populate_sql_trigger_definition(sql_trigger_resource,
+                                       trigger_body,
+                                       trigger_operation,
+                                       trigger_type):
+    if all(arg is None for arg in
+           [trigger_body, trigger_operation, trigger_type]):
+        return False
+
+    if trigger_body is not None:
+        sql_trigger_resource.body = trigger_body
+
+    if trigger_body is not None:
+        sql_trigger_resource.trigger_operation = trigger_operation
+
+    if trigger_body is not None:
+        sql_trigger_resource.trigger_type = trigger_type
+
+    return True
+
+def cli_cosmosdb_sql_trigger_create(client,
+                                      resource_group_name,
+                                      account_name,
+                                      database_name,
+                                      container_name,
+                                      trigger_name,
+                                      trigger_body,
+                                      trigger_type=None,
+                                      trigger_operation=None):
+
+    """Creates an Azure Cosmos DB SQL trigger """
+    if trigger_operation is None:
+        trigger_operation = "All"
+
+    if trigger_type is None:
+        trigger_type = "Pre"
+
+    sql_trigger_resource = SqlTriggerResource(id=trigger_name)
+    sql_trigger_resource.body = trigger_body
+    sql_trigger_resource.trigger_type = trigger_type
+    sql_trigger_resource.trigger_operation = trigger_operation
+
+    sql_trigger_create_update_resource = SqlTriggerCreateUpdateParameters(
+        resource=sql_trigger_resource,
+        options={})
+
+    return client.create_update_sql_trigger(resource_group_name,
+                                              account_name,
+                                              database_name,
+                                              container_name,
+                                              trigger_name,
+                                              sql_trigger_create_update_resource)
+
+def cli_cosmosdb_sql_trigger_update(client,
+                                      resource_group_name,
+                                      account_name,
+                                      database_name,
+                                      container_name,
+                                      trigger_name,
+                                      trigger_body=None,
+                                      trigger_type=None,
+                                      trigger_operation=None):
+
+    """Updates an Azure Cosmos DB SQL trigger """
+    logger.debug('reading SQL trigger')
+    sql_trigger = client.get_sql_trigger(resource_group_name, account_name, database_name, container_name, trigger_name)
+
+    sql_trigger_resource = SqlTriggerResource(id=trigger_name)
+    sql_trigger_resource.body = sql_trigger.resource.body
+    sql_trigger_resource.trigger_operation = sql_trigger.resource.trigger_operation
+    sql_trigger_resource.trigger_type = sql_trigger.resource.trigger_type
+
+    if _populate_sql_trigger_definition(sql_trigger_resource,
+                                          trigger_body,
+                                          trigger_operation,
+                                          trigger_type):
+        logger.debug('replacing SQL trigger')
+
+    sql_trigger_create_update_resource = SqlTriggerCreateUpdateParameters(
+        resource=sql_trigger_resource,
+        options={})
+
+    return client.create_update_sql_trigger(resource_group_name,
+                                              account_name,
+                                              database_name,
+                                              container_name,
+                                              trigger_name,
+                                              sql_trigger_create_update_resource)
 
 def cli_cosmosdb_gremlin_database_create(client,
                                          resource_group_name,