From a68486670808d60c681144f1f5dae3755f5b5498 Mon Sep 17 00:00:00 2001
From: Kellen Dye <kellend@spotify.com>
Date: Fri, 3 Jan 2025 12:43:11 -0500
Subject: [PATCH] Add mariadb support

---
 dbeam-core/pom.xml                             |  8 ++++++++
 .../spotify/dbeam/args/JdbcConnectionUtil.java | 18 ++++++++++++++----
 e2e/ddl.sql                                    |  2 +-
 pom.xml                                        | 11 +++++++++++
 4 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/dbeam-core/pom.xml b/dbeam-core/pom.xml
index ea4980a4..b58e9371 100644
--- a/dbeam-core/pom.xml
+++ b/dbeam-core/pom.xml
@@ -56,6 +56,10 @@
       <groupId>com.mysql</groupId>
       <artifactId>mysql-connector-j</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.mariadb.jdbc</groupId>
+      <artifactId>mariadb-java-client</artifactId>
+    </dependency>
     <dependency>
       <groupId>com.google.cloud.sql</groupId>
       <artifactId>postgres-socket-factory</artifactId>
@@ -64,6 +68,10 @@
       <groupId>com.google.cloud.sql</groupId>
       <artifactId>mysql-socket-factory-connector-j-8</artifactId>
     </dependency>
+    <dependency>
+      <groupId>com.google.cloud.sql</groupId>
+      <artifactId>mariadb-socket-factory</artifactId>
+    </dependency>
 
     <dependency>
       <groupId>com.google.guava</groupId>
diff --git a/dbeam-core/src/main/java/com/spotify/dbeam/args/JdbcConnectionUtil.java b/dbeam-core/src/main/java/com/spotify/dbeam/args/JdbcConnectionUtil.java
index 6f70b7bd..1c5b3b88 100644
--- a/dbeam-core/src/main/java/com/spotify/dbeam/args/JdbcConnectionUtil.java
+++ b/dbeam-core/src/main/java/com/spotify/dbeam/args/JdbcConnectionUtil.java
@@ -23,6 +23,7 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 public class JdbcConnectionUtil {
 
@@ -30,14 +31,23 @@ public class JdbcConnectionUtil {
       ImmutableMap.of(
           "postgresql", "org.postgresql.Driver",
           "mysql", "com.mysql.cj.jdbc.Driver",
-          "h2", "org.h2.Driver");
+          "mariadb", "org.mariadb.jdbc.Driver",
+          "h2", "org.h2.Driver"
+      );
+
+  private static String driverPrefixes =
+      driverMapping.keySet().stream()
+          .map(k -> "jdbc:" + k)
+          .collect(Collectors.joining(", "));
 
   public static String getDriverClass(final String url) throws ClassNotFoundException {
     final String[] parts = url.split(":", 3);
     Preconditions.checkArgument(
-        parts.length > 1 && "jdbc".equals(parts[0]) && driverMapping.get(parts[1]) != null,
-        "Invalid jdbc connection URL: %s. Expect jdbc:postgresql or jdbc:mysql as prefix.",
-        url);
+        parts.length > 1 && "jdbc".equals(parts[0]) && driverMapping.containsKey(parts[1]),
+        "Invalid jdbc connection URL: %s. Expected one of %s as prefix.",
+        url,
+        driverPrefixes
+    );
     return Class.forName(driverMapping.get(parts[1])).getCanonicalName();
   }
 }
diff --git a/e2e/ddl.sql b/e2e/ddl.sql
index 29d2f866..760eff4f 100644
--- a/e2e/ddl.sql
+++ b/e2e/ddl.sql
@@ -1,5 +1,5 @@
 
--- This file contatins psql views with complex types to validate and troubleshoot dbeam
+-- This file contains psql views with complex types to validate and troubleshoot dbeam
 -- import with:
 -- psql -f ddl.sql postgres
 
diff --git a/pom.xml b/pom.xml
index 5fec16c1..20976458 100644
--- a/pom.xml
+++ b/pom.xml
@@ -135,6 +135,7 @@
     <bouncycastle.version>1.78.1</bouncycastle.version>
     <junit.version>4.13.2</junit.version>
     <mysql.version>8.4.0</mysql.version>
+    <mariadb.version>3.5.1</mariadb.version>
     <postgresql.version>42.7.4</postgresql.version>
     <socket-factory.version>1.18.0</socket-factory.version>
   </properties>
@@ -365,6 +366,11 @@
         <artifactId>mysql-connector-j</artifactId>
         <version>${mysql.version}</version>
       </dependency>
+      <dependency>
+        <groupId>org.mariadb.jdbc</groupId>
+        <artifactId>mariadb-java-client</artifactId>
+        <version>${mariadb.version}</version>
+      </dependency>
       <dependency>
         <groupId>com.google.cloud.sql</groupId>
         <artifactId>postgres-socket-factory</artifactId>
@@ -376,6 +382,11 @@
         <artifactId>mysql-socket-factory-connector-j-8</artifactId>
         <version>${socket-factory.version}</version>
       </dependency>
+      <dependency>
+        <groupId>com.google.cloud.sql</groupId>
+        <artifactId>mariadb-socket-factory</artifactId>
+        <version>${socket-factory.version}</version>
+      </dependency>
       <dependency>
         <groupId>org.hamcrest</groupId>
         <artifactId>hamcrest</artifactId>