Fix JWT reset after snapshot restore

Author: tomach

crate/operator/operations.py

@@ -1165,6 +1165,7 @@ async def set_user_jwt(
     :param namespace: The Kubernetes namespace of the CrateDB cluster.
     :param name: The CrateDB custom resource name defining the CrateDB cluster.
     :param username: The name of the user the JWT properties should be set for.
+    :param logger: Logger for operation tracking.
     """
     cratedb = await get_cratedb_resource(namespace, name)
     await cursor.execute(
@@ -1175,36 +1176,67 @@ async def set_user_jwt(
 
     username_ident = quote_ident(username, cursor._impl)
     iss = cratedb["spec"].get("grandCentral", {}).get("jwkUrl")
-    logger.info("Setting JWT auth properties for user '%s'", username)
-    if user_exists and iss:
-        query = (
-            f"ALTER USER {username_ident} SET "
-            f"""(jwt = {{"iss" = '{iss}', "username" = '{username}', """
-            f""""aud" = '{name}'}})"""
-        )
-        pod_name = get_crash_pod_name(cratedb, name)
-        scheme = get_crash_scheme(cratedb)
-        logger.info("... executing query: %s", query)
-        result = await execute_sql(
-            namespace=namespace,
-            name=name,
-            pod_name=pod_name,
-            scheme=scheme,
-            sql=query,
-            args=None,
-            logger=logger,
+
+    if not (user_exists and iss):
+        return
+
+    pod_name = get_crash_pod_name(cratedb, name)
+    scheme = get_crash_scheme(cratedb)
+
+    # Step 1: Reset JWT to NULL first
+    # This prevents RoleAlreadyExistsException when restoring snapshots
+    # where the user might have JWT properties with different aud/iss
+    logger.info("Resetting JWT auth properties for user '%s'", username)
+    reset_query = f"ALTER USER {username_ident} SET (jwt = NULL)"
+    logger.info("... executing query: %s", reset_query)
+
+    reset_result = await execute_sql(
+        namespace=namespace,
+        name=name,
+        pod_name=pod_name,
+        scheme=scheme,
+        sql=reset_query,
+        args=None,
+        logger=logger,
+    )
+    logger.info("... result: %s", reset_result)
+
+    if (reset_result.rowcount or 0) > 0:
+        logger.info("... JWT reset successful")
+    else:
+        logger.info(
+            "... JWT reset had no effect (might not have been set). %s", reset_result
         )
-        logger.info("... result: %s", result)
-        if (result.rowcount or 0) > 0:
-            logger.info("... success")
-        else:
-            logger.info("... error. %s", result)
-            # Continue if the same JWT properties are already set
-            if (
-                not result.error_message
-                or "RoleAlreadyExistsException" not in result.error_message
-            ):
-                raise kopf.TemporaryError(delay=config.BOOTSTRAP_RETRY_DELAY)
+
+    # Step 2: Set new JWT properties with current cluster's aud
+    logger.info("Setting JWT auth properties for user '%s'", username)
+    set_query = (
+        f"ALTER USER {username_ident} SET "
+        f"""(jwt = {{"iss" = '{iss}', "username" = '{username}', """
+        f""""aud" = '{name}'}})"""
+    )
+    logger.info("... executing query: %s", set_query)
+
+    result = await execute_sql(
+        namespace=namespace,
+        name=name,
+        pod_name=pod_name,
+        scheme=scheme,
+        sql=set_query,
+        args=None,
+        logger=logger,
+    )
+    logger.info("... result: %s", result)
+
+    if (result.rowcount or 0) > 0:
+        logger.info("... success")
+    else:
+        logger.info("... error. %s", result)
+        if (
+            not result.error_message
+            or "RoleAlreadyExistsException" not in result.error_message
+        ):
+            raise kopf.TemporaryError(delay=config.BOOTSTRAP_RETRY_DELAY)
 
 
 class StartClusterSubHandler(StateBasedSubHandler):