content: /*
   * 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 static org.apache.beam.it.jdbc.JDBCResourceManagerUtils.checkValidTableName;
  import static org.apache.beam.it.jdbc.JDBCResourceManagerUtils.generateDatabaseName;
  import static org.apache.beam.it.jdbc.JDBCResourceManagerUtils.generateJdbcPassword;
  
  import java.sql.Connection;
  import java.sql.ResultSet;
  import java.sql.ResultSetMetaData;
  import java.sql.SQLException;
  import java.sql.Statement;
  import java.util.ArrayList;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import org.apache.beam.it.testcontainers.TestContainerResourceManager;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
  import org.apache.commons.lang3.math.NumberUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.testcontainers.containers.JdbcDatabaseContainer;
  
  /**
   * Abstract class for implementation of {@link JDBCResourceManager} interface.
   *
   * <p>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.
   *
   * <p>The database name is formed using testId. The database name will be "{testId}-{ISO8601 time,
   * microsecond precision}", with additional formatting.
   *
   * <p>The class is thread-safe.
   */
  public abstract class AbstractJDBCResourceManager<T extends JdbcDatabaseContainer<?>>
      extends TestContainerResourceManager<JdbcDatabaseContainer<?>> implements JDBCResourceManager {
  
    private static final Logger LOG = LoggerFactory.getLogger(AbstractJDBCResourceManager.class);
  
    protected static final String DEFAULT_JDBC_USERNAME = "root";
  
    protected final JDBCDriverFactory driver;
    protected final String databaseName;
    protected final String username;
    protected final String password;
  
    private final Map<String, String> tableIds;
  
    @VisibleForTesting
    AbstractJDBCResourceManager(T container, Builder<T> builder, JDBCDriverFactory driver) {
      super(
          container
              .withUsername(builder.username)
              .withPassword(builder.password)
              .withDatabaseName(builder.databaseName),
          builder);
  
      this.databaseName = container.getDatabaseName();
      this.username = container.getUsername();
      this.password = container.getPassword();
      this.tableIds = new HashMap<>();
      this.driver = driver;
    }
  
    protected AbstractJDBCResourceManager(T container, Builder<T> builder) {
      this(container, builder, new JDBCDriverFactory());
    }
  
    /**
     * Return the default port that this JDBC implementation listens on.
     *
     * @return the JDBC port.
     */
    protected abstract int getJDBCPort();
  
    @Override
    public String getUsername() {
      return username;
    }
  
    @Override
    public String getPassword() {
      return password;
    }
  
    public int getPort() {
      return this.getPort(getJDBCPort());
    }
  
    @Override
    public synchronized String getUri() {
      return String.format(
          "jdbc:%s://%s:%d/%s",
          getJDBCPrefix(), this.getHost(), this.getPort(), this.getDatabaseName());
    }
  
    public abstract String getJDBCPrefix();
  
    @Override
    public synchronized String getDatabaseName() {
      return databaseName;
    }
  
    @Override
    public boolean createTable(String tableName, JDBCSchema schema) {
      // Check table ID
      checkValidTableName(tableName);
  
      // Check if table already exists
      if (tableIds.containsKey(tableName)) {
        throw new IllegalStateException(
            "Table " + tableName + " already exists for database " + databaseName + ".");
      }
  
      LOG.info("Creating table using tableName '{}'.", tableName);
  
      StringBuilder sql = new StringBuilder();
      try (Connection con = driver.getConnection(getUri(), username, password)) {
        Statement stmt = con.createStatement();
        sql.append("CREATE TABLE ")
            .append(tableName)
            .append(" (")
            .append(schema.toSqlStatement())
            .append(")");
  
        stmt.executeUpdate(sql.toString());
        stmt.close();
      } catch (Exception e) {
        throw new JDBCResourceManagerException(
            "Error creating table with SQL statement: "
                + sql
                + " (for connection with URL "
                + getUri()
                + ")",
            e);
      }
  
      tableIds.put(tableName, schema.getIdColumn());
      LOG.info("Successfully created table {}.{}", databaseName, tableName);
  
      return true;
    }
  
    /**
     * Writes the given mapped rows into the specified columns. This method requires {@link
     * JDBCResourceManager#createTable(String, JDBCSchema)} to be called for the target table
     * beforehand.
     *
     * <p>The rows map must use the row id as the key, and the values will be inserted into the
     * columns at the row with that id. i.e. {0: [val1, val2, ...], 1: [val1, val2, ...], ...}
     *
     * @param tableName The name of the table to insert the given rows into.
     * @param rows A map representing the rows to be inserted into the table.
     * @throws JDBCResourceManagerException if method is called after resources have been cleaned up,
     *     if the manager object has no dataset, if the table does not exist or if there is an
     *     Exception when attempting to insert the rows.
     */
    @Override
    @SuppressWarnings("nullness")
    public boolean write(String tableName, List<Map<String, Object>> rows)
        throws JDBCResourceManagerException {
      if (rows.size() == 0) {
        return false;
      }
  
      LOG.info("Attempting to write {} rows to {}.{}.", rows.size(), databaseName, tableName);
  
      try (Connection con = driver.getConnection(getUri(), username, password)) {
        Statement stmt = con.createStatement();
  
        for (Map<String, Object> row : rows) {
          List<String> columns = new ArrayList<>(row.keySet());
          StringBuilder sql =
              new StringBuilder("INSERT INTO ")
                  .append(tableName)
                  .append("(")
                  .append(String.join(",", columns))
                  .append(") VALUES (");
  
          List<String> valueList = new ArrayList<>();
          for (String colName : columns) {
            Object value = row.get(colName);
            if (value == null) {
              valueList.add(null);
            } else if (NumberUtils.isCreatable(value.toString())
                || "true".equalsIgnoreCase(value.toString())
                || "false".equalsIgnoreCase(value.toString())
                || value.toString().startsWith("ARRAY[")) {
              valueList.add(String.valueOf(value));
            } else {
              valueList.add("'" + value + "'");
            }
          }
          sql.append(String.join(",", valueList)).append(")");
  
          try {
            LOG.info("Running SQL statement: " + sql);
            stmt.executeUpdate(sql.toString());
          } catch (SQLException e) {
            throw new JDBCResourceManagerException(
                "Failed to insert values into table with SQL statement: " + sql, e);
          }
        }
        stmt.close();
      } catch (SQLException e) {
        throw new JDBCResourceManagerException(
            String.format("Exception occurred when trying to write records to %s.", tableName), e);
      }
  
      LOG.info("Successfully wrote {} rows to {}.{}.", rows.size(), databaseName, tableName);
      return true;
    }
  
    @Override
    @SuppressWarnings("nullness")
    public List<Map<String, Object>> readTable(String tableName) {
      LOG.info("Reading all rows from {}.{}", databaseName, tableName);
      List<Map<String, Object>> result = runSQLQuery(String.format("SELECT * FROM %s", tableName));
      LOG.info("Successfully loaded rows from {}.{}", databaseName, tableName);
      return result;
    }
  
    @Override
    public synchronized List<String> getTableSchema(String tableName) {
      String sql = "";
      try (Connection con = driver.getConnection(getUri(), username, password)) {
        Statement stmt = con.createStatement();
        sql = getFirstRow(tableName);
        ResultSet result = stmt.executeQuery(sql);
  
        ResultSetMetaData metadata = result.getMetaData();
        List<String> columnNames = new ArrayList<>();
        // Columns list in table metadata is 1-indexed
        for (int i = 1; i <= metadata.getColumnCount(); i++) {
          columnNames.add(metadata.getColumnName(i));
        }
        result.close();
        stmt.close();
        return columnNames;
      } catch (Exception e) {
        throw new JDBCResourceManagerException(
            "Failed to fetch table schema. SQL statement: " + sql, e);
      }
    }
  
    /**
     * Retrieves the first row from the table.
     *
     * @param tableName the name of the table to query.
     * @return the first row from the table.
     */
    protected String getFirstRow(String tableName) {
      return "SELECT * FROM " + tableName + " LIMIT 1";
    }
  
    @Override
    @SuppressWarnings("nullness")
    public synchronized List<Map<String, Object>> runSQLQuery(String sql) {
      try (Statement stmt = driver.getConnection(getUri(), username, password).createStatement()) {
        List<Map<String, Object>> result = new ArrayList<>();
        ResultSet resultSet = stmt.executeQuery(sql);
        while (resultSet.next()) {
          Map<String, Object> row = new HashMap<>();
          ResultSetMetaData metadata = resultSet.getMetaData();
          // Columns list in table metadata is 1-indexed
          for (int i = 1; i <= metadata.getColumnCount(); i++) {
            row.put(metadata.getColumnName(i), resultSet.getObject(i));
          }
          result.add(row);
        }
        return result;
      } catch (Exception e) {
        throw new JDBCResourceManagerException("Failed to execute SQL statement: " + sql, e);
      }
    }
  
    @Override
    public synchronized void runSQLUpdate(String sql) {
      try (Statement stmt = driver.getConnection(getUri(), username, password).createStatement()) {
        stmt.executeUpdate(sql);
      } catch (Exception e) {
        throw new JDBCResourceManagerException("Failed to execute SQL statement: " + sql, e);
      }
    }
  
    @Override
    public synchronized long getRowCount(String tableName) {
      try (Connection con = driver.getConnection(getUri(), username, password)) {
        Statement stmt = con.createStatement();
        ResultSet resultSet = stmt.executeQuery(String.format("SELECT count(*) FROM %s", tableName));
        resultSet.next();
        long rows = resultSet.getLong(1);
        resultSet.close();
        stmt.close();
        return rows;
      } catch (Exception e) {
        throw new JDBCResourceManagerException("Failed to get row count from " + tableName, e);
      }
    }
  
    /**
     * Builder for {@link AbstractJDBCResourceManager}.
     *
     * @param <T> A class that extends {@link JdbcDatabaseContainer} for specific JDBC
     *     implementations.
     */
    public abstract static class Builder<T extends JdbcDatabaseContainer<?>>
        extends TestContainerResourceManager.Builder<AbstractJDBCResourceManager<T>> {
  
      protected String databaseName;
      protected String username;
      protected String password;
  
      public Builder(String testId, String containerImageName, String containerImageTag) {
        super(testId, containerImageName, containerImageTag);
  
        this.username = DEFAULT_JDBC_USERNAME;
        this.password = generateJdbcPassword();
        this.databaseName = generateDatabaseName(testId);
      }
  
      /**
       * Sets the database name to that of a static database instance. Use this method only when
       * attempting to operate on a pre-existing JDBC database.
       *
       * @param databaseName The database name.
       * @return this builder object with the database name set.
       */
      public Builder<T> setDatabaseName(String databaseName) {
        this.databaseName = databaseName;
        return this;
      }
  
      /**
       * Manually set the JDBC database username to the given username. This username will be used by
       * the resource manager to authenticate with the JDBC database.
       *
       * @param username the username for the JDBC database.
       * @return this builder with the username manually set.
       */
      public Builder<T> setUsername(String username) {
        this.username = username;
        return this;
      }
  
      /**
       * Manually set the JDBC database password to the given password. This password will be used by
       * the resource manager to authenticate with the JDBC database.
       *
       * @param password the password for the JDBC database.
       * @return this builder with the password manually set.
       */
      public Builder<T> setPassword(String password) {
        this.password = password;
        return this;
      }
    }
  }
  