From 6fad14ead731bb9eab09a15f5ada66d51e68e137 Mon Sep 17 00:00:00 2001
From: Glenn Renfro <grenfro@pivotal.io>
Date: Wed, 26 Apr 2017 16:24:03 -0400
Subject: [PATCH] Support DB2 on Mainframe as DB2ZOS

resolves #303

Added support for DB2ZOS DatabaseType.
Added DB2AS400, DB2VSE, based on how Spring Batch implemeted DB2 identification via metadata.
All DB2 database types will use the DB2 repository initialize sql script

Tested this against a local DB2.
---
 .../support/SqlPagingQueryProviderFactoryBean.java |  6 ++++++
 .../task/repository/support/DatabaseType.java      | 14 ++++++++++++--
 .../support/TaskRepositoryInitializer.java         |  3 ++-
 3 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/database/support/SqlPagingQueryProviderFactoryBean.java b/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/database/support/SqlPagingQueryProviderFactoryBean.java
index 6ec2d206..b9282dc9 100755
--- a/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/database/support/SqlPagingQueryProviderFactoryBean.java
+++ b/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/database/support/SqlPagingQueryProviderFactoryBean.java
@@ -17,6 +17,9 @@
 package org.springframework.cloud.task.repository.database.support;
 
 import static org.springframework.cloud.task.repository.support.DatabaseType.DB2;
+import static org.springframework.cloud.task.repository.support.DatabaseType.DB2AS400;
+import static org.springframework.cloud.task.repository.support.DatabaseType.DB2VSE;
+import static org.springframework.cloud.task.repository.support.DatabaseType.DB2ZOS;
 import static org.springframework.cloud.task.repository.support.DatabaseType.HSQL;
 import static org.springframework.cloud.task.repository.support.DatabaseType.H2;
 import static org.springframework.cloud.task.repository.support.DatabaseType.MYSQL;
@@ -69,6 +72,9 @@ public class SqlPagingQueryProviderFactoryBean implements FactoryBean<PagingQuer
 		providers.put(ORACLE, new OraclePagingQueryProvider());
 		providers.put(SQLSERVER, new SqlServerPagingQueryProvider());
 		providers.put(DB2, new Db2PagingQueryProvider());
+		providers.put(DB2VSE, new Db2PagingQueryProvider());
+		providers.put(DB2ZOS, new Db2PagingQueryProvider());
+		providers.put(DB2AS400, new Db2PagingQueryProvider());
 	}
 
 	/**
diff --git a/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/support/DatabaseType.java b/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/support/DatabaseType.java
index 72dff1bb..f8721a3a 100644
--- a/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/support/DatabaseType.java
+++ b/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/support/DatabaseType.java
@@ -41,7 +41,10 @@ public enum DatabaseType {
 	MYSQL("MySQL"),
 	POSTGRES("PostgreSQL"),
 	SQLSERVER("Microsoft SQL Server"),
-	DB2("DB2");
+	DB2("DB2"),
+	DB2VSE("DB2VSE"),
+	DB2ZOS("DB2ZOS"),
+	DB2AS400("DB2AS400");
 
 	private static final Map<String, DatabaseType> dbNameMap;
 
@@ -71,9 +74,16 @@ public enum DatabaseType {
 		if (StringUtils.hasText(databaseProductName) && !databaseProductName.equals("DB2/Linux") && databaseProductName.startsWith("DB2")) {
 			String databaseProductVersion =
 					JdbcUtils.extractDatabaseMetaData(dataSource, "getDatabaseProductVersion").toString();
-			if (!databaseProductVersion.startsWith("SQL")) {
+			if (databaseProductVersion.startsWith("ARI")) {
+				databaseProductName = "DB2VSE";
+			}
+			else if (databaseProductVersion.startsWith("DSN")) {
 				databaseProductName = "DB2ZOS";
 			}
+			else if (databaseProductName.indexOf("AS") != -1 && (databaseProductVersion.startsWith("QSQ") ||
+					databaseProductVersion.substring(databaseProductVersion.indexOf('V')).matches("V\\dR\\d[mM]\\d"))) {
+				databaseProductName = "DB2AS400";
+			}
 			else {
 				databaseProductName = JdbcUtils.commonDatabaseName(databaseProductName);
 			}
diff --git a/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/support/TaskRepositoryInitializer.java b/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/support/TaskRepositoryInitializer.java
index d4ff9ad3..7ba21bb4 100644
--- a/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/support/TaskRepositoryInitializer.java
+++ b/spring-cloud-task-core/src/main/java/org/springframework/cloud/task/repository/support/TaskRepositoryInitializer.java
@@ -27,6 +27,7 @@ import org.springframework.beans.factory.annotation.Value;
 import org.springframework.core.io.ResourceLoader;
 import org.springframework.jdbc.datasource.init.DatabasePopulatorUtils;
 import org.springframework.jdbc.datasource.init.ResourceDatabasePopulator;
+import org.springframework.jdbc.support.JdbcUtils;
 import org.springframework.jdbc.support.MetaDataAccessException;
 
 /**
@@ -76,7 +77,7 @@ public final class TaskRepositoryInitializer implements InitializingBean {
 
 	private String getDatabaseType(DataSource dataSource) {
 		try {
-			return DatabaseType.fromMetaData(dataSource).toString().toLowerCase();
+			return JdbcUtils.commonDatabaseName(DatabaseType.fromMetaData(dataSource).toString()).toLowerCase();
 		}
 		catch (MetaDataAccessException ex) {
 			throw new IllegalStateException("Unable to detect database type", ex);