Merge pull request #133 from Ryszard-Trojnacki/feature/pgvector-container

Add PGvector test container

Author: rbygrave

README.md

@@ -164,6 +164,19 @@ occurring automatically on JVM shutdown.
 
 ```
 
+#### Postgres PGvector
+```java
+    var container = PGvectorContainer.builder("pg18")
+      .start();
+```
+
+#### Postgres Postgis
+```java
+    PostgisContainer container = PostgisContainer.builder("15")
+      .useLW(true) // use LW compression JDBC urls
+      .build();
+```
+
 #### SqlServer
 
 ```java

pom.xml

@@ -151,6 +151,13 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>com.pgvector</groupId>
+      <artifactId>pgvector</artifactId>
+      <version>0.1.6</version>
+      <scope>test</scope>
+    </dependency>
+
     <dependency>
       <groupId>ru.yandex.clickhouse</groupId>
       <artifactId>clickhouse-jdbc</artifactId>

src/main/java/io/ebean/test/containers/ContainerFactory.java

@@ -61,6 +61,10 @@ private void init() {
     if (pgisVersion != null) {
       containers.add(PostgisContainer.builder(pgisVersion).properties(properties).build());
     }
+    String pgvectorVersion = runWithVersion("pgvector");
+    if (pgvectorVersion != null) {
+      containers.add(PGvectorContainer.builder(pgvectorVersion).properties(properties).build());
+    }
     String mysqlVersion = runWithVersion("mysql");
     if (mysqlVersion != null) {
       containers.add(MySqlContainer.builder(mysqlVersion).properties(properties).build());

src/main/java/io/ebean/test/containers/PGvectorContainer.java

@@ -0,0 +1,65 @@
+package io.ebean.test.containers;
+
+/**
+ * Commands for controlling a pgvector docker container.
+ */
+public class PGvectorContainer extends BasePostgresContainer<PGvectorContainer> {
+
+  @Override
+  public PGvectorContainer start() {
+    startOrThrow();
+    return this;
+  }
+
+  /**
+   * Create a builder for PGvectorContainer.
+   */
+  public static Builder builder(String version) {
+    return new Builder(version);
+  }
+
+  private PGvectorContainer(Builder config) {
+    super(config);
+  }
+
+  /**
+   * Builder for Postgis container.
+   */
+  public static class Builder extends BaseDbBuilder<PGvectorContainer, Builder> {
+
+    private Builder(String version) {
+      super("pgvector", 6435, 5432, version);
+      this.image = "pgvector/pgvector:" + version;
+      this.adminUsername = "postgres";
+      this.tmpfs = "/var/lib/postgresql/data:rw";
+      this.extensions = "vector";
+      this.extra.extensions = extensions;
+      this.extra2.extensions = extensions;
+    }
+
+    @Override
+    protected String buildJdbcUrl() {
+      return "jdbc:postgresql://" + host + ":" + port + "/" + dbName;
+    }
+
+    @Override
+    protected String buildJdbcAdminUrl() {
+      return "jdbc:postgresql://" + host + ":" + port + "/postgres";
+    }
+
+    @Override
+    protected String buildExtraJdbcUrl(String dbName) {
+      return "jdbc:postgresql://" + host + ":" + port + "/" + dbName;
+    }
+
+    @Override
+    public PGvectorContainer build() {
+      return new PGvectorContainer(this);
+    }
+
+    @Override
+    public PGvectorContainer start() {
+      return build().start();
+    }
+  }
+}

src/test/java/io/ebean/test/containers/PGvectorContainerTest.java

@@ -0,0 +1,82 @@
+package io.ebean.test.containers;
+
+import io.ebean.Database;
+import io.ebean.datasource.DataSourcePool;
+import org.junit.jupiter.api.Test;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
+import java.util.HashSet;
+
+import com.pgvector.PGvector;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class PGvectorContainerTest {
+  private final HashSet<Connection> connections = new HashSet<>();
+
+  /**
+   * Helper function to register the PGvector types only once per connection.
+   * @param connection the connection
+   * @return the same connection
+   * @throws SQLException if an error occurs
+   */
+  private Connection wrapConnection(Connection connection) throws SQLException {
+    if(connections.add(connection)) {
+      PGvector.registerTypes(connection);
+    }
+    return connection;
+  }
+
+  @Test
+  void extraDb() throws java.sql.SQLException {
+    PGvectorContainer container = PGvectorContainer.builder("pg18")
+      .port(0)
+      .extraDb("myextra")
+      .build();
+
+    container.startMaybe();
+    assertThat(container.port()).isGreaterThan(0);
+
+    ContainerConfig containerConfig = container.config();
+    assertThat(containerConfig.port()).isEqualTo(container.port());
+
+    String jdbcUrl = container.config().jdbcUrl();
+    assertThat(jdbcUrl).contains(":" + containerConfig.port());
+    runSomeSql(container);
+
+    DataSourcePool dataSource = container.ebean().dataSourceBuilder().build();
+    try (Connection connection = wrapConnection(dataSource.getConnection())) {
+      exeSql(connection, "INSERT INTO items (embedding) values ('[7,8,9]')");
+    }
+    dataSource.shutdown();
+
+    Database ebean = container.ebean().builder()
+      .register(false)
+      .defaultDatabase(false)
+      .build();
+    // This can't be done yet, because Ebean doesn't know about PGvector type
+//    ebean.sqlUpdate("insert into items (embedding) values (?)")
+//      .setParameter(new PGvector(new float[] { 10f, 11f, 12f }))
+//      .execute();
+
+    ebean.shutdown();
+  }
+
+  private void runSomeSql(PGvectorContainer container) {
+    try {
+      Connection connection = wrapConnection(container.createConnection());
+      exeSql(connection, "CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))");
+      exeSql(connection, "INSERT INTO items (embedding) VALUES ('[1,2,3]'), ('[4,5,6]')");
+    } catch (SQLException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private static void exeSql(Connection connection, String sql) throws SQLException {
+    try (PreparedStatement st = connection.prepareStatement(sql)) {
+      st.execute();
+    }
+  }
+}