From d90d653d1449affd8ef272141cf0af59e635b115 Mon Sep 17 00:00:00 2001
From: Anthony Sena <anthonysena@users.noreply.github.com>
Date: Mon, 26 Feb 2024 15:18:36 -0500
Subject: [PATCH] Add Databricks profile (#2339)

* Create new webapi-databricks profile
* Check for spark and databricks connection string info
---
 pom.xml                                               | 11 +++++++++++
 src/main/java/org/ohdsi/webapi/DataAccessConfig.java  |  2 +-
 .../org/ohdsi/webapi/util/CancelableJdbcTemplate.java |  3 ++-
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/pom.xml b/pom.xml
index bf5b8deeeb..f799a72a71 100644
--- a/pom.xml
+++ b/pom.xml
@@ -1482,6 +1482,17 @@
         </plugins>
       </build>
     </profile>
+    <profile>
+      <id>webapi-databricks</id>
+      <dependencies>
+        <dependency>
+          <groupId>com.databricks</groupId>
+          <artifactId>databricks-jdbc</artifactId>
+          <version>2.6.34</version>
+          <scope>runtime</scope>
+        </dependency>
+      </dependencies>
+    </profile>
     <profile>
       <id>webapi-bigquery</id>
       <properties>
diff --git a/src/main/java/org/ohdsi/webapi/DataAccessConfig.java b/src/main/java/org/ohdsi/webapi/DataAccessConfig.java
index 4391198906..8bc4e438dc 100644
--- a/src/main/java/org/ohdsi/webapi/DataAccessConfig.java
+++ b/src/main/java/org/ohdsi/webapi/DataAccessConfig.java
@@ -82,7 +82,7 @@ public DataSource primaryDataSource() {
         //note autocommit defaults vary across vendors. use provided @Autowired TransactionTemplate
 
         String[] supportedDrivers;
-        supportedDrivers = new String[]{"org.postgresql.Driver", "com.microsoft.sqlserver.jdbc.SQLServerDriver", "oracle.jdbc.driver.OracleDriver", "com.amazon.redshift.jdbc.Driver", "com.cloudera.impala.jdbc.Driver", "net.starschema.clouddb.jdbc.BQDriver", "org.netezza.Driver", "com.simba.googlebigquery.jdbc42.Driver", "org.apache.hive.jdbc.HiveDriver", "com.simba.spark.jdbc.Driver", "net.snowflake.client.jdbc.SnowflakeDriver"};
+        supportedDrivers = new String[]{"org.postgresql.Driver", "com.microsoft.sqlserver.jdbc.SQLServerDriver", "oracle.jdbc.driver.OracleDriver", "com.amazon.redshift.jdbc.Driver", "com.cloudera.impala.jdbc.Driver", "net.starschema.clouddb.jdbc.BQDriver", "org.netezza.Driver", "com.simba.googlebigquery.jdbc42.Driver", "org.apache.hive.jdbc.HiveDriver", "com.simba.spark.jdbc.Driver", "net.snowflake.client.jdbc.SnowflakeDriver", "com.databricks.client.jdbc.Driver"};
         for (String driverName : supportedDrivers) {
             try {
                 Class.forName(driverName);
diff --git a/src/main/java/org/ohdsi/webapi/util/CancelableJdbcTemplate.java b/src/main/java/org/ohdsi/webapi/util/CancelableJdbcTemplate.java
index 00a29b973d..a7ebb3d39e 100644
--- a/src/main/java/org/ohdsi/webapi/util/CancelableJdbcTemplate.java
+++ b/src/main/java/org/ohdsi/webapi/util/CancelableJdbcTemplate.java
@@ -78,7 +78,8 @@ public int[] doInStatement(Statement stmt) throws SQLException, DataAccessExcept
         }
         else {
           for (int i = 0; i < sql.length; i++) {
-            if (stmt.getConnection().getMetaData().getURL().startsWith("jdbc:spark")) {
+            String connectionString = stmt.getConnection().getMetaData().getURL();
+            if (connectionString.startsWith("jdbc:spark") || connectionString.startsWith("jdbc:databricks")) {
               this.currSql = BigQuerySparkTranslate.sparkHandleInsert(sql[i], stmt.getConnection());
               if (this.currSql == "" || this.currSql.isEmpty() || this.currSql == null) {
                 rowsAffected[i] = -1;