[SQL] Add new commands `sql db classification show/list/update/delete…

…` and `sql db classification recommendation list/enable/disable` to manage sensitivity classifications for SQL databases. (#11597)

* New commands `sql db sensitivity-labels show/list/list-recommended/update/delete/enable-recommendation/disable-recommendation` to manage sensitivity labels for SQL databases.

* Fix history file

* Fix failing style checks

* Fix help

* Use g.command() for list, list-recommended, delete, enable-recommendation and disable-recommendation

* rename sensitivity-labels to sensitivity-classification and add another group for recommendation

* Fix sensitivity classifications test after sync

* Rename command from `sensitivity-classification` to `classification`

* information_type and label_name should not be required

* SQl classification - fix comments

* SQL classification Additional fixes

* Update help

* SQL classification - show command should be seperated for current/recommended

Co-authored-by: Zunli Hu <zuh@microsoft.com>

src/azure-cli/HISTORY.rst

@@ -79,6 +79,7 @@ Release History
 **SQL**
 
 * [BREAKING CHANGE] `az sql db create`: Remove "WideWorldImportersStd" and "WideWorldImportersFull" as documented allowed values for "az sql db create --sample-name". These sample databases would always cause creation to fail.
+* Add New commands `sql db classification show/list/update/delete` and `sql db classification recommendation list/enable/disable` to manage sensitivity classifications for SQL databases.
 
 **Storage**
 

src/azure-cli/azure/cli/command_modules/sql/_help.py

@@ -825,3 +825,69 @@
   - name: Get the details for a virtual cluster
     text: az sql virtual-cluster show -g mygroup -n mycluster
 """
+
+helps['sql db classification'] = """
+type: group
+short-summary: Manage sensitivity classifications.
+"""
+
+helps['sql db classification update'] = """
+type: command
+short-summary: Update a columns's sensitivity classification.
+examples:
+  - name: Update sensitivity classification for a given column.
+    text: az sql db classification update -g mygroup -s myserver -n mydb --schema dbo --table mytable --column mycolumn --information-type Name --label "Confidential - GDPR"
+"""
+
+helps['sql db classification list'] = """
+type: command
+short-summary: Get the sensitivity classifications of a given database.
+examples:
+  - name: List the sensitivity classification of a given database.
+    text: az sql db classification list -g mygroup -s myserver -n mydb
+"""
+
+helps['sql db classification show'] = """
+type: command
+short-summary: Get the sensitivity classification of a given column.
+examples:
+  - name: Get the sensitivity classification of a given column.
+    text: az sql db classification show -g mygroup -s myserver -n mydb --schema dbo --table mytable --column mycolumn
+"""
+
+helps['sql db classification delete'] = """
+type: command
+short-summary: Delete the sensitivity classification of a given column.
+examples:
+  - name: Delete the sensitivity classification of a given column.
+    text: az sql db classification delete -g mygroup -s myserver -n mydb --schema dbo --table mytable --column mycolumn
+"""
+
+helps['sql db classification recommendation'] = """
+type: group
+short-summary: Manage sensitivity classification recommendations.
+"""
+
+helps['sql db classification recommendation list'] = """
+type: command
+short-summary: List the recommended sensitivity classifications of a given database.
+examples:
+  - name: List the recommended sensitivity classifications of a given database.
+    text: az sql db classification recommendation list -g mygroup -s myserver -n mydb
+"""
+
+helps['sql db classification recommendation enable'] = """
+type: command
+short-summary: Enable sensitivity recommendations for a given column (recommendations are enabled by default on all columns).
+examples:
+  - name: Enable sensitivity recommendations for a given column.
+    text: az sql db classification recommendation enable -g mygroup -s myserver -n mydb --schema dbo --table mytable --column mycolumn
+"""
+
+helps['sql db classification recommendation disable'] = """
+type: command
+short-summary: Disable sensitivity recommendations for a given column (recommendations are enabled by default on all columns).
+examples:
+  - name: Disable sensitivity recommendations for a given column.
+    text: az sql db classification recommendation disable -g mygroup -s myserver -n mydb --schema dbo --table mytable --column mycolumn
+"""

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

@@ -1428,3 +1428,34 @@ def _configure_security_policy_storage_params(arg_ctx):
 
         c.argument('allow_data_loss',
                    arg_type=allow_data_loss_param_type)
+
+    ###################################################
+    #             sql sensitivity classification      #
+    ###################################################
+    with self.argument_context('sql db classification') as c:
+        c.argument('schema_name',
+                   required=True,
+                   help='The name of the schema.',
+                   options_list=['--schema'])
+
+        c.argument('table_name',
+                   required=True,
+                   help='The name of the table.',
+                   options_list=['--table'])
+
+        c.argument('column_name',
+                   required=True,
+                   help='The name of the column.',
+                   options_list=['--column'])
+
+        c.argument('information_type',
+                   required=False,
+                   help='The information type.')
+
+        c.argument('label_name',
+                   required=False,
+                   help='The label name.',
+                   options_list=['--label'])
+
+    with self.argument_context('sql db classification recommendation list') as c:
+        c.ignore('skip_token')

src/azure-cli/azure/cli/command_modules/sql/_util.py

@@ -40,6 +40,10 @@ def get_sql_database_blob_auditing_policies_operations(cli_ctx, _):
     return get_sql_management_client(cli_ctx).database_blob_auditing_policies
 
 
+def get_sql_database_sensitivity_labels_operations(cli_ctx, _):
+    return get_sql_management_client(cli_ctx).sensitivity_labels
+
+
 def get_sql_database_threat_detection_policies_operations(cli_ctx, _):
     return get_sql_management_client(cli_ctx).database_threat_detection_policies
 

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

@@ -25,6 +25,7 @@
     get_sql_capabilities_operations,
     get_sql_databases_operations,
     get_sql_database_blob_auditing_policies_operations,
+    get_sql_database_sensitivity_labels_operations,
     get_sql_database_operations_operations,
     get_sql_database_threat_detection_policies_operations,
     get_sql_database_transparent_data_encryption_activities_operations,
@@ -231,6 +232,27 @@ def load_command_table(self, _):
         g.generic_update_command('update',
                                  custom_func_name='db_audit_policy_update')
 
+    database_sensitivity_labels_operations = CliCommandType(
+        operations_tmpl='azure.mgmt.sql.operations#SensitivityLabelsOperations.{}',
+        client_factory=get_sql_database_sensitivity_labels_operations)
+
+    with self.command_group('sql db classification',
+                            database_sensitivity_labels_operations,
+                            client_factory=get_sql_database_sensitivity_labels_operations) as g:
+
+        g.command('list', 'list_current_by_database')
+        g.custom_command('show', 'db_sensitivity_label_show')
+        g.command('delete', 'delete')
+        g.custom_command('update', 'db_sensitivity_label_update')
+
+    with self.command_group('sql db classification recommendation',
+                            database_sensitivity_labels_operations,
+                            client_factory=get_sql_database_sensitivity_labels_operations) as g:
+
+        g.command('list', 'list_recommended_by_database')
+        g.command('enable', 'enable_recommendation')
+        g.command('disable', 'disable_recommendation')
+
     database_threat_detection_policies_operations = CliCommandType(
         operations_tmpl='azure.mgmt.sql.operations#DatabaseThreatDetectionPoliciesOperations.{}',
         client_factory=get_sql_database_threat_detection_policies_operations)

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

@@ -26,6 +26,8 @@
     ReplicationRole,
     ResourceIdentity,
     SecurityAlertPolicyState,
+    SensitivityLabel,
+    SensitivityLabelSource,
     ServerKey,
     ServerKeyType,
     ServiceObjectiveName,
@@ -1484,6 +1486,99 @@ def db_threat_detection_policy_update(
     return instance
 
 
+def db_sensitivity_label_show(
+        client,
+        database_name,
+        server_name,
+        schema_name,
+        table_name,
+        column_name,
+        resource_group_name):
+
+    return client.get(
+        resource_group_name,
+        server_name,
+        database_name,
+        schema_name,
+        table_name,
+        column_name,
+        SensitivityLabelSource.current)
+
+
+def db_sensitivity_label_update(
+        cmd,
+        client,
+        database_name,
+        server_name,
+        schema_name,
+        table_name,
+        column_name,
+        resource_group_name,
+        label_name=None,
+        information_type=None):
+    '''
+    Updates a sensitivity label. Custom update function to apply parameters to instance.
+    '''
+
+    # Get the information protection policy
+    from azure.cli.core.commands.client_factory import get_mgmt_service_client
+    from azure.mgmt.security import SecurityCenter
+    from msrestazure.azure_exceptions import CloudError
+
+    security_center_client = get_mgmt_service_client(cmd.cli_ctx, SecurityCenter, asc_location="centralus")
+
+    information_protection_policy = security_center_client.information_protection_policies.get(
+        scope='/providers/Microsoft.Management/managementGroups/{}'.format(_get_tenant_id()),
+        information_protection_policy_name="effective")
+
+    sensitivity_label = SensitivityLabel()
+
+    # Get the current label
+    try:
+        current_label = client.get(
+            resource_group_name,
+            server_name,
+            database_name,
+            schema_name,
+            table_name,
+            column_name,
+            SensitivityLabelSource.current)
+        # Initialize with existing values
+        sensitivity_label.label_name = current_label.label_name
+        sensitivity_label.label_id = current_label.label_id
+        sensitivity_label.information_type = current_label.information_type
+        sensitivity_label.information_type_id = current_label.information_type_id
+
+    except CloudError as ex:
+        if not(ex.error and ex.error.error and 'SensitivityLabelsLabelNotFound' in ex.error.error):
+            raise ex
+
+    # Find the label id and information type id in the policy by the label name provided
+    label_id = None
+    if label_name:
+        label_id = next((id for id in information_protection_policy.labels
+                         if information_protection_policy.labels[id].display_name.lower() ==
+                         label_name.lower()),
+                        None)
+        if label_id is None:
+            raise CLIError('The provided label name was not found in the information protection policy.')
+        sensitivity_label.label_id = label_id
+        sensitivity_label.label_name = label_name
+    information_type_id = None
+    if information_type:
+        information_type_id = next((id for id in information_protection_policy.information_types
+                                    if information_protection_policy.information_types[id].display_name.lower() ==
+                                    information_type.lower()),
+                                   None)
+        if information_type_id is None:
+            raise CLIError('The provided information type was not found in the information protection policy.')
+        sensitivity_label.information_type_id = information_type_id
+        sensitivity_label.information_type = information_type
+
+    return client.create_or_update(
+        resource_group_name, server_name, database_name, schema_name, table_name, column_name, sensitivity_label)
+
+
 ###############################################
 #                sql dw                       #
 ###############################################