[Synapse] az synapse kql-script: New command group to support Kusto script

src/azure-cli/azure/cli/command_modules/synapse/manual/_client_factory.py

@@ -242,3 +242,11 @@ def cf_synapse_managedprivateendpoints_factory(cli_ctx, workspace_name):
 
 def cf_synapse_spark_job_definition(cli_ctx, workspace_name):
     return cf_synapse_client_artifacts_factory(cli_ctx, workspace_name).spark_job_definition
+
+
+def cf_kusto_script(cli_ctx, workspace_name):
+    return cf_synapse_client_artifacts_factory(cli_ctx, workspace_name).kql_script
+
+
+def cf_kusto_scripts(cli_ctx, workspace_name):
+    return cf_synapse_client_artifacts_factory(cli_ctx, workspace_name).kql_scripts

src/azure-cli/azure/cli/command_modules/synapse/manual/_help.py

@@ -2135,3 +2135,71 @@
     type: group
     short-summary: "Manage synapse kusto"
 """
+
+helps['synapse kql-script'] = """
+    type: group
+    short-summary: Manage script with kusto
+"""
+
+helps['synapse kql-script show'] = """
+    type: command
+    short-summary: "Gets a KQL script."
+    examples:
+      - name: KustoScriptsGet
+        text: |-
+               az synapse kql-script show --workspace-name "kustoWorkspaceName" --name "kustoScript1"
+"""
+
+helps['synapse kql-script list'] = """
+    type: command
+    short-summary: "List KQL scripts."
+    examples:
+      - name: KustoScriptsList
+        text: |-
+               az synapse kql-script list --workspace-name "kustoWorkspaceName"
+"""
+
+helps['synapse kql-script create'] = """
+    type: command
+    short-summary: "Creates a KQL script."
+    examples:
+      - name: KustoScriptsCreateOrUpdate
+        text: |-
+               az synapse kql-script create --resource-group "kustorptest" --workspace-name "kustoWorkspaceName" \
+               --kusto-pool-name kustopooltest --kusto-database-name kustodbtest --file C:\\samples\\KqlScript.kql \
+               --name "kustoScript1"
+"""
+
+helps['synapse kql-script import'] = """
+    type: command
+    short-summary: "Creates a KQL script."
+    examples:
+      - name: KustoScriptsCreateOrUpdate
+        text: |-
+               az synapse kql-script import --resource-group "kustorptest" --workspace-name "kustoWorkspaceName" \
+               --kusto-pool-name kustopooltest --kusto-database-name kustodbtest --file C:\\samples\\KqlScript.kql \
+               --name "kustoScript1"
+"""
+
+helps['synapse kql-script export'] = """
+    type: command
+    short-summary: "Export KQL scripts."
+    examples:
+      - name: KustoScriptsExport
+        text: |-
+               az synapse kql-script export --workspace-name "kustoWorkspaceName" --output-folder "C:\\KqlScirpt"
+"""
+
+helps['synapse kql-script delete'] = """
+    type: command
+    short-summary: "Deletes a KQL script"
+    examples:
+      - name: KustoScriptsDelete
+        text: |-
+               az synapse kql-script delete --workspace-name "kustoWorkspaceName" --name "kustoScript1"
+"""
+
+helps['synapse kql-script wait'] = """
+    type: command
+    short-summary: "Place the CLI in a waiting state until a condition of a KQL script is met."
+"""

src/azure-cli/azure/cli/command_modules/synapse/manual/_params.py

@@ -1042,3 +1042,26 @@ def load_arguments(self, _):
                    options_list=['--attached-database-configuration-name', '--adcn'],
                    type=str, help='Resource name of the attached database '
                    'configuration in the follower cluster.')
+
+    for scope in ['import', 'create']:
+        with self.argument_context('synapse kql-script ' + scope) as c:
+            c.argument('workspace_name', arg_type=workspace_name_arg_type, help='The name of the workspace')
+            c.argument('kusto_pool_name', type=str, help='The name of the Kusto pool.')
+            c.argument('kusto_database_name', type=str, help='The name of the Kusto database.')
+            c.argument('script_name', arg_type=name_type, help='The name of the KQL script.')
+            c.argument('definition_file', options_list=['--file', '-f'], type=file_type, completer=FilesCompleter(),
+                       help='The KQL query file path')
+
+    for scope in ['show', 'wait', 'delete']:
+        with self.argument_context('synapse kql-script ' + scope) as c:
+            c.argument('workspace_name', arg_type=workspace_name_arg_type, help='The name of the workspace')
+            c.argument('script_name', arg_type=name_type,
+                       help='The name of the KQL script.')
+
+    with self.argument_context('synapse kql-script list') as c:
+        c.argument('workspace_name', arg_type=workspace_name_arg_type, help='The name of the workspace')
+
+    with self.argument_context('synapse kql-script export') as c:
+        c.argument('workspace_name', arg_type=workspace_name_arg_type, help='The name of the workspace')
+        c.argument('output_folder', type=str, help='The name of the output folder')
+        c.argument('script_name', arg_type=name_type, help='The name of the KQL script.')

src/azure-cli/azure/cli/command_modules/synapse/manual/commands.py

@@ -30,6 +30,8 @@ def load_command_table(self, _):
     from ._client_factory import cf_synapse_client_integrationruntimeconnectioninfos_factory
     from ._client_factory import cf_synapse_client_integrationruntimestatus_factory
     from ._client_factory import cf_kusto_pool
+    from ._client_factory import cf_kusto_script
+    from ._client_factory import cf_kusto_scripts
 
     def get_custom_sdk(custom_module, client_factory):
         return CliCommandType(
@@ -195,6 +197,12 @@ def get_custom_sdk(custom_module, client_factory):
         operations_tmpl='azure.mgmt.synapse.operations._kusto_pools_operations#KustoPoolsOperations.{}',
         client_factory=cf_kusto_pool,
     )
+
+    synapse_kusto_script_sdk = CliCommandType(
+        operations_tmpl='azure.synapse.artifacts.operations#KqlScriptOperations.{}',
+        client_factory=cf_kusto_script,
+    )
+
     # Management Plane Commands --Workspace
     with self.command_group('synapse workspace', command_type=synapse_workspace_sdk,
                             custom_command_type=get_custom_sdk('workspace', cf_synapse_client_workspace_factory),
@@ -544,3 +552,14 @@ def get_custom_sdk(custom_module, client_factory):
         g.custom_command('add-language-extension', 'synapse_kusto_pool_add_language_extension', supports_no_wait=True)
         g.custom_command('detach-follower-database', 'synapse_kusto_pool_detach_follower_database', supports_no_wait=True)
         g.custom_command('remove-language-extension', 'synapse_kusto_pool_remove_language_extension', supports_no_wait=True)
+
+    with self.command_group('synapse kql-script', command_type=synapse_kusto_script_sdk,
+                            custom_command_type=get_custom_sdk('kustopool', cf_kusto_script),
+                            client_factory=cf_kusto_script) as g:
+        g.custom_show_command('show', 'synapse_kusto_script_show')
+        g.custom_command('create', 'synapse_kusto_script_create', supports_no_wait=True)
+        g.custom_command('import', 'synapse_kusto_script_create', supports_no_wait=True)
+        g.custom_command('delete', 'synapse_kusto_script_delete', supports_no_wait=True, confirmation=True)
+        g.custom_command('list', 'synapse_kusto_script_list', client_factory=cf_kusto_scripts)
+        g.custom_command('export', 'synapse_kusto_script_export')
+        g.custom_wait_command('wait', 'synapse_kusto_script_show')

src/azure-cli/azure/cli/command_modules/synapse/manual/operations/kustopool.py

@@ -3,8 +3,10 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 # pylint: disable=unused-argument, line-too-long
-from azure.cli.core.util import sdk_no_wait
-from .._client_factory import cf_synapse_client_workspace_factory
+from azure.cli.core.util import sdk_no_wait, read_file_content, CLIError
+from .._client_factory import cf_synapse_client_workspace_factory, cf_kusto_script, cf_kusto_scripts
+from azure.synapse.artifacts.models import KqlScriptResource, KqlScript, KqlScriptContent, KqlScriptContentMetadata, KqlScriptContentCurrentConnection
+import os
 
 
 # Synapse kustopool
@@ -147,3 +149,76 @@ def synapse_kusto_pool_detach_follower_database(client,
                        kusto_pool_name=kusto_pool_name,
                        resource_group_name=resource_group_name,
                        follower_database_to_remove=follower_database_to_remove)
+
+
+def synapse_kusto_script_show(cmd, client,
+                              workspace_name,
+                              script_name):
+    client = cf_kusto_script(cmd.cli_ctx, workspace_name)
+    return client.get_by_name(kql_script_name=script_name)
+
+
+def synapse_kusto_script_create(cmd, client,
+                                resource_group_name,
+                                workspace_name,
+                                script_name,
+                                definition_file,
+                                kusto_pool_name=None,
+                                kusto_database_name=None,
+                                no_wait=False):
+
+    client = cf_kusto_script(cmd.cli_ctx, workspace_name)
+    query = read_file_content(definition_file)
+    metadata = KqlScriptContentMetadata(language="kql")
+    current_connection = KqlScriptContentCurrentConnection(pool_name=kusto_pool_name,
+                                                           database_name=kusto_database_name)
+    script_content = KqlScriptContent(query=query, metadata=metadata, current_connection=current_connection)
+    properties = KqlScript(content=script_content)
+    kql_script = KqlScriptResource(name=script_name, properties=properties)
+    return sdk_no_wait(no_wait,
+                       client.begin_create_or_update,
+                       kql_script_name=script_name,
+                       kql_script=kql_script
+                       )
+
+
+def synapse_kusto_script_delete(cmd, client,
+                                workspace_name,
+                                script_name,
+                                no_wait=False):
+    client = cf_kusto_script(cmd.cli_ctx, workspace_name)
+
+    return sdk_no_wait(no_wait,
+                       client.begin_delete_by_name,
+                       kql_script_name=script_name)
+
+
+def synapse_kusto_script_list(cmd, client,
+                              workspace_name):
+    client = cf_kusto_scripts(cmd.cli_ctx, workspace_name)
+    return client.get_all()
+
+
+def synapse_kusto_script_export(cmd, workspace_name, output_folder, script_name=None):
+    if script_name is not None:
+        kusto_script_client = cf_kusto_script(cmd.cli_ctx, workspace_name)
+        kql_script = kusto_script_client.get_by_name(script_name)
+        path = os.path.join(output_folder, script_name + '.kql')
+        write_to_file(kql_script, path)
+    else:
+        kusto_script_client = cf_kusto_scripts(cmd.cli_ctx, workspace_name)
+        kql_scripts = kusto_script_client.get_all()
+        for kql_script in kql_scripts:
+            path = os.path.join(output_folder, kql_script.name + '.kql')
+            write_to_file(kql_script, path)
+
+
+def write_to_file(kql_script, path):
+    try:
+        query = ''
+        if hasattr(kql_script.properties.content, 'query'):
+            query = kql_script.properties.content.query
+        with open(path, 'w') as f:
+            f.write(query)
+    except IOError:
+        raise CLIError('Unable to export to file: {}'.format(path))

src/azure-cli/azure/cli/command_modules/synapse/tests/latest/assets/kqlScript.kql

@@ -0,0 +1 @@
+Logs | where TimeStamp > ago(1h) | take 10