enable user and IDP reset in Keycloak via flag (#7950)

Co-authored-by: pmossman <[email protected]>
duffelhq · Aug 10, 2023 · 8ee9fd2 · 8ee9fd2
1 parent 40ce36b
commit 8ee9fd2
Show file tree

Hide file tree

Showing 11 changed files with 119 additions and 9 deletions.
diff --git a/...mmons-auth/src/main/java/io/airbyte/commons/auth/config/AirbyteKeycloakConfiguration.java b/...mmons-auth/src/main/java/io/airbyte/commons/auth/config/AirbyteKeycloakConfiguration.java
@@ -35,6 +35,7 @@ public class AirbyteKeycloakConfiguration {
   String accountClientId;
   String username;
   String password;
+  Boolean resetRealm;
 
   public String getKeycloakUserInfoEndpoint() {
     final String hostWithoutTrailingSlash = host.endsWith("/") ? host.substring(0, host.length() - 1) : host;

diff --git a/airbyte-keycloak-setup/src/main/java/io/airbyte/keycloak/setup/IdentityProvidersCreator.java b/airbyte-keycloak-setup/src/main/java/io/airbyte/keycloak/setup/IdentityProvidersCreator.java
@@ -79,4 +79,16 @@ private void createIdp(final RealmResource keycloakRealm, final IdentityProvider
     }
   }
 
+  /**
+   * Delete existing identity providers and re-create them. This is useful when we want to update the
+   * configuration of an existing identity provider.
+   */
+  public void resetIdentityProviders(RealmResource realmResource) {
+    realmResource.identityProviders().findAll().forEach(idpRepresentation -> {
+      realmResource.identityProviders().get(idpRepresentation.getInternalId()).remove();
+    });
+
+    createIdps(realmResource);
+  }
+
 }
diff --git a/airbyte-keycloak-setup/src/main/java/io/airbyte/keycloak/setup/KeycloakServer.java b/airbyte-keycloak-setup/src/main/java/io/airbyte/keycloak/setup/KeycloakServer.java
@@ -47,6 +47,14 @@ public KeycloakServer(final KeycloakAdminClientProvider keycloakAdminClientProvi
   public void createAirbyteRealm() {
     if (doesRealmExist()) {
       log.info("Realm {} already exists.", keycloakConfiguration.getAirbyteRealm());
+
+      if (keycloakConfiguration.getResetRealm()) {
+        final RealmResource airbyteRealm = keycloakAdminClient.realm(keycloakConfiguration.getAirbyteRealm());
+        userCreator.resetUser(airbyteRealm);
+        identityProvidersCreator.resetIdentityProviders(airbyteRealm);
+        log.info("Reset user and identity providers for realm {}.", keycloakConfiguration.getAirbyteRealm());
+      }
+
       return;
     }
     createRealm();
@@ -61,12 +69,12 @@ private boolean doesRealmExist() {
 
   private void createRealm() {
     log.info("Creating realm {}...", keycloakConfiguration.getAirbyteRealm());
-    RealmRepresentation airbyteRealmRepresentation = buildRealmRepresentation();
+    final RealmRepresentation airbyteRealmRepresentation = buildRealmRepresentation();
     keycloakAdminClient.realms().create(airbyteRealmRepresentation);
   }
 
   private RealmRepresentation buildRealmRepresentation() {
-    RealmRepresentation airbyteRealmRepresentation = new RealmRepresentation();
+    final RealmRepresentation airbyteRealmRepresentation = new RealmRepresentation();
     airbyteRealmRepresentation.setRealm(keycloakConfiguration.getAirbyteRealm());
     airbyteRealmRepresentation.setEnabled(true);
     airbyteRealmRepresentation.setLoginTheme("airbyte-keycloak-theme");

diff --git a/airbyte-keycloak-setup/src/main/java/io/airbyte/keycloak/setup/UserCreator.java b/airbyte-keycloak-setup/src/main/java/io/airbyte/keycloak/setup/UserCreator.java
@@ -10,6 +10,7 @@
 import javax.ws.rs.core.Response;
 import lombok.extern.slf4j.Slf4j;
 import org.keycloak.admin.client.resource.RealmResource;
+import org.keycloak.admin.client.resource.UsersResource;
 import org.keycloak.representations.idm.CredentialRepresentation;
 import org.keycloak.representations.idm.UserRepresentation;
 
@@ -59,4 +60,15 @@ CredentialRepresentation createCredentialRepresentation() {
     return password;
   }
 
+  /**
+   * This method resets the user by deleting all users and re-creating them.
+   */
+  public void resetUser(final RealmResource realmResource) {
+    UsersResource usersResource = realmResource.users();
+
+    usersResource.list().forEach(user -> usersResource.delete(user.getId()));
+
+    createUser(realmResource);
+  }
+
 }
diff --git a/airbyte-keycloak-setup/src/main/resources/application.yml b/airbyte-keycloak-setup/src/main/resources/application.yml
@@ -17,3 +17,4 @@ airbyte:
     account-client-id: ${KEYCLOAK_ACCOUNT_CLIENT_ID:account}
     username: ${KEYCLOAK_ADMIN_USER:}
     password: ${KEYCLOAK_ADMIN_PASSWORD:}
+    reset-realm: ${KEYCLOAK_SETUP_RESET_REALM:false}
diff --git a/...-keycloak-setup/src/test/java/io/airbyte/keycloak/setup/IdentityProvidersCreatorTest.java b/...-keycloak-setup/src/test/java/io/airbyte/keycloak/setup/IdentityProvidersCreatorTest.java
@@ -6,18 +6,21 @@
 
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import io.airbyte.commons.auth.config.IdentityProviderConfiguration;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import javax.ws.rs.core.Response;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.keycloak.admin.client.resource.IdentityProviderResource;
 import org.keycloak.admin.client.resource.IdentityProvidersResource;
 import org.keycloak.admin.client.resource.RealmResource;
 import org.keycloak.representations.idm.IdentityProviderRepresentation;
@@ -38,6 +41,8 @@ class IdentityProvidersCreatorTest {
   private IdentityProviderConfiguration identityProviderConfiguration;
   @Mock
   private IdentityProvidersResource identityProvidersResource;
+  @Mock
+  private IdentityProviderRepresentation idpRepresentation;
   @InjectMocks
   private IdentityProvidersCreator identityProvidersCreator;
 
@@ -50,8 +55,10 @@ void setUp() {
   @Test
   void testCreateIdps() {
     when(realmResource.identityProviders()).thenReturn(identityProvidersResource);
-    when(identityProvidersResource.create(any(IdentityProviderRepresentation.class))).thenReturn(response);
+
+    Response response = mock(Response.class);
     when(response.getStatus()).thenReturn(201);
+    when(identityProvidersResource.create(any(IdentityProviderRepresentation.class))).thenReturn(response);
 
     identityProvidersCreator.createIdps(realmResource);
 
@@ -95,4 +102,23 @@ void testCreateIdps_Failure() {
     });
   }
 
+  @Test
+  void testResetIdentityProviders() {
+    IdentityProviderRepresentation identityProviderRepresentation = mock(IdentityProviderRepresentation.class);
+    IdentityProviderResource identityProvider = mock(IdentityProviderResource.class);
+    Response response = mock(Response.class);
+
+    when(realmResource.identityProviders()).thenReturn(identityProvidersResource);
+    when(identityProvidersResource.findAll()).thenReturn(Arrays.asList(identityProviderRepresentation));
+    when(identityProvidersResource.get(identityProviderRepresentation.getInternalId())).thenReturn(identityProvider);
+    when(identityProvidersResource.create(any(IdentityProviderRepresentation.class))).thenReturn(response);
+    when(response.getStatus()).thenReturn(201);
+
+    identityProvidersCreator.resetIdentityProviders(realmResource);
+
+    verify(identityProvidersResource).findAll();
+    verify(identityProvidersResource).get(identityProviderRepresentation.getInternalId());
+    verify(identityProvider).remove();
+  }
+
 }
diff --git a/airbyte-keycloak-setup/src/test/java/io/airbyte/keycloak/setup/KeycloakServerTest.java b/airbyte-keycloak-setup/src/test/java/io/airbyte/keycloak/setup/KeycloakServerTest.java
@@ -83,7 +83,7 @@ void testCreateAirbyteRealm() {
 
   @Test
   void testCreateAirbyteRealmWhenRealmAlreadyExists() {
-    RealmRepresentation existingRealm = new RealmRepresentation();
+    final RealmRepresentation existingRealm = new RealmRepresentation();
     existingRealm.setRealm(REALM_NAME);
 
     when(realmsResource.findAll()).thenReturn(Collections.singletonList(existingRealm));
@@ -103,13 +103,30 @@ void testCreateAirbyteRealmWhenRealmAlreadyExists() {
   void testBuildRealmRepresentation() {
     keycloakServer.createAirbyteRealm();
 
-    ArgumentCaptor<RealmRepresentation> realmRepresentationCaptor = ArgumentCaptor.forClass(RealmRepresentation.class);
+    final ArgumentCaptor<RealmRepresentation> realmRepresentationCaptor = ArgumentCaptor.forClass(RealmRepresentation.class);
     verify(realmsResource).create(realmRepresentationCaptor.capture());
 
-    RealmRepresentation createdRealm = realmRepresentationCaptor.getValue();
+    final RealmRepresentation createdRealm = realmRepresentationCaptor.getValue();
     assertEquals(REALM_NAME, createdRealm.getRealm());
     assertTrue(createdRealm.isEnabled());
     assertEquals("airbyte-keycloak-theme", createdRealm.getLoginTheme());
   }
 
+  @Test
+  void testCreateAirbyteRealmWhenRealmAlreadyExistsAndResetUserIsTrue() {
+    final RealmRepresentation existingRealm = new RealmRepresentation();
+    existingRealm.setRealm(REALM_NAME);
+
+    when(realmsResource.findAll()).thenReturn(Collections.singletonList(existingRealm));
+    when(keycloakConfiguration.getAirbyteRealm()).thenReturn(REALM_NAME);
+    when(keycloakConfiguration.getResetRealm()).thenReturn(true);
+
+    keycloakServer.createAirbyteRealm();
+
+    verify(realmsResource, times(1)).findAll();
+    verify(realmsResource, times(0)).create(any());
+    verify(userCreator, times(1)).resetUser(any());
+    verify(identityProvidersCreator, times(1)).resetIdentityProviders(any());
+  }
+
 }
diff --git a/airbyte-keycloak-setup/src/test/java/io/airbyte/keycloak/setup/KeycloakSetupTest.java b/airbyte-keycloak-setup/src/test/java/io/airbyte/keycloak/setup/KeycloakSetupTest.java
@@ -59,7 +59,7 @@ void testRun() {
 
   @Test
   void testRunThrowsException() {
-    String keycloakUrl = "http://localhost:8180/auth";
+    final String keycloakUrl = "http://localhost:8180/auth";
     when(keycloakServer.getKeycloakServerUrl()).thenReturn(keycloakUrl);
     when(httpClient.toBlocking().exchange(any(HttpRequest.class), eq(String.class)))
         .thenThrow(new HttpClientResponseException("Error", HttpResponse.serverError()));

diff --git a/airbyte-keycloak-setup/src/test/java/io/airbyte/keycloak/setup/UserCreatorTest.java b/airbyte-keycloak-setup/src/test/java/io/airbyte/keycloak/setup/UserCreatorTest.java
@@ -6,15 +6,19 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import io.airbyte.commons.auth.config.InitialUserConfiguration;
+import java.util.Collections;
 import javax.ws.rs.core.Response;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.keycloak.admin.client.resource.RealmResource;
+import org.keycloak.admin.client.resource.UserResource;
 import org.keycloak.admin.client.resource.UsersResource;
 import org.keycloak.representations.idm.CredentialRepresentation;
 import org.keycloak.representations.idm.UserRepresentation;
@@ -36,6 +40,8 @@ class UserCreatorTest {
   @Mock
   private UsersResource usersResource;
   @Mock
+  private UserResource userResource;
+  @Mock
   private Response response;
 
   @BeforeEach
@@ -51,6 +57,10 @@ void setUp() {
     when(realmResource.users()).thenReturn(usersResource);
     when(usersResource.create(any(UserRepresentation.class))).thenReturn(response);
 
+    when(usersResource.get(anyString())).thenReturn(userResource);
+    when(usersResource.search(anyString(), anyInt(), anyInt())).thenReturn(Collections.singletonList(new UserRepresentation()));
+    when(response.getStatusInfo()).thenReturn(Response.Status.OK);
+
     userCreator = new UserCreator(initialUserConfiguration);
   }
 
@@ -89,4 +99,19 @@ void testCreateCredentialRepresentation() {
     assertEquals(PASSWORD, credentialRepresentation.getValue());
   }
 
+  @Test
+  void testResetUser() {
+    UserRepresentation userRepresentation = new UserRepresentation();
+    userRepresentation.setId("id1");
+    when(usersResource.list()).thenReturn(Collections.singletonList(userRepresentation));
+
+    when(response.getStatus()).thenReturn(201);
+
+    userCreator.resetUser(realmResource);
+
+    verify(usersResource).list();
+    verify(usersResource).delete("id1");
+    verify(usersResource).create(any(UserRepresentation.class));
+  }
+
 }
diff --git a/charts/airbyte-keycloak-setup/templates/pod.yaml b/charts/airbyte-keycloak-setup/templates/pod.yaml
@@ -63,6 +63,8 @@ spec:
             configMapKeyRef:
               name: {{ .Values.global.configMapName | default (printf "%s-airbyte-env" .Release.Name) }}
               key: KEYCLOAK_INTERNAL_HOST
+        - name: KEYCLOAK_SETUP_RESET_REALM
+          value: {{ .Values.resetRealm | default "false" | quote }}
         # Values from secret
         {{- if .Values.secrets }}
         {{- range $k, $v := .Values.secrets }}

diff --git a/tools/bin/install_airbyte_pro_on_helm.sh b/tools/bin/install_airbyte_pro_on_helm.sh
@@ -23,6 +23,12 @@ else
   set_dev_image_tag_if_true=""
 fi
 
+if [ "$RESET_KEYCLOAK_REALM" == "true" ]; then
+  set_reset_realm_if_true="--set keycloak-setup.resetRealm=true"
+else
+  set_reset_realm_if_true=""
+fi
+
 # run this script with LOCAL=true to install using the local version of the chart.
 # Otherwise, it will install using the latest version of the chart from the airbyte/airbyte repo
 if [ "$LOCAL" == "true" ]; then
@@ -32,13 +38,13 @@ if [ "$LOCAL" == "true" ]; then
   helm dep update
 
   # additional arguments to this script are appended at the end of the command, so that further customization can be done
-  helm upgrade --install "$airbyte_pro_release_name" . $set_dev_image_tag_if_true --set-file airbyteYml="$airbyte_yml_file_path" --values "$airbyte_pro_values_yml_file_path" "$@"
+  helm upgrade --install "$airbyte_pro_release_name" . $set_dev_image_tag_if_true $set_reset_realm_if_true --set-file airbyteYml="$airbyte_yml_file_path" --values "$airbyte_pro_values_yml_file_path" "$@"
   popd
 else
   airbyte_chart="airbyte/airbyte"
 
   # additional arguments to this script are appended at the end of the command, so that further customization can be done
-  helm upgrade --install "$airbyte_pro_release_name" "$airbyte_chart" $set_dev_image_tag_if_true --set-file airbyteYml="$airbyte_yml_file_path" --values "$airbyte_pro_values_yml_file_path" "$@"
+  helm upgrade --install "$airbyte_pro_release_name" "$airbyte_chart" $set_dev_image_tag_if_true $set_reset_realm_if_true --set-file airbyteYml="$airbyte_yml_file_path" --values "$airbyte_pro_values_yml_file_path" "$@"
 fi
 
 echo "Airbyte Pro installation complete!"