keycloak-aplcache

diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/session/JpaUserSessionPersisterProvider.java b/model/jpa/src/main/java/org/keycloak/models/jpa/session/JpaUserSessionPersisterProvider.java
index a4c02de..587ed7e 100644
--- a/model/jpa/src/main/java/org/keycloak/models/jpa/session/JpaUserSessionPersisterProvider.java
+++ b/model/jpa/src/main/java/org/keycloak/models/jpa/session/JpaUserSessionPersisterProvider.java
@@ -30,6 +30,7 @@ import org.keycloak.models.session.PersistentClientSessionModel;
 import org.keycloak.models.session.PersistentUserSessionAdapter;
 import org.keycloak.models.session.PersistentUserSessionModel;
 import org.keycloak.models.session.UserSessionPersisterProvider;
+import org.keycloak.storage.StorageId;
 
 import javax.persistence.EntityManager;
 import javax.persistence.Query;
@@ -78,7 +79,17 @@ public class JpaUserSessionPersisterProvider implements UserSessionPersisterProv
         PersistentClientSessionModel model = adapter.getUpdatedModel();
 
         PersistentClientSessionEntity entity = new PersistentClientSessionEntity();
-        entity.setClientId(clientSession.getClient().getId());
+        StorageId clientStorageId = new StorageId(clientSession.getClient().getId());
+        if (clientStorageId.isLocal()) {
+            entity.setClientId(clientStorageId.getId());
+            entity.setClientStorageProvider(PersistentClientSessionEntity.LOCAL);
+            entity.setExternalClientId(PersistentClientSessionEntity.LOCAL);
+
+        } else {
+            entity.setClientId(PersistentClientSessionEntity.EXTERNAL);
+            entity.setClientStorageProvider(clientStorageId.getProviderId());
+            entity.setExternalClientId(clientStorageId.getExternalId());
+        }
         entity.setTimestamp(clientSession.getTimestamp());
         String offlineStr = offlineToString(offline);
         entity.setOffline(offlineStr);
@@ -127,7 +138,18 @@ public class JpaUserSessionPersisterProvider implements UserSessionPersisterProv
     @Override
     public void removeClientSession(String userSessionId, String clientUUID, boolean offline) {
         String offlineStr = offlineToString(offline);
-        PersistentClientSessionEntity sessionEntity = em.find(PersistentClientSessionEntity.class, new PersistentClientSessionEntity.Key(userSessionId, clientUUID, offlineStr));
+        StorageId clientStorageId = new StorageId(clientUUID);
+        String clientId = PersistentClientSessionEntity.EXTERNAL;
+        String clientStorageProvider = PersistentClientSessionEntity.LOCAL;
+        String externalId = PersistentClientSessionEntity.LOCAL;
+        if (clientStorageId.isLocal()) {
+            clientId = clientUUID;
+        } else {
+            clientStorageProvider = clientStorageId.getProviderId();
+            externalId = clientStorageId.getExternalId();
+
+        }
+        PersistentClientSessionEntity sessionEntity = em.find(PersistentClientSessionEntity.class, new PersistentClientSessionEntity.Key(userSessionId, clientId, clientStorageProvider, externalId, offlineStr));
         if (sessionEntity != null) {
             em.remove(sessionEntity);
 
@@ -168,7 +190,16 @@ public class JpaUserSessionPersisterProvider implements UserSessionPersisterProv
     }
 
     private void onClientRemoved(String clientUUID) {
-        int num = em.createNamedQuery("deleteClientSessionsByClient").setParameter("clientId", clientUUID).executeUpdate();
+        int num = 0;
+        StorageId clientStorageId = new StorageId(clientUUID);
+        if (clientStorageId.isLocal()) {
+            num = em.createNamedQuery("deleteClientSessionsByClient").setParameter("clientId", clientUUID).executeUpdate();
+        } else {
+            num = em.createNamedQuery("deleteClientSessionsByExternalClient")
+                    .setParameter("clientStorageProvider", clientStorageId.getProviderId())
+                    .setParameter("externalClientId", clientStorageId.getExternalId())
+                    .executeUpdate();
+        }
         num = em.createNamedQuery("deleteDetachedUserSessions").executeUpdate();
     }
 
@@ -282,10 +313,14 @@ public class JpaUserSessionPersisterProvider implements UserSessionPersisterProv
     }
 
     private PersistentAuthenticatedClientSessionAdapter toAdapter(RealmModel realm, PersistentUserSessionAdapter userSession, PersistentClientSessionEntity entity) {
-        ClientModel client = realm.getClientById(entity.getClientId());
+        String clientId = entity.getClientId();
+        if (!entity.getExternalClientId().equals("local")) {
+            clientId = new StorageId(entity.getClientId(), entity.getExternalClientId()).getId();
+        }
+        ClientModel client = realm.getClientById(clientId);
 
         PersistentClientSessionModel model = new PersistentClientSessionModel();
-        model.setClientId(entity.getClientId());
+        model.setClientId(clientId);
         model.setUserSessionId(userSession.getId());
         model.setUserId(userSession.getUserId());
         model.setTimestamp(entity.getTimestamp());

diff --git a/model/jpa/src/main/java/org/keycloak/models/jpa/session/PersistentClientSessionEntity.java b/model/jpa/src/main/java/org/keycloak/models/jpa/session/PersistentClientSessionEntity.java
index 3ae17b2..44c3188 100644
--- a/model/jpa/src/main/java/org/keycloak/models/jpa/session/PersistentClientSessionEntity.java
+++ b/model/jpa/src/main/java/org/keycloak/models/jpa/session/PersistentClientSessionEntity.java
@@ -32,6 +32,8 @@ import java.io.Serializable;
 @NamedQueries({
         @NamedQuery(name="deleteClientSessionsByRealm", query="delete from PersistentClientSessionEntity sess where sess.userSessionId IN (select u.userSessionId from PersistentUserSessionEntity u where u.realmId = :realmId)"),
         @NamedQuery(name="deleteClientSessionsByClient", query="delete from PersistentClientSessionEntity sess where sess.clientId = :clientId"),
+        @NamedQuery(name="deleteClientSessionsByExternalClient", query="delete from PersistentClientSessionEntity sess where sess.clientStorageProvider = :clientStorageProvider and sess.externalClientId = :externalClientId"),
+        @NamedQuery(name="deleteClientSessionsByClientStorageProvider", query="delete from PersistentClientSessionEntity sess where sess.clientStorageProvider = :clientStorageProvider"),
         @NamedQuery(name="deleteClientSessionsByUser", query="delete from PersistentClientSessionEntity sess where sess.userSessionId IN (select u.userSessionId from PersistentUserSessionEntity u where u.userId = :userId)"),
         @NamedQuery(name="deleteClientSessionsByUserSession", query="delete from PersistentClientSessionEntity sess where sess.userSessionId = :userSessionId and sess.offline = :offline"),
         @NamedQuery(name="deleteDetachedClientSessions", query="delete from PersistentClientSessionEntity sess where NOT EXISTS (select u.userSessionId from PersistentUserSessionEntity u where u.userSessionId = sess.userSessionId )"),
@@ -44,6 +46,8 @@ import java.io.Serializable;
 @IdClass(PersistentClientSessionEntity.Key.class)
 public class PersistentClientSessionEntity {
 
+    public static final String LOCAL = "local";
+    public static final String EXTERNAL = "external";
     @Id
     @Column(name = "USER_SESSION_ID", length = 36)
     protected String userSessionId;
@@ -52,6 +56,14 @@ public class PersistentClientSessionEntity {
     @Column(name="CLIENT_ID", length = 36)
     protected String clientId;
 
+    @Id
+    @Column(name="CLIENT_STORAGE_PROVIDER", length = 36)
+    protected String clientStorageProvider;
+
+    @Id
+    @Column(name="EXTERNAL_CLIENT_ID", length = 255)
+    protected String externalClientId;
+
     @Column(name="TIMESTAMP")
     protected int timestamp;
 
@@ -78,6 +90,22 @@ public class PersistentClientSessionEntity {
         this.clientId = clientId;
     }
 
+    public String getClientStorageProvider() {
+        return clientStorageProvider;
+    }
+
+    public void setClientStorageProvider(String clientStorageProvider) {
+        this.clientStorageProvider = clientStorageProvider;
+    }
+
+    public String getExternalClientId() {
+        return externalClientId;
+    }
+
+    public void setExternalClientId(String externalClientId) {
+        this.externalClientId = externalClientId;
+    }
+
     public int getTimestamp() {
         return timestamp;
     }
@@ -107,15 +135,19 @@ public class PersistentClientSessionEntity {
         protected String userSessionId;
 
         protected String clientId;
+        protected String clientStorageProvider;
+        protected String externalClientId;
 
         protected String offline;
 
         public Key() {
         }
 
-        public Key(String userSessionId, String clientId, String offline) {
+        public Key(String userSessionId, String clientId, String clientStorageProvider, String externalClientId, String offline) {
             this.userSessionId = userSessionId;
             this.clientId = clientId;
+            this.externalClientId = externalClientId;
+            this.clientStorageProvider = clientStorageProvider;
             this.offline = offline;
         }
 
@@ -131,6 +163,14 @@ public class PersistentClientSessionEntity {
             return offline;
         }
 
+        public String getClientStorageProvider() {
+            return clientStorageProvider;
+        }
+
+        public String getExternalClientId() {
+            return externalClientId;
+        }
+
         @Override
         public boolean equals(Object o) {
             if (this == o) return true;
@@ -140,6 +180,8 @@ public class PersistentClientSessionEntity {
 
             if (this.userSessionId != null ? !this.userSessionId.equals(key.userSessionId) : key.userSessionId != null) return false;
             if (this.clientId != null ? !this.clientId.equals(key.clientId) : key.clientId != null) return false;
+            if (this.externalClientId != null ? !this.externalClientId.equals(key.clientId) : key.externalClientId != null) return false;
+            if (this.clientStorageProvider != null ? !this.clientStorageProvider.equals(key.clientId) : key.clientStorageProvider != null) return false;
             if (this.offline != null ? !this.offline.equals(key.offline) : key.offline != null) return false;
 
             return true;
@@ -149,6 +191,8 @@ public class PersistentClientSessionEntity {
         public int hashCode() {
             int result = this.userSessionId != null ? this.userSessionId.hashCode() : 0;
             result = 37 * result + (this.clientId != null ? this.clientId.hashCode() : 0);
+            result = 37 * result + (this.externalClientId != null ? this.externalClientId.hashCode() : 0);
+            result = 37 * result + (this.clientStorageProvider != null ? this.clientStorageProvider.hashCode() : 0);
             result = 31 * result + (this.offline != null ? this.offline.hashCode() : 0);
             return result;
         }

diff --git a/model/jpa/src/main/resources/META-INF/jpa-changelog-4.0.0.xml b/model/jpa/src/main/resources/META-INF/jpa-changelog-4.0.0.xml
index 09a1d80..3a81cba 100644
--- a/model/jpa/src/main/resources/META-INF/jpa-changelog-4.0.0.xml
+++ b/model/jpa/src/main/resources/META-INF/jpa-changelog-4.0.0.xml
@@ -68,14 +68,18 @@
             <column name="EXTERNAL_CLIENT_ID" type="VARCHAR(255)" />
         </createIndex>
 
-        <!-- Modify CLIENT_NODE_REGISTRATIONS -->
-        <dropForeignKeyConstraint constraintName="FK4129723BA992F594" baseTableName="CLIENT"/>
-        <modifyDataType tableName="CLIENT_NODE_REGISTRATIONS" columnName="CLIENT_ID" newDataType="VARCHAR(255)"/>
-
         <!-- Modify OFFLINE_CLIENT_SESSION -->
         <dropPrimaryKey tableName="OFFLINE_CLIENT_SESSION" constraintName="CONSTRAINT_OFFL_CL_SES_PK3"/>
+        <addColumn tableName="OFFLINE_CLIENT_SESSION">
+            <column name="CLIENT_STORAGE_PROVIDER" type="VARCHAR(36)" defaultValue="local">
+                <constraints nullable="false"/>
+            </column>
+             <column name="EXTERNAL_CLIENT_ID" type="VARCHAR(255)" defaultValue="local">
+                <constraints nullable="false"/>
+            </column>
+        </addColumn>
         <modifyDataType tableName="OFFLINE_CLIENT_SESSION" columnName="CLIENT_ID" newDataType="VARCHAR(255)"/>
-        <addPrimaryKey columnNames="USER_SESSION_ID,CLIENT_ID, OFFLINE_FLAG" constraintName="CONSTRAINT_OFFL_CL_SES_PK3" tableName="OFFLINE_CLIENT_SESSION"/>
+        <addPrimaryKey columnNames="USER_SESSION_ID,CLIENT_ID, CLIENT_STORAGE_PROVIDER, EXTERNAL_CLIENT_ID, OFFLINE_FLAG" constraintName="CONSTRAINT_OFFL_CL_SES_PK3" tableName="OFFLINE_CLIENT_SESSION"/>
 
      </changeSet>
 </databaseChangeLog>