add test-connection command

Author: yugangw-msft

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,13 @@ 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)
+        c.argument('cluster_name', options_list=['--cluster-name', '--name', '-n'], type=str, help='The name of the '
+                   'RedisEnterprise cluster.', id_part='name')
+        c.argument('auth', arg_type=get_enum_type(['entra', 'access-key']), help='The authentication method to use '
+                   'for the connection test. Allowed values: entra, access-key.')
+
 
 class AddPersistence(argparse.Action):
     def __call__(self, parser, namespace, values, option_string=None):

src/redisenterprise/azext_redisenterprise/commands.py

@@ -16,6 +16,7 @@ def load_command_table(self, _):  # pylint: disable=unused-argument
         g.custom_command('create', 'redisenterprise_create', supports_no_wait=True)
         g.custom_command('list', 'redisenterprise_list')
         g.custom_show_command('show', 'redisenterprise_show')
+        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

@@ -27,10 +27,243 @@
 from azure.cli.core.azclierror import (
     MutuallyExclusiveArgumentError,
 )
+from azure.cli.core.azclierror import ValidationError
 
 logger = get_logger(__name__)
 
 
+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_write(redis_client, test_key_prefix="az_cli_test_connection"):
+    """
+    Test Redis connection by writing and reading a test key.
+
+    :param redis_client: The Redis client instance.
+    :param test_key_prefix: Prefix for the test key.
+    :return: Tuple of (success: bool, message: str, steps: list).
+    """
+    import uuid
+    import time
+
+    test_logger = get_logger(__name__)
+    test_key = f"{test_key_prefix}:{uuid.uuid4()}"
+    test_value = f"test_value_{int(time.time())}"
+    steps = []
+
+    try:
+        # Step 1: Write a test key
+        test_logger.warning("Step 1: Writing test key '%s' with value '%s'...", test_key, test_value)
+        redis_client.set(test_key, test_value, ex=60)  # Expire in 60 seconds
+        steps.append({'step': 1, 'action': 'write', 'status': 'success', 'key': test_key, 'value': test_value,
+                      'message': f"Successfully wrote key '{test_key}'"})
+        test_logger.warning("Step 1: Successfully wrote test key")
+
+        # Step 2: Read the test key back
+        test_logger.warning("Step 2: Reading test key '%s'...", test_key)
+        retrieved_value = redis_client.get(test_key)
+        steps.append({'step': 2, 'action': 'read', 'status': 'success', 'key': test_key, 'value': retrieved_value,
+                      'message': f"Successfully read key '{test_key}', value: '{retrieved_value}'"})
+        test_logger.warning("Step 2: Successfully read test key, value: '%s'", retrieved_value)
+
+        # Step 3: Verify the value
+        test_logger.warning("Step 3: Verifying value matches...")
+        if retrieved_value == test_value:
+            steps.append({'step': 3, 'action': 'verify', 'status': 'success',
+                          'message': 'Value verification passed'})
+            test_logger.warning("Step 3: Value verification passed")
+        else:
+            steps.append({'step': 3, 'action': 'verify', 'status': 'failed',
+                          'message': f"Value mismatch: expected '{test_value}', got '{retrieved_value}'"})
+            test_logger.warning("Step 3: Value verification failed")
+            return False, f"Value mismatch: expected '{test_value}', got '{retrieved_value}'.", steps
+
+        # Step 4: Delete the test key
+        test_logger.warning("Step 4: Deleting test key '%s'...", test_key)
+        redis_client.delete(test_key)
+        steps.append({'step': 4, 'action': 'delete', 'status': 'success', 'key': test_key,
+                      'message': f"Successfully deleted key '{test_key}'"})
+        test_logger.warning("Step 4: Successfully deleted test key")
+
+        return True, "Successfully connected and verified write/read/delete operations.", steps
+
+    except Exception as e:  # pylint: disable=broad-except
+        error_msg = f"Failed during test operation: {str(e)}"
+        test_logger.warning("Error: %s", error_msg)
+        steps.append({'step': len(steps) + 1, 'action': 'error', 'status': 'failed', 'message': error_msg})
+        return False, error_msg, steps
+
+
+def redisenterprise_test_connection(cmd,
+                                    resource_group_name,
+                                    cluster_name,
+                                    auth):
+    """
+    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.
+    :param cluster_name: The name of the Redis Enterprise cluster.
+    :param auth: The authentication method to use ('entra' or 'access-key').
+    :return: Connection test result.
+    """
+    # 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}'.")
+
+    # Get the list of databases in the cluster
+    databases = list(_DatabaseList(cli_ctx=cmd.cli_ctx)(command_args={
+        "cluster_name": cluster_name,
+        "resource_group": resource_group_name}))
+
+    if not databases:
+        raise ValidationError(f"No databases found in cluster '{cluster_name}'. "
+                              "Please create a database before testing the connection.")
+
+    # Use the first database
+    first_database = databases[0]
+    database_name = first_database.get('name', 'default')
+    port = first_database.get('port', 10000)
+
+    logger.warning("Using database: %s", database_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})
+            user_name = decoded_token.get('oid', decoded_token.get('sub', 'default'))
+
+            logger.warning("Connecting with Entra ID user (oid): %s", user_name)
+
+            redis_client = _get_redis_connection(
+                host_name=host_name,
+                port=port,
+                password=access_token,
+                ssl=True,
+                username=user_name
+            )
+
+            logger.warning("Successfully connected to Redis at %s:%s", host_name, port)
+
+            # Test the connection with a write operation
+            success, message, _ = _test_redis_connection_with_write(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:
+            keys = _DatabaseListKey(cli_ctx=cmd.cli_ctx)(command_args={
+                "cluster_name": cluster_name,
+                "resource_group": resource_group_name,
+                "database_name": database_name})
+
+            access_key = None
+            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 write operation
+            success, message, _ = _test_redis_connection_with_write(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)}'
+
+    return result
+
+
 class DatabaseFlush(_DatabaseFlush):
 
     @classmethod

src/redisenterprise/setup.py

@@ -26,7 +26,7 @@
     'License :: OSI Approved :: MIT License',
 ]
 
-DEPENDENCIES = []
+DEPENDENCIES = ['redis~=7.1.0']
 
 with open('README.md', 'r', encoding='utf-8') as f:
     README = f.read()