org.apache.beam.it.jdbc.MSSQLResourceManager Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of beam-it-jdbc Show documentation
Integration test utilities suites.
There is a newer version: 2.61.0
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.it.jdbc;

import java.sql.Connection;
import java.sql.Statement;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testcontainers.containers.MSSQLServerContainer;
import org.testcontainers.utility.DockerImageName;

/**
 * Default class for the MS SQL implementation of {@link AbstractJDBCResourceManager} abstract
 * class.
 *
 * The class supports one database, and multiple tables per database object. A database is *
 * created when the container first spins up, if one is not given.
 *
 * The class is thread-safe.
 */
public class MSSQLResourceManager
    extends AbstractJDBCResourceManager> {

  private static final Logger LOG = LoggerFactory.getLogger(MSSQLResourceManager.class);

  private static final String DEFAULT_MSSQL_CONTAINER_NAME = "mcr.microsoft.com/azure-sql-edge";

  // A list of available MSSQL Docker image tags can be found at
  // https://hub.docker.com/_/microsoft-azure-sql-edge
  private static final String DEFAULT_MSSQL_CONTAINER_TAG = "1.0.6";

  private boolean initialized;

  @SuppressWarnings("nullness")
  private MSSQLResourceManager(Builder builder) {
    super(
        new DefaultMSSQLServerContainer<>(
            DockerImageName.parse(builder.containerImageName)
                .withTag(builder.containerImageTag)
                .asCompatibleSubstituteFor("mcr.microsoft.com/mssql/server"),
            builder.databaseName,
            builder.useStaticContainer),
        builder);
    createDatabase(builder.databaseName);
  }

  @VisibleForTesting
  > MSSQLResourceManager(T container, Builder builder) {
    super(container, builder);
    initialized = true;
  }

  public static Builder builder(String testId) {
    return new Builder(testId);
  }

  private synchronized void createDatabase(String databaseName) {
    LOG.info("Creating database using databaseName '{}'.", databaseName);

    StringBuilder sql = new StringBuilder();
    try (Connection con = driver.getConnection(getUri(), username, password)) {
      Statement stmt = con.createStatement();
      sql.append("CREATE DATABASE ").append(databaseName);
      stmt.executeUpdate(sql.toString());
      stmt.close();
    } catch (Exception e) {
      throw new JDBCResourceManagerException(
          "Error creating database with SQL statement: " + sql, e);
    }

    initialized = true;
    LOG.info("Successfully created database {}.{}", databaseName, databaseName);
  }

  @Override
  public synchronized String getUri() {
    return String.format(
        "jdbc:%s://%s:%d%s%s",
        getJDBCPrefix(),
        this.getHost(),
        this.getPort(getJDBCPort()),
        initialized ? ";DatabaseName=" + databaseName : "",
        ";encrypt=true;trustServerCertificate=true;");
  }

  @Override
  public String getJDBCPrefix() {
    return "sqlserver";
  }

  @Override
  protected int getJDBCPort() {
    return DefaultMSSQLServerContainer.MS_SQL_SERVER_PORT;
  }

  @Override
  protected String getFirstRow(String tableName) {
    return "SELECT TOP 1 * FROM " + tableName;
  }

  /** Builder for {@link MSSQLResourceManager}. */
  public static final class Builder
      extends AbstractJDBCResourceManager.Builder> {

    public Builder(String testId) {
      super(testId, DEFAULT_MSSQL_CONTAINER_NAME, DEFAULT_MSSQL_CONTAINER_TAG);
    }

    @Override
    public MSSQLResourceManager build() {
      return new MSSQLResourceManager(this);
    }
  }

  static final class DefaultMSSQLServerContainer>
      extends MSSQLServerContainer {

    private final String databaseName;
    private final boolean usingStaticDatabase;

    @SuppressWarnings("nullness")
    DefaultMSSQLServerContainer(
        DockerImageName dockerImageName, String databaseName, boolean usingStaticDatabase) {
      super(dockerImageName);
      this.addEnv("ACCEPT_EULA", "Y");
      this.databaseName = databaseName;
      this.usingStaticDatabase = usingStaticDatabase;
    }

    @Override
    public T withUsername(String username) {
      if (!username.equals(AbstractJDBCResourceManager.DEFAULT_JDBC_USERNAME)
          && !usingStaticDatabase) {
        throw new UnsupportedOperationException("The username cannot be overridden to " + username);
      }
      return this.self();
    }

    @Override
    public T withDatabaseName(String dbName) {
      // AbstractJDBCResourceManager constructor calls this method, but MSSQLResourceManager
      // does
      // database creation after class initialization, so just pass container instance back as-is.
      return this.self();
    }

    @Override
    public String getDatabaseName() {
      return databaseName;
    }

    @Override
    @SuppressWarnings("nullness")
    public boolean equals(Object o) {
      return super.equals(o);
    }

    @Override
    public int hashCode() {
      return super.hashCode();
    }
  }
}