private Iterator fetchIdentifiers()

in google-cloud-spanner-hibernate-tools/src/main/java/com/google/cloud/spanner/hibernate/PooledBitReversedSequenceStyleGenerator.java [397:464]


  private Iterator<Long> fetchIdentifiers(SharedSessionContractImplementor session)
      throws HibernateException {
    // Prefix all 'set ...' statements with 'spanner.' if the dialect is PostgreSQL.
    // The safest way to determine that is by looking at the quote character for identifiers.
    String extensionPrefix = dialect.openQuote() == '"' ? "spanner." : "";
    Connection connection = null;
    Boolean retryAbortsInternally = null;
    try {
      // Use a separate connection to get new sequence values. This ensures that it also uses a
      // separate read/write transaction, which again means that it will not interfere with any
      // retries of the actual business transaction.
      connection = session.getJdbcConnectionAccess().obtainConnection();
      connection.setAutoCommit(false);
      try (Statement statement = connection.createStatement()) {
        // TODO: Use 'set local spanner.retry_aborts_internally=false' when that has been
        //       implemented.
        retryAbortsInternally = isRetryAbortsInternally(statement);
        connection.commit();
        statement.execute(String.format("set %sretry_aborts_internally=false", extensionPrefix));
        List<Long> identifiers = new ArrayList<>(this.fetchSize);
        try (ResultSet resultSet = statement.executeQuery(this.select)) {
          while (resultSet.next()) {
            for (int col = 1; col <= resultSet.getMetaData().getColumnCount(); col++) {
              identifiers.add(resultSet.getLong(col));
            }
          }
        }
        // Do a rollback instead of a commit here because:
        // 1. We have only accessed a bit-reversed sequence during the transaction.
        // 2. Committing or rolling back the transaction does not make any difference for the
        //    sequence. Its state has been updated in both cases.
        // 3. Committing the transaction on the emulator would cause it to be aborted, as the
        //    emulator only supports one transaction at any time. Rolling back is however allowed.
        connection.rollback();
        return identifiers.iterator();
      }
    } catch (SQLException sqlException) {
      if (connection != null) {
        ignoreSqlException(connection::rollback);
      }
      if (isAbortedError(sqlException)) {
        // Return an empty iterator to force a retry.
        return EMPTY_ITERATOR;
      }
      throw session
          .getJdbcServices()
          .getSqlExceptionHelper()
          .convert(sqlException, "could not get next sequence values", this.select);
    } finally {
      if (connection != null) {
        Connection finalConnection = connection;
        if (retryAbortsInternally != null) {
          Boolean finalRetryAbortsInternally = retryAbortsInternally;
          ignoreSqlException(
              () ->
                  finalConnection
                      .createStatement()
                      .execute(
                          String.format(
                              "set %sretry_aborts_internally=%s",
                              extensionPrefix, finalRetryAbortsInternally)));
          ignoreSqlException(connection::commit);
        }
        ignoreSqlException(
            () -> session.getJdbcConnectionAccess().releaseConnection(finalConnection));
      }
    }
  }