Fix Redis MI auth: call super().on_connect() BEFORE AUTH command, add debug logging

Author: Chen, Vivien

application/single_app/app.py

@@ -136,22 +136,32 @@ def __init__(self, *args, **kwargs):
                             super().__init__(*args, **kwargs)
 
                         def on_connect(self):
-                            super().on_connect()
                             if hasattr(ManagedIdentityConnection, '_credential') and ManagedIdentityConnection._credential:
                                 # Get fresh token and extract username
-                                token = ManagedIdentityConnection._credential.get_token(ManagedIdentityConnection._token_scope)
-                                import base64, json
-                                token_parts = token.token.split('.')
-                                if len(token_parts) > 1:
-                                    payload = token_parts[1] + '=' * (4 - len(token_parts[1]) % 4)
-                                    token_data = json.loads(base64.urlsafe_b64decode(payload))
-                                    username = token_data.get('oid', '')
-                                else:
-                                    username = ''
-                                # Authenticate with fresh token
-                                if username and token.token:
-                                    self.send_command('AUTH', username, token.token)
-                                    self.read_response()
+                                try:
+                                    token = ManagedIdentityConnection._credential.get_token(ManagedIdentityConnection._token_scope)
+                                    import base64, json
+                                    token_parts = token.token.split('.')
+                                    if len(token_parts) > 1:
+                                        payload = token_parts[1] + '=' * (4 - len(token_parts[1]) % 4)
+                                        token_data = json.loads(base64.urlsafe_b64decode(payload))
+                                        username = token_data.get('oid', '')
+                                    else:
+                                        username = ''
+                                    
+                                    # Call parent to establish connection
+                                    super().on_connect()
+                                    
+                                    # Authenticate with fresh token after connection established
+                                    if username and token.token:
+                                        self.send_command('AUTH', username, token.token)
+                                        auth_response = self.read_response()
+                                        print(f"Redis AUTH response (startup): {auth_response}")
+                                except Exception as e:
+                                    print(f"Redis MI auth error in on_connect (startup): {e}")
+                                    raise
+                            else:
+                                super().on_connect()
 
                     # Create connection pool with custom connection class
                     pool = ConnectionPool(
@@ -253,22 +263,32 @@ def __init__(self, *args, **kwargs):
                             super().__init__(*args, **kwargs)
 
                         def on_connect(self):
-                            super().on_connect()
                             if hasattr(ManagedIdentityConnection, '_credential') and ManagedIdentityConnection._credential:
                                 # Get fresh token and extract username
-                                token = ManagedIdentityConnection._credential.get_token(ManagedIdentityConnection._token_scope)
-                                import base64, json
-                                token_parts = token.token.split('.')
-                                if len(token_parts) > 1:
-                                    payload = token_parts[1] + '=' * (4 - len(token_parts[1]) % 4)
-                                    token_data = json.loads(base64.urlsafe_b64decode(payload))
-                                    username = token_data.get('oid', '')
-                                else:
-                                    username = ''
-                                # Authenticate with fresh token
-                                if username and token.token:
-                                    self.send_command('AUTH', username, token.token)
-                                    self.read_response()
+                                try:
+                                    token = ManagedIdentityConnection._credential.get_token(ManagedIdentityConnection._token_scope)
+                                    import base64, json
+                                    token_parts = token.token.split('.')
+                                    if len(token_parts) > 1:
+                                        payload = token_parts[1] + '=' * (4 - len(token_parts[1]) % 4)
+                                        token_data = json.loads(base64.urlsafe_b64decode(payload))
+                                        username = token_data.get('oid', '')
+                                    else:
+                                        username = ''
+                                    
+                                    # Call parent to establish connection
+                                    super().on_connect()
+                                    
+                                    # Authenticate with fresh token after connection established
+                                    if username and token.token:
+                                        self.send_command('AUTH', username, token.token)
+                                        auth_response = self.read_response()
+                                        print(f"Redis AUTH response (before-first-request): {auth_response}")
+                                except Exception as e:
+                                    print(f"Redis MI auth error in on_connect (before-first-request): {e}")
+                                    raise
+                            else:
+                                super().on_connect()
 
                     # Create connection pool with custom connection class
                     pool = ConnectionPool(