multi tenancy aware sequences

ebean-core/src/main/java/io/ebeaninternal/server/core/InternalConfiguration.java

@@ -97,6 +97,7 @@ public final class InternalConfiguration {
   private final ExtraMetrics extraMetrics = new ExtraMetrics();
   private ServerCacheNotify cacheNotify;
   private boolean localL2Caching;
+  private final DataSourceSupplier dataSourceSupplier;
 
   InternalConfiguration(boolean online, ClusterManager clusterManager, SpiBackgroundExecutor backgroundExecutor,
                         DatabaseConfig config, BootupClasses bootupClasses) {
@@ -124,6 +125,7 @@ public final class InternalConfiguration {
 
     final InternalConfigXmlMap xmlMap = initExternalMapping();
     this.dtoBeanManager = new DtoBeanManager(typeManager, xmlMap.readDtoMapping());
+    this.dataSourceSupplier = createDataSourceSupplier();
     this.beanDescriptorManager = new BeanDescriptorManager(this);
     Map<String, String> asOfTableMapping = beanDescriptorManager.deploy(xmlMap.xmlDeployment());
     Map<String, String> draftTableMap = beanDescriptorManager.draftTableMap();
@@ -311,6 +313,10 @@ Persister createPersister(SpiEbeanServer server) {
     return new DefaultPersister(server, binder, beanDescriptorManager);
   }
 
+  public DataSourceSupplier getDataSourceSupplier() {
+    return dataSourceSupplier;
+  }
+
   public SpiCacheManager getCacheManager() {
     return cacheManager;
   }
@@ -385,7 +391,7 @@ TransactionManager createTransactionManager(SpiServer server, DocStoreUpdateProc
 
     TransactionManagerOptions options =
       new TransactionManagerOptions(server, notifyL2CacheInForeground, config, scopeManager, clusterManager, backgroundExecutor,
-        indexUpdateProcessor, beanDescriptorManager, dataSource(), profileHandler(), logManager,
+        indexUpdateProcessor, beanDescriptorManager, getDataSourceSupplier(), profileHandler(), logManager,
         tableModState, cacheNotify, clockService);
 
     if (config.isDocStoreOnly()) {
@@ -410,7 +416,7 @@ private SpiProfileHandler profileHandler() {
   /**
    * Return the DataSource supplier based on the tenancy mode.
    */
-  private DataSourceSupplier dataSource() {
+  private DataSourceSupplier createDataSourceSupplier() {
     switch (config.getTenantMode()) {
       case DB:
       case DB_WITH_MASTER:

ebean-core/src/main/java/io/ebeaninternal/server/deploy/BeanDescriptorManager.java

@@ -3,6 +3,7 @@
 import io.ebean.BackgroundExecutor;
 import io.ebean.Model;
 import io.ebean.RawSqlBuilder;
+import io.ebean.Transaction;
 import io.ebean.annotation.ConstraintMode;
 import io.ebean.bean.BeanCollection;
 import io.ebean.bean.EntityBean;
@@ -34,6 +35,7 @@
 import io.ebeaninternal.server.properties.BeanPropertiesReader;
 import io.ebeaninternal.server.properties.BeanPropertyAccess;
 import io.ebeaninternal.server.properties.EnhanceBeanPropertyAccess;
+import io.ebeaninternal.server.transaction.DataSourceSupplier;
 import io.ebeaninternal.server.type.TypeManager;
 import io.ebeaninternal.xmapping.api.XmapEbean;
 import io.ebeaninternal.xmapping.api.XmapEntity;
@@ -94,7 +96,7 @@ public final class BeanDescriptorManager implements BeanDescriptorMap, SpiBeanTy
   private final Map<String, List<BeanDescriptor<?>>> tableToDescMap = new HashMap<>();
   private final Map<String, List<BeanDescriptor<?>>> tableToViewDescMap = new HashMap<>();
   private final DbIdentity dbIdentity;
-  private final DataSource dataSource;
+  private final DataSourceSupplier dataSourceSupplier;
   private final DatabasePlatform databasePlatform;
   private final SpiCacheManager cacheManager;
   private final BackgroundExecutor backgroundExecutor;
@@ -129,7 +131,7 @@ public BeanDescriptorManager(InternalConfiguration config) {
     this.cacheManager = config.getCacheManager();
     this.docStoreFactory = config.getDocStoreFactory();
     this.backgroundExecutor = config.getBackgroundExecutor();
-    this.dataSource = this.config.getDataSource();
+    this.dataSourceSupplier = config.getDataSourceSupplier();
     this.encryptKeyManager = this.config.getEncryptKeyManager();
     this.databasePlatform = this.config.getDatabasePlatform();
     this.multiValueBind = config.getMultiValueBind();
@@ -1234,7 +1236,39 @@ private <T> void setIdGeneration(DeployBeanDescriptor<T> desc) {
   }
 
   private PlatformIdGenerator createSequenceIdGenerator(String seqName, int stepSize) {
-    return databasePlatform.createSequenceIdGenerator(backgroundExecutor, dataSource, stepSize, seqName);
+    return new PlatformIdGenerator() {
+
+      private Map<DataSource, PlatformIdGenerator> map = Collections.synchronizedMap(new WeakHashMap<>());
+
+      private PlatformIdGenerator create() {
+        return databasePlatform.createSequenceIdGenerator(backgroundExecutor, dataSourceSupplier.getDataSource(), stepSize, seqName);
+      }
+
+      private PlatformIdGenerator get() {
+        return map.computeIfAbsent(dataSourceSupplier.getDataSource(), k -> create());
+      }
+
+      @Override
+      public void preAllocateIds(int allocateSize) {
+        get().preAllocateIds(allocateSize);
+
+      }
+
+      @Override
+      public Object nextId(Transaction transaction) {
+        return get().nextId(transaction);
+      }
+
+      @Override
+      public boolean isDbSequence() {
+        return get().isDbSequence();
+      }
+
+      @Override
+      public String getName() {
+        return get().getName();
+      }
+    };
   }
 
   private void createByteCode(DeployBeanDescriptor<?> deploy) {

ebean-test/src/test/java/org/tests/idkeys/TestSequenceMultiTenant.java

@@ -0,0 +1,196 @@
+package org.tests.idkeys;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import javax.sql.DataSource;
+
+import org.junit.jupiter.api.Test;
+import org.multitenant.partition.UserContext;
+import org.tests.idkeys.db.GenKeySeqA;
+
+import io.ebean.Database;
+import io.ebean.DatabaseFactory;
+import io.ebean.config.DatabaseConfig;
+import io.ebean.config.TenantDataSourceProvider;
+import io.ebean.config.TenantMode;
+import io.ebean.config.TenantSchemaProvider;
+import io.ebean.config.dbplatform.h2.H2Platform;
+
+/**
+ * Tests sequences for multiple tenants.
+ */
+class TestSequenceMultiTenant {
+
+  /**
+   * Tests sequences using multi tenancy per database
+   */
+  @Test
+  void test_multi_tenant_db_sequences() {
+
+    Database db = setupDb();
+
+    UserContext.set("4711", "1");
+    assertEquals(1L, db.nextId(GenKeySeqA.class));
+    assertEquals(2L, db.nextId(GenKeySeqA.class));
+
+    UserContext.set("5711", "2");
+    assertEquals(1L, db.nextId(GenKeySeqA.class));
+
+    UserContext.set("4711", "1");
+    assertEquals(3L, db.nextId(GenKeySeqA.class));
+
+  }
+
+  private static Database setupDb() {
+
+    DatabaseConfig config = new DatabaseConfig();
+
+    config.setName("h2multitenantseq");
+    config.loadFromProperties();
+    config.setDdlGenerate(true);
+    config.setDdlRun(true);
+    config.setDdlExtra(false);
+    config.setRegister(false);
+    config.setDefaultServer(false);
+    config.setCurrentTenantProvider(() -> UserContext.get().getTenantId());
+    config.setTenantMode(TenantMode.DB);
+    config.setDatabasePlatform(new H2Platform());
+    config.setTenantDataSourceProvider(new TenantDataSourceProvider() {
+
+      Map<Object, DataSource> map = new ConcurrentHashMap<>();
+
+      @Override
+      public DataSource dataSource(Object tenantId) {
+        if (tenantId == null) {
+          tenantId = "1";
+        }
+        return map.computeIfAbsent(tenantId, this::createDataSource);
+      }
+
+      private DataSource createDataSource(Object tenantId) {
+
+        DatabaseConfig config = new DatabaseConfig();
+
+        config.setName("h2multitenantseq");
+        config.loadFromProperties();
+        config.setDdlRun(true);
+        config.setDdlExtra(false);
+        config.setRegister(false);
+        config.setDefaultServer(false);
+        config.getDataSourceConfig().setUrl("jdbc:h2:mem:h2multitenantseq-" + tenantId);
+
+        return DatabaseFactory.create(config).pluginApi().dataSource();
+      }
+    });
+
+    config.getClasses().add(GenKeySeqA.class);
+
+    return DatabaseFactory.create(config);
+  }
+
+  /**
+   * Tests sequences using multi tenancy per schema
+   */
+  @Test
+  void test_multi_tenant_schema_sequences() throws SQLException {
+    createDDl("PUBLIC");
+    createDDl("TENANT_SCHEMA_1");
+    createDDl("TENANT_SCHEMA_2");
+
+    Database db = setupSchema();
+
+    Connection connection = db.dataSource().getConnection();
+
+    // debugging schemas
+    // see org.h2.jdbc.JdbcDatabaseMetaData.getSchemas()
+    ResultSet rs = connection.getMetaData().getSchemas();
+    String[] schemaData = new String[4];
+    int schemaCnt = 0;
+    while(rs.next()) {
+      schemaData[schemaCnt++] = String.format("%s, %s, %s", rs.getString(1), rs.getString(2), rs.getBoolean(3));
+    }
+    assertArrayEquals(schemaData, new String[]{"INFORMATION_SCHEMA, H2MULTITENANTSEQ, false", 
+      "PUBLIC, H2MULTITENANTSEQ, true", 
+      "TENANT_SCHEMA_1, H2MULTITENANTSEQ, false", 
+      "TENANT_SCHEMA_2, H2MULTITENANTSEQ, false"});
+
+    // debugging sequences
+    rs = connection.prepareStatement("SELECT * FROM INFORMATION_SCHEMA.SEQUENCES;").executeQuery();
+    String[] sequenceData = new String[4];
+    int sequenceCnt = 0;
+    while (rs.next()) {
+      // SEQUENCE_CATALOG,SEQUENCE_SCHEMA,SEQUENCE_NAME,CURRENT_VALUE,INCREMENT,IS_GENERATED,REMARKS,CACHE,MIN_VALUE,MAX_VALUE,IS_CYCLE,ID
+      String format = String.format("%s, %s, %s", rs.getString(1), rs.getString(2), rs.getString(3));
+      System.out.println(format);
+      sequenceData[sequenceCnt++] = format;
+    }
+
+    UserContext.set("4711", "1");
+    assertEquals(1L, db.nextId(GenKeySeqA.class));
+    assertEquals(2L, db.nextId(GenKeySeqA.class));
+
+    UserContext.set("5711", "2");
+    assertEquals(1L, db.nextId(GenKeySeqA.class));
+
+    UserContext.set("4711", "1");
+    assertEquals(3L, db.nextId(GenKeySeqA.class));
+  }
+
+  private Database setupSchema() {
+
+    DatabaseConfig config = new DatabaseConfig();
+
+    config.setName("h2multitenantseq");
+    config.loadFromProperties();
+    config.setDdlGenerate(true);
+    config.setDdlRun(false);
+    config.setDdlExtra(false);
+    config.setRegister(false);
+    config.setDefaultServer(false);
+    config.setCurrentTenantProvider(() -> UserContext.get().getTenantId());
+    config.setTenantMode(TenantMode.SCHEMA);
+    config.setDatabasePlatform(new H2Platform());
+    config.getDataSourceConfig().setUrl("jdbc:h2:mem:h2multitenantseq;DB_CLOSE_ON_EXIT=FALSE;");
+    config.setTenantSchemaProvider(new TenantSchemaProvider() {
+
+      @Override
+      public String schema(Object tenantId) {
+        return tenantId == null
+            ? "PUBLIC"
+            : "TENANT_SCHEMA_" + tenantId;
+      }
+    });
+
+    config.getClasses().add(GenKeySeqA.class);
+
+    return DatabaseFactory.create(config);
+  }
+
+  private void createDDl(String schema) {
+    DatabaseConfig config = new DatabaseConfig();
+
+    config.setName("h2multitenantseq");
+    config.loadFromProperties();
+    config.setDdlGenerate(true);
+    config.setDdlRun(true);
+    config.setDdlExtra(false);
+    config.setRegister(false);
+    config.setDefaultServer(false);
+    config.getDataSourceConfig().setUrl(
+        "jdbc:h2:mem:h2multitenantseq;DB_CLOSE_ON_EXIT=FALSE;INIT=CREATE SCHEMA IF NOT EXISTS "
+            + schema + "\\;SET SCHEMA " + schema);
+    config.setDbSchema(schema); // see io.ebeaninternal.dbmigration.DdlGenerator.createSchemaIfRequired(Connection)
+
+    config.getClasses().add(GenKeySeqA.class);
+
+    DatabaseFactory.create(config);
+  }
+
+}

ebean-test/src/test/resources/ebean.properties

@@ -88,6 +88,12 @@ datasource.h2multitenant.username=sa
 datasource.h2multitenant.password=
 datasource.h2multitenant.url=jdbc:h2:mem:h2multitenant
 
+ebean.h2multitenantseq.idType=SEQUENCE
+datasource.h2multitenantseq.username=sa
+datasource.h2multitenantseq.password=
+datasource.h2multitenantseq.url=jdbc:h2:mem:h2multitenantseq
+datasource.h2multitenantseq.driver=org.h2.Driver
+
 datasource.h2autocommit.autoCommit=true
 datasource.h2autocommit.username=sa
 datasource.h2autocommit.password=