redisenterprise: add test-connection command to address a common supportability ask (#9490)

* add test-connection command

* update version

Author: yugangw-msft

src/redisenterprise/HISTORY.rst

@@ -2,6 +2,9 @@
 
 Release History
 ===============
+1.4.0
+- add a new command `az redisenterprise test-connection` to test the connection to a cluster.
+
 1.3.1
 - Fixed an issue where updating sku from Azure Cache for Enterprise to Azure Managed Redis SKU was not working as expected.
 

src/redisenterprise/azext_redisenterprise/_help.py

@@ -33,3 +33,16 @@
     type: command
     short-summary: Get information about a RedisEnterprise cluster
 """
+
+helps['redisenterprise test-connection'] = """
+    type: command
+    short-summary: Test connection to a Redis Enterprise cluster
+    long-summary: Test the connection to a Redis Enterprise cluster using the specified authentication method.
+    examples:
+      - name: Test connection using Entra authentication
+        text: |-
+               az redisenterprise test-connection --cluster-name "cache1" --resource-group "rg1" --auth entra
+      - name: Test connection using access key authentication
+        text: |-
+               az redisenterprise test-connection --cluster-name "cache1" --resource-group "rg1" --auth access-key
+"""

src/redisenterprise/azext_redisenterprise/_params.py

@@ -247,6 +247,17 @@ def load_arguments(self, _):
         c.argument('cluster_name', options_list=['--cluster-name', '--name', '-n'], type=str, help='The name of the '
                    'RedisEnterprise cluster.', id_part='name')
 
+    with self.argument_context('redisenterprise test-connection') as c:
+        c.argument('resource_group_name', resource_group_name_type, required=False)
+        c.argument('cluster_name', options_list=['--cluster-name', '--name', '-n'], type=str, help='The name of the '
+                   'RedisEnterprise cluster.', id_part='name', required=False)
+        c.argument('auth', arg_type=get_enum_type(['entra', 'access-key']), help='The authentication method to use '
+                   'for the connection test.')
+        c.argument('host_name', options_list=['--host', '--host-name'], type=str, help='The Redis host name (FQDN) to connect to. '
+                   'If provided, --resource-group and --cluster-name are optional.')
+        c.argument('access_key', options_list=['--access-key'], type=str, help='The access key for authentication. '
+                   'Required when using --host with access-key auth.')
+
 
 class AddPersistence(argparse.Action):
     def __call__(self, parser, namespace, values, option_string=None):

src/redisenterprise/azext_redisenterprise/commands.py

@@ -18,6 +18,7 @@ def load_command_table(self, _):  # pylint: disable=unused-argument
         g.custom_show_command('show', 'redisenterprise_show')
         from .custom import RedisEnterpriseUpdate
         self.command_table["redisenterprise update"] = RedisEnterpriseUpdate(loader=self)
+        g.custom_command('test-connection', 'redisenterprise_test_connection')
     with self.command_group("redisenterprise database"):
         from .custom import DatabaseFlush, DatabaseCreate, DatabaseDelete, DatabaseExport, DatabaseForceUnlink
         from .custom import DatabaseImport, DatabaseListKey, DatabaseRegenerateKey

src/redisenterprise/azext_redisenterprise/custom.py

@@ -28,6 +28,7 @@
 from azure.cli.core.azclierror import (
     MutuallyExclusiveArgumentError,
 )
+from azure.cli.core.azclierror import ValidationError
 
 logger = get_logger(__name__)
 
@@ -69,6 +70,241 @@ def _handle_sku_change(self, current_sku, new_sku, instance):
                 pass
 
 
+def _get_redis_connection(host_name, port, password, ssl=True, username=None):
+    """
+    Create a Redis connection using the provided credentials.
+
+    :param host_name: The Redis host name.
+    :param port: The Redis port.
+    :param password: The password or token for authentication.
+    :param ssl: Whether to use SSL connection.
+    :param username: The username for authentication (required for Entra ID auth).
+    :return: Redis client instance.
+    """
+    import redis
+    return redis.Redis(
+        host=host_name,
+        port=port,
+        username=username,
+        password=password,
+        ssl=ssl,
+        ssl_cert_reqs=None,
+        decode_responses=True
+    )
+
+
+def _test_redis_connection_with_ping(redis_client):
+    """
+    Test Redis connection by sending a PING command.
+
+    :param redis_client: The Redis client instance.
+    :return: Tuple of (success: bool, message: str).
+    """
+    test_logger = get_logger(__name__)
+
+    try:
+        test_logger.warning("Sending PING command to Redis server...")
+        response = redis_client.ping()
+
+        if response:
+            test_logger.warning("PING successful, received PONG response.")
+            return True, "Successfully connected and verified with PING command."
+
+        test_logger.warning("PING command did not receive expected response.")
+        return False, "PING command did not receive expected response."
+
+    except Exception as e:  # pylint: disable=broad-except
+        error_msg = f"Failed during PING operation: {str(e)}"
+        test_logger.warning("Error: %s", error_msg)
+        return False, error_msg
+
+
+def redisenterprise_test_connection(cmd,
+                                    resource_group_name=None,
+                                    cluster_name=None,
+                                    auth=None,
+                                    host_name=None,
+                                    access_key=None):
+    # pylint: disable=too-many-branches, line-too-long
+    """
+    Test connection to a Redis Enterprise cluster using the specified authentication method.
+
+    :param cmd: The command instance.
+    :param resource_group_name: The name of the resource group (optional if host is provided).
+    :param cluster_name: The name of the Redis Enterprise cluster (optional if host is provided).
+    :param auth: The authentication method to use ('entra' or 'access-key').
+    :param host_name: The Redis host name (optional, will be retrieved from cluster if not provided).
+    :param access_key: The access key for authentication (optional, only used with access-key auth).
+    :return: Connection test result.
+    """
+    # Default port for Redis Enterprise
+    port = 10000
+    database_name = 'default'
+
+    # Parse port from hostname if provided (e.g., "myhost.redis.cache.windows.net:10000")
+    if host_name and ':' in host_name:
+        host_parts = host_name.rsplit(':', 1)
+        host_name = host_parts[0]
+        try:
+            port = int(host_parts[1])
+        except ValueError:
+            raise ValidationError(f"Invalid port number in hostname: {host_parts[1]}")
+
+    # Infer auth method from provided parameters
+    if access_key and not auth:
+        auth = 'access-key'
+
+    # Validate parameters - host and cluster-name are mutually exclusive
+    if host_name and cluster_name:
+        raise MutuallyExclusiveArgumentError(
+            "--host and --cluster-name cannot be used together.",
+            "Use either --host (with --access-key or --auth entra) or --resource-group and --cluster-name.")
+
+    # Validate parameters
+    if host_name:
+        # Direct connection mode - host is provided
+        # Default to Entra auth if neither access_key nor auth is provided
+        if not access_key and not auth:
+            auth = 'entra'
+    else:
+        # Azure resource mode - need resource group and cluster name
+        if not resource_group_name or not cluster_name:
+            raise ValidationError("Either --host or both --resource-group and --cluster-name must be provided.")
+
+        # Default to Entra auth if auth is not specified
+        if not auth:
+            auth = 'entra'
+
+        # Get cluster information
+        cluster = _ClusterShow(cli_ctx=cmd.cli_ctx)(command_args={
+            "cluster_name": cluster_name,
+            "resource_group": resource_group_name})
+
+        if not cluster:
+            raise ValidationError(f"Cluster '{cluster_name}' not found in resource group '{resource_group_name}'.")
+
+        # Get the hostname from the cluster
+        host_name = cluster.get('hostName')
+        if not host_name:
+            raise ValidationError(f"Unable to retrieve hostname for cluster '{cluster_name}'.")
+
+    result = {
+        'clusterName': cluster_name,
+        'resourceGroup': resource_group_name,
+        'hostName': host_name,
+        'port': port,
+        'databaseName': database_name,
+        'authMethod': auth,
+        'connectionStatus': 'NotTested',
+        'message': ''
+    }
+
+    if auth == 'entra':
+        # Get token from current Azure CLI credentials for Redis
+        try:
+            from azure.cli.core._profile import Profile
+
+            profile = Profile(cli_ctx=cmd.cli_ctx)
+            # Use get_raw_token with the Redis resource
+            creds, _, _ = profile.get_raw_token(resource="https://redis.azure.com")
+            access_token = creds[1]
+
+            logger.debug("Successfully obtained Entra ID token for Redis.")
+
+            # Create Redis connection with the token
+            # For Entra auth, username is the object ID (oid) from the token
+            # The password is the access token itself
+            import jwt
+            decoded_token = jwt.decode(access_token, options={"verify_signature": False})
+            entra_object_id = decoded_token.get('oid', decoded_token.get('sub', 'default'))
+
+            logger.warning("Connecting with Entra ID user (oid): %s", entra_object_id)
+
+            redis_client = _get_redis_connection(
+                host_name=host_name,
+                port=port,
+                password=access_token,
+                ssl=True,
+                username=entra_object_id
+            )
+
+            logger.warning("Successfully connected to Redis at %s:%s", host_name, port)
+
+            # Test the connection with a PING command
+            success, message = _test_redis_connection_with_ping(redis_client)
+
+            if success:
+                result['connectionStatus'] = 'Success'
+                result['message'] = message
+            else:
+                result['connectionStatus'] = 'Failed'
+                result['message'] = message
+
+        except ImportError as ie:
+            result['connectionStatus'] = 'Failed'
+            result['message'] = (f"Required package not installed: {str(ie)}. "
+                                 "Please install 'redis' and 'PyJWT' packages.")
+        except Exception as e:  # pylint: disable=broad-except
+            result['connectionStatus'] = 'Failed'
+            result['message'] = f'Entra authentication failed: {str(e)}'
+
+    elif auth == 'access-key':
+        # Get access keys for the database
+        try:
+            # Use provided access key or retrieve from Azure
+            if not access_key:
+                if not resource_group_name or not cluster_name:
+                    raise ValidationError("--access-key is required when using --host without --resource-group and --cluster-name.")
+
+                keys = _DatabaseListKey(cli_ctx=cmd.cli_ctx)(command_args={
+                    "cluster_name": cluster_name,
+                    "resource_group": resource_group_name,
+                    "database_name": database_name})
+
+                if keys:
+                    access_key = keys.get('primaryKey') or keys.get('secondaryKey')
+
+            if not access_key:
+                result['connectionStatus'] = 'Failed'
+                result['message'] = ('Access keys authentication may be disabled. '
+                                     'Enable access keys authentication or use Entra authentication.')
+                return result
+
+            # Create Redis connection with the access key
+            redis_client = _get_redis_connection(
+                host_name=host_name,
+                port=port,
+                password=access_key,
+                ssl=True
+            )
+
+            logger.warning("Successfully connected to Redis at %s:%s", host_name, port)
+
+            # Test the connection with a PING command
+            success, message = _test_redis_connection_with_ping(redis_client)
+
+            if success:
+                result['connectionStatus'] = 'Success'
+                result['message'] = message
+            else:
+                result['connectionStatus'] = 'Failed'
+                result['message'] = message
+
+        except ImportError as ie:
+            result['connectionStatus'] = 'Failed'
+            result['message'] = f"Required package not installed: {str(ie)}. Please install 'redis' package."
+        except Exception as e:  # pylint: disable=broad-except
+            result['connectionStatus'] = 'Failed'
+            result['message'] = f'Failed to connect with access key: {str(e)}'
+
+    # Raise error on connection failure to return non-zero exit code
+    if result['connectionStatus'] == 'Failed':
+        from azure.cli.core.azclierror import CLIError
+        raise CLIError(result['message'])
+
+    return result
+
+
 class DatabaseFlush(_DatabaseFlush):
 
     @classmethod