From 69d36fcd65c81d28d35daef60c3f8ce692ca20b0 Mon Sep 17 00:00:00 2001
From: Alexander Schwartz <aschwart@redhat.com>
Date: Tue, 14 Jan 2025 10:24:52 +0100
Subject: [PATCH] Connect to primary PostgreSQL instance by default (#36330)

Closes #24493

Signed-off-by: Alexander Schwartz <aschwart@redhat.com>
---
 .../DefaultJpaConnectionProviderFactory.java  | 36 +++++++++++++++----
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/model/jpa/src/main/java/org/keycloak/connections/jpa/DefaultJpaConnectionProviderFactory.java b/model/jpa/src/main/java/org/keycloak/connections/jpa/DefaultJpaConnectionProviderFactory.java
index 3132631bff3..dc804286dde 100755
--- a/model/jpa/src/main/java/org/keycloak/connections/jpa/DefaultJpaConnectionProviderFactory.java
+++ b/model/jpa/src/main/java/org/keycloak/connections/jpa/DefaultJpaConnectionProviderFactory.java
@@ -178,9 +178,7 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
                         } else {
                             String url = config.get("url");
                             String driver = config.get("driver");
-                            if (driver.equals("org.h2.Driver")) {
-                                url = addH2NonKeywords(url);
-                            }
+                            url = augmentJdbcUrl(driver, url);
                             properties.put(AvailableSettings.JAKARTA_JDBC_URL, url);
                             properties.put(AvailableSettings.JAKARTA_JDBC_DRIVER, driver);
 
@@ -373,9 +371,7 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
             } else {
                 String url = config.get("url");
                 String driver = config.get("driver");
-                if (driver.equals("org.h2.Driver")) {
-                    url = addH2NonKeywords(url);
-                }
+                url = augmentJdbcUrl(driver, url);
                 Class.forName(driver);
                 return DriverManager.getConnection(StringPropertyReplacer.replaceProperties(url, System.getProperties()), config.get("user"), config.get("password"));
             }
@@ -384,6 +380,16 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
         }
     }
 
+    private String augmentJdbcUrl(String driver, String url) {
+        if (driver.equals("org.postgresql.xa.PGXADataSource") || driver.equals("org.postgresql.Driver")) {
+            url = addPostgreSQLKeywords(url);
+        }
+        if (driver.equals("org.h2.Driver")) {
+            url = addH2NonKeywords(url);
+        }
+        return url;
+    }
+
     @Override
     public String getSchema() {
         String schema = config.get("schema");
@@ -435,4 +441,22 @@ public class DefaultJpaConnectionProviderFactory implements JpaConnectionProvide
         return jdbcUrl;
     }
 
+    /**
+     * For a PostgreSQL cluster, Keycloak would need to connect to the primary node that is writable.
+     * The `targetServerType` should avoid connecting to a reader instance accidentally during node failover.
+
+     * @return JDBC URL with <code>targetServerType=primary</code> appended if the URL doesn't contain <code>targetServerType=</code> yet
+     */
+    private String addPostgreSQLKeywords(String jdbcUrl) {
+        if (!jdbcUrl.contains("targetServerType=")) {
+            if (jdbcUrl.contains("?")) {
+                jdbcUrl = jdbcUrl + "&";
+            } else {
+                jdbcUrl = jdbcUrl + "?";
+            }
+            jdbcUrl = jdbcUrl + "targetServerType=primary";
+        }
+        return jdbcUrl;
+    }
+
 }