From 58072a7c03d849a1145ed409e8a7aebfb5728753 Mon Sep 17 00:00:00 2001
From: Thijs Lemmens <thijs.lemmens@xenit.eu>
Date: Thu, 16 May 2024 15:22:09 +0200
Subject: [PATCH] Make pgvector available with JDBC syntax (#8633)

---
 .../jdbc/AbstractJDBCDriverTest.java          |  3 +-
 .../containers/PgVectorContainerProvider.java | 44 +++++++++++++++++++
 ...s.containers.JdbcDatabaseContainerProvider |  1 +
 .../jdbc/pgvector/PgVectorJDBCDriverTest.java | 28 ++++++++++++
 4 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 modules/postgresql/src/main/java/org/testcontainers/containers/PgVectorContainerProvider.java
 create mode 100644 modules/postgresql/src/test/java/org/testcontainers/jdbc/pgvector/PgVectorJDBCDriverTest.java

diff --git a/modules/jdbc-test/src/main/java/org/testcontainers/jdbc/AbstractJDBCDriverTest.java b/modules/jdbc-test/src/main/java/org/testcontainers/jdbc/AbstractJDBCDriverTest.java
index a22b4e24ed0..5a22e37ecfc 100644
--- a/modules/jdbc-test/src/main/java/org/testcontainers/jdbc/AbstractJDBCDriverTest.java
+++ b/modules/jdbc-test/src/main/java/org/testcontainers/jdbc/AbstractJDBCDriverTest.java
@@ -130,7 +130,8 @@ private void performTestForJDBCParamUsage(HikariDataSource dataSource) throws SQ
         if (
             databaseType.equalsIgnoreCase("postgresql") ||
             databaseType.equalsIgnoreCase("postgis") ||
-            databaseType.equalsIgnoreCase("timescaledb")
+            databaseType.equalsIgnoreCase("timescaledb") ||
+            databaseType.equalsIgnoreCase("pgvector")
         ) {
             databaseQuery = "SELECT CURRENT_DATABASE()";
         }
diff --git a/modules/postgresql/src/main/java/org/testcontainers/containers/PgVectorContainerProvider.java b/modules/postgresql/src/main/java/org/testcontainers/containers/PgVectorContainerProvider.java
new file mode 100644
index 00000000000..5aebbbefbc3
--- /dev/null
+++ b/modules/postgresql/src/main/java/org/testcontainers/containers/PgVectorContainerProvider.java
@@ -0,0 +1,44 @@
+package org.testcontainers.containers;
+
+import org.testcontainers.jdbc.ConnectionUrl;
+import org.testcontainers.utility.DockerImageName;
+
+/**
+ * Factory for PgVector containers.
+ *
+ * @see <a href="https://github.com/pgvector/pgvector">https://github.com/pgvector/pgvector</a>
+ */
+public class PgVectorContainerProvider extends JdbcDatabaseContainerProvider {
+
+    private static final String NAME = "pgvector";
+
+    private static final String DEFAULT_TAG = "pg16";
+
+    private static final DockerImageName DEFAULT_IMAGE = DockerImageName
+        .parse("pgvector/pgvector")
+        .asCompatibleSubstituteFor("postgres");
+
+    public static final String USER_PARAM = "user";
+
+    public static final String PASSWORD_PARAM = "password";
+
+    @Override
+    public boolean supports(String databaseType) {
+        return databaseType.equals(NAME);
+    }
+
+    @Override
+    public JdbcDatabaseContainer newInstance() {
+        return newInstance(DEFAULT_TAG);
+    }
+
+    @Override
+    public JdbcDatabaseContainer newInstance(String tag) {
+        return new PostgreSQLContainer(DEFAULT_IMAGE.withTag(tag));
+    }
+
+    @Override
+    public JdbcDatabaseContainer newInstance(ConnectionUrl connectionUrl) {
+        return newInstanceFromConnectionUrl(connectionUrl, USER_PARAM, PASSWORD_PARAM);
+    }
+}
diff --git a/modules/postgresql/src/main/resources/META-INF/services/org.testcontainers.containers.JdbcDatabaseContainerProvider b/modules/postgresql/src/main/resources/META-INF/services/org.testcontainers.containers.JdbcDatabaseContainerProvider
index 05df9054522..33429fa8b1e 100644
--- a/modules/postgresql/src/main/resources/META-INF/services/org.testcontainers.containers.JdbcDatabaseContainerProvider
+++ b/modules/postgresql/src/main/resources/META-INF/services/org.testcontainers.containers.JdbcDatabaseContainerProvider
@@ -1,3 +1,4 @@
 org.testcontainers.containers.PostgreSQLContainerProvider
 org.testcontainers.containers.PostgisContainerProvider
 org.testcontainers.containers.TimescaleDBContainerProvider
+org.testcontainers.containers.PgVectorContainerProvider
diff --git a/modules/postgresql/src/test/java/org/testcontainers/jdbc/pgvector/PgVectorJDBCDriverTest.java b/modules/postgresql/src/test/java/org/testcontainers/jdbc/pgvector/PgVectorJDBCDriverTest.java
new file mode 100644
index 00000000000..2d7e67f4923
--- /dev/null
+++ b/modules/postgresql/src/test/java/org/testcontainers/jdbc/pgvector/PgVectorJDBCDriverTest.java
@@ -0,0 +1,28 @@
+package org.testcontainers.jdbc.pgvector;
+
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.testcontainers.jdbc.AbstractJDBCDriverTest;
+
+import java.util.Arrays;
+import java.util.EnumSet;
+
+@RunWith(Parameterized.class)
+public class PgVectorJDBCDriverTest extends AbstractJDBCDriverTest {
+
+    @Parameterized.Parameters(name = "{index} - {0}")
+    public static Iterable<Object[]> data() {
+        return Arrays.asList(
+            new Object[][] {
+                {
+                    "jdbc:tc:pgvector://hostname/databasename?user=someuser&password=somepwd",
+                    EnumSet.of(Options.JDBCParams),
+                },
+                {
+                    "jdbc:tc:pgvector:pg14://hostname/databasename?user=someuser&password=somepwd",
+                    EnumSet.of(Options.JDBCParams),
+                },
+            }
+        );
+    }
+}