From 658988ccd008ac303e99dd8e88a02370213ab110 Mon Sep 17 00:00:00 2001
From: Hynek Mlnarik <hmlnarik@redhat.com>
Date: Tue, 18 Oct 2016 13:54:31 +0200
Subject: [PATCH] KEYCLOAK-3588 Make manual update scripts play nicely with
 custom JpaEntityProviders

---
 .../LiquibaseJpaUpdaterProvider.java          | 31 +++++++++++++------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/LiquibaseJpaUpdaterProvider.java b/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/LiquibaseJpaUpdaterProvider.java
index f3b304d7c9..cd2d2537c6 100755
--- a/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/LiquibaseJpaUpdaterProvider.java
+++ b/model/jpa/src/main/java/org/keycloak/connections/jpa/updater/liquibase/LiquibaseJpaUpdaterProvider.java
@@ -79,10 +79,19 @@ public class LiquibaseJpaUpdaterProvider implements JpaUpdaterProvider {
         // Need ThreadLocal as liquibase doesn't seem to have API to inject custom objects into tasks
         ThreadLocalSessionContext.setCurrentSession(session);
 
+        Writer exportWriter = null;
         try {
             // Run update with keycloak master changelog first
             Liquibase liquibase = getLiquibaseForKeycloakUpdate(connection, defaultSchema);
-            updateChangeSet(liquibase, liquibase.getChangeLogFile(), file);
+            if (file != null) {
+                exportWriter = new FileWriter(file);
+                List<ChangeSet> changeSets = getChangeSets(liquibase);
+                if (! changeSets.isEmpty() && liquibase.getDatabase().getRanChangeSetList().isEmpty()) {
+                    outputChangeLogTableCreationScript(liquibase, exportWriter);
+                }
+            }
+
+            updateChangeSet(liquibase, liquibase.getChangeLogFile(), exportWriter);
 
             // Run update for each custom JpaEntityProvider
             Set<JpaEntityProvider> jpaProviders = session.getAllProviders(JpaEntityProvider.class);
@@ -92,18 +101,25 @@ public class LiquibaseJpaUpdaterProvider implements JpaUpdaterProvider {
                     String factoryId = jpaProvider.getFactoryId();
                     String changelogTableName = JpaUtils.getCustomChangelogTableName(factoryId);
                     liquibase = getLiquibaseForCustomProviderUpdate(connection, defaultSchema, customChangelog, jpaProvider.getClass().getClassLoader(), changelogTableName);
-                    updateChangeSet(liquibase, liquibase.getChangeLogFile(), file);
+                    updateChangeSet(liquibase, liquibase.getChangeLogFile(), exportWriter);
                 }
             }
         } catch (Exception e) {
             throw new RuntimeException("Failed to update database", e);
         } finally {
             ThreadLocalSessionContext.removeCurrentSession();
+            if (exportWriter != null) {
+                try {
+                    exportWriter.close();
+                } catch (IOException ioe) {
+                    // ignore
+                }
+            }
         }
     }
 
 
-    protected void updateChangeSet(Liquibase liquibase, String changelog, File exportFile) throws LiquibaseException, IOException {
+    protected void updateChangeSet(Liquibase liquibase, String changelog, Writer exportWriter) throws LiquibaseException, IOException {
         List<ChangeSet> changeSets = getChangeSets(liquibase);
         if (!changeSets.isEmpty()) {
             List<RanChangeSet> ranChangeSets = liquibase.getDatabase().getRanChangeSetList();
@@ -117,13 +133,8 @@ public class LiquibaseJpaUpdaterProvider implements JpaUpdaterProvider {
                 }
             }
 
-            if (exportFile != null) {
-                try (Writer exportWriter = new FileWriter(exportFile)) {
-                    if (ranChangeSets.isEmpty()) {
-                        outputChangeLogTableCreationScript(liquibase, exportWriter);
-                    }
-                    liquibase.update((Contexts) null, new LabelExpression(), exportWriter, false);
-                }
+            if (exportWriter != null) {
+                liquibase.update((Contexts) null, new LabelExpression(), exportWriter, false);
             } else {
                 liquibase.update((Contexts) null);
             }