From a2b418df8b41b97159d49023a87270b5567b2f98 Mon Sep 17 00:00:00 2001
From: rob bygrave <robin.bygrave@gmail.com>
Date: Thu, 12 Sep 2019 20:59:52 +1200
Subject: [PATCH] #1820 - ENH: For DB Migration add setLockTimeout(10); ... to
 add Postgres `set lock_timeout = ` to generated migrations

---
 .../java/io/ebean/dbmigration/DbMigration.java |  9 +++++++++
 .../dbmigration/DefaultDbMigration.java        |  8 +++++++-
 .../ddlgeneration/platform/PlatformDdl.java    |  7 +++++++
 .../ddlgeneration/platform/PostgresDdl.java    |  4 ++++
 .../dbmigration/model/PlatformDdlWriter.java   | 12 +++++++++++-
 .../platform/PostgresDdlTest.java              | 18 ++++++++++++++++++
 6 files changed, 56 insertions(+), 2 deletions(-)
 create mode 100644 src/test/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PostgresDdlTest.java

diff --git a/src/main/java/io/ebean/dbmigration/DbMigration.java b/src/main/java/io/ebean/dbmigration/DbMigration.java
index 5319719ba..fd0fe7e4e 100644
--- a/src/main/java/io/ebean/dbmigration/DbMigration.java
+++ b/src/main/java/io/ebean/dbmigration/DbMigration.java
@@ -144,6 +144,15 @@ public interface DbMigration {
    */
   void setAddForeignKeySkipCheck(boolean addForeignKeySkipCheck);
 
+  /**
+   * Set the lock timeout to be included with the DDL generation.
+   * <p>
+   * Currently this is only useful for Postgres migrations adding a <code>set lock_timeout</code>
+   * statement to the generated database migration.
+   * </p>
+   */
+  void setLockTimeout(int seconds);
+
   /**
    * Add an additional platform to write the migration DDL.
    * <p>
diff --git a/src/main/java/io/ebeaninternal/dbmigration/DefaultDbMigration.java b/src/main/java/io/ebeaninternal/dbmigration/DefaultDbMigration.java
index 2c2016e8c..48e705cc6 100644
--- a/src/main/java/io/ebeaninternal/dbmigration/DefaultDbMigration.java
+++ b/src/main/java/io/ebeaninternal/dbmigration/DefaultDbMigration.java
@@ -108,6 +108,7 @@ public class DefaultDbMigration implements DbMigration {
   protected String name;
   protected String generatePendingDrop;
   private boolean addForeignKeySkipCheck;
+  private int lockTimeoutSeconds;
 
   protected boolean includeBuiltInPartitioning = true;
 
@@ -187,6 +188,11 @@ public class DefaultDbMigration implements DbMigration {
     this.addForeignKeySkipCheck = addForeignKeySkipCheck;
   }
 
+  @Override
+  public void setLockTimeout(int seconds) {
+    this.lockTimeoutSeconds = seconds;
+  }
+
   @Override
   public void setGeneratePendingDrop(String generatePendingDrop) {
     this.generatePendingDrop = generatePendingDrop;
@@ -641,7 +647,7 @@ public class DefaultDbMigration implements DbMigration {
   }
 
   private PlatformDdlWriter createDdlWriter(DatabasePlatform platform) {
-    return new PlatformDdlWriter(platform, serverConfig, migrationConfig);
+    return new PlatformDdlWriter(platform, serverConfig, migrationConfig, lockTimeoutSeconds);
   }
 
   /**
diff --git a/src/main/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PlatformDdl.java b/src/main/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PlatformDdl.java
index 209df5276..a78a80b58 100644
--- a/src/main/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PlatformDdl.java
+++ b/src/main/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PlatformDdl.java
@@ -211,6 +211,13 @@ public class PlatformDdl {
     return inlineForeignKeys;
   }
 
+  /**
+   * By default this does nothing returning null / no lock timeout.
+   */
+  public String setLockTimeout(int lockTimeoutSeconds) {
+    return null;
+  }
+
   /**
    * Write all the table columns converting to platform types as necessary.
    */
diff --git a/src/main/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PostgresDdl.java b/src/main/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PostgresDdl.java
index 56d307892..017f8f4fb 100644
--- a/src/main/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PostgresDdl.java
+++ b/src/main/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PostgresDdl.java
@@ -20,6 +20,10 @@ public class PostgresDdl extends PlatformDdl {
     this.addForeignKeySkipCheck = " not valid";
   }
 
+  public String setLockTimeout(int lockTimeoutSeconds) {
+    return "set lock_timeout = " + (lockTimeoutSeconds * 1000);
+  }
+
   @Override
   public boolean suppressPrimaryKeyOnPartition() {
     return true;
diff --git a/src/main/java/io/ebeaninternal/dbmigration/model/PlatformDdlWriter.java b/src/main/java/io/ebeaninternal/dbmigration/model/PlatformDdlWriter.java
index 8b2bd71ab..43ce498b9 100644
--- a/src/main/java/io/ebeaninternal/dbmigration/model/PlatformDdlWriter.java
+++ b/src/main/java/io/ebeaninternal/dbmigration/model/PlatformDdlWriter.java
@@ -33,10 +33,13 @@ public class PlatformDdlWriter {
 
   private final PlatformDdl platformDdl;
 
-  public PlatformDdlWriter(DatabasePlatform platform, ServerConfig serverConfig, DbMigrationConfig config) {
+  private final int lockTimeoutSeconds;
+
+  public PlatformDdlWriter(DatabasePlatform platform, ServerConfig serverConfig, DbMigrationConfig config, int lockTimeoutSeconds) {
     this.platformDdl = PlatformDdlBuilder.create(platform);
     this.serverConfig = serverConfig;
     this.config = config;
+    this.lockTimeoutSeconds = lockTimeoutSeconds;
   }
 
   /**
@@ -46,6 +49,13 @@ public class PlatformDdlWriter {
 
     DdlHandler handler = handler();
     handler.generateProlog(write);
+    if (lockTimeoutSeconds > 0) {
+      String lockSql = platformDdl.setLockTimeout(lockTimeoutSeconds);
+      if (lockSql != null) {
+        write.apply().append(lockSql).endOfStatement().newLine();
+      }
+    }
+
     List<ChangeSet> changeSets = dbMigration.getChangeSet();
     for (ChangeSet changeSet : changeSets) {
       if (isApply(changeSet)) {
diff --git a/src/test/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PostgresDdlTest.java b/src/test/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PostgresDdlTest.java
new file mode 100644
index 000000000..611ec94dd
--- /dev/null
+++ b/src/test/java/io/ebeaninternal/dbmigration/ddlgeneration/platform/PostgresDdlTest.java
@@ -0,0 +1,18 @@
+package io.ebeaninternal.dbmigration.ddlgeneration.platform;
+
+import io.ebean.config.dbplatform.postgres.PostgresPlatform;
+import org.junit.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class PostgresDdlTest {
+
+  private PostgresDdl postgresDdl = new PostgresDdl(new PostgresPlatform());
+
+  @Test
+  public void setLockTimeout() {
+
+    final String sql = postgresDdl.setLockTimeout(5);
+    assertThat(sql).isEqualTo("set lock_timeout = 5000");
+  }
+}