All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.vertx.ext.jdbc.impl.actions.JDBCSQLRowStream Maven / Gradle / Ivy

There is a newer version: 5.0.0.CR2
Show newest version
/*
 * Copyright (c) 2011-2014 The original author or authors
 * ------------------------------------------------------
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * and Apache License v2.0 which accompanies this distribution.
 *
 *     The Eclipse Public License is available at
 *     http://www.eclipse.org/legal/epl-v10.html
 *
 *     The Apache License v2.0 is available at
 *     http://www.opensource.org/licenses/apache2.0.php
 *
 * You may elect to redistribute this code under either of these licenses.
 */
package io.vertx.ext.jdbc.impl.actions;

import io.vertx.core.AsyncResult;
import io.vertx.core.Future;
import io.vertx.core.Handler;
import io.vertx.core.Promise;
import io.vertx.core.impl.ContextInternal;
import io.vertx.core.impl.TaskQueue;
import io.vertx.core.impl.logging.Logger;
import io.vertx.core.impl.logging.LoggerFactory;
import io.vertx.core.json.JsonArray;
import io.vertx.core.streams.ReadStream;
import io.vertx.core.streams.impl.InboundBuffer;
import io.vertx.ext.jdbc.spi.JDBCDecoder;
import io.vertx.ext.sql.SQLRowStream;
import io.vertx.ext.jdbc.spi.JDBCColumnDescriptorProvider;

import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author Paulo Lopes
 */
class JDBCSQLRowStream implements SQLRowStream {

  private static final Logger log = LoggerFactory.getLogger(JDBCSQLRowStream.class);

  private final ContextInternal ctx;
  private final TaskQueue statementsQueue;
  private final Statement st;
  private final int fetchSize;
  private final InboundBuffer pending;
  private final AtomicBoolean ended = new AtomicBoolean(false);
  private final AtomicBoolean stClosed = new AtomicBoolean(false);
  private final AtomicBoolean rsClosed = new AtomicBoolean(false);
  private final AtomicBoolean more = new AtomicBoolean(false);

  private ResultSet rs;
  private ResultSetMetaData metaData;
  private JDBCDecoder decoder;
  private List columns;
  private int cols;

  private Handler exceptionHandler;
  private Handler endHandler;
  private Handler rsClosedHandler;

  JDBCSQLRowStream(ContextInternal ctx, TaskQueue statementsQueue, Statement st, ResultSet rs, JDBCDecoder decoder, int fetchSize) throws SQLException {
    this.ctx = ctx;
    this.statementsQueue = statementsQueue;
    this.st = st;
    this.fetchSize = fetchSize;
    this.rs = rs;
    this.decoder = decoder;
    this.pending = new InboundBuffer(ctx, fetchSize)
      .drainHandler(v -> readBatch())
      .emptyHandler(v -> checkEndHandler());

    metaData = rs.getMetaData();
    cols = metaData.getColumnCount();
    stClosed.set(false);
    rsClosed.set(false);
    // the first rs is populated in the constructor
    more.set(true);
  }

  private void checkEndHandler() {
    if (ended.get() && pending.isEmpty()) {
      if (endHandler != null) {
        endHandler.handle(null);
      }
    }
  }

  @Override
  public int column(String name) {
    try {
      return rs.findColumn(name) - 1;
    } catch (SQLException e) {
      return -1;
    }
  }

  @Override
  public List columns() {
    if (columns == null) {
      try {
        if (cols > 0) {
          final List columns = new ArrayList<>(cols);
          for (int i = 0; i < cols; i++) {
            columns.add(i, metaData.getColumnLabel(i + 1));
          }
          this.columns = Collections.unmodifiableList(columns);
        } else {
          this.columns = Collections.emptyList();
        }
      } catch (SQLException e) {
        throw new RuntimeException(e);
      }
    }

    return columns;
  }

  @Override
  public SQLRowStream exceptionHandler(Handler handler) {
    this.exceptionHandler = handler;
    return this;
  }

  @Override
  public SQLRowStream handler(Handler handler) {
    pending.handler(handler);
    if (handler != null && pending.isWritable()) {
      readBatch();
    }
    return this;
  }

  @Override
  public SQLRowStream pause() {
    pending.pause();
    return this;
  }

  @Override
  public ReadStream fetch(long amount) {
    pending.fetch(amount);
    return this;
  }

  @Override
  public SQLRowStream resume() {
    pending.resume();
    return this;
  }

  private void readBatch() {
    if (!rsClosed.get()) {
      ctx.>executeBlocking(fut -> {
        try {
          JDBCColumnDescriptorProvider provider = JDBCColumnDescriptorProvider.fromResultMetaData(metaData);
          List rows = new ArrayList<>(fetchSize);
          for (int i = 0; i < fetchSize && rs.next(); i++) {
            JsonArray result = new JsonArray();
            for (int j = 1; j <= cols; j++) {
              Object res = decoder.parse(rs, j, provider);
              if (res != null) {
                result.add(res);
              } else {
                result.addNull();
              }
            }
            rows.add(result);
          }
          fut.complete(rows);
        } catch (SQLException e) {
          fut.fail(e);
        }
      }, statementsQueue, ar -> {
        if (ar.succeeded()) {
          List rows = ar.result();
          if (rows.isEmpty()) {
            empty(null);
          } else if (pending.write(rows)) {
            readBatch();
          }
        } else {
          empty(ar.cause());
        }
      });
    }
  }

  private void empty(Throwable err) {
    // mark as ended if the handler was registered too late
    ended.set(true);
    // automatically close resources

    if (rsClosedHandler != null) {
      // only close the result set and notify
      closeRS(pending.isEmpty(), c -> {
        if (err != null) {
          if (exceptionHandler != null) {
            exceptionHandler.handle(err);
          } else {
            log.debug(err);
          }
        } else {
          rsClosedHandler.handle(null);
        }
      });
    } else {
      // default behavior close result set + statement
      closeAll(pending.isEmpty(), c -> {
        if (err != null) {
          if (exceptionHandler != null) {
            exceptionHandler.handle(err);
          } else {
            log.debug(err);
          }
        } else {
          checkEndHandler();
        }
      });
    }
  }

  @Override
  public SQLRowStream endHandler(Handler handler) {
    this.endHandler = handler;
    checkEndHandler();
    return this;
  }

  private void closeRS(boolean pause, Handler> handler) {
    // stop pumping data if needed
    if (pause) {
      pause();
    }
    // close the cursor
    close(rs, rsClosed, handler);
  }

  @Override
  public void close() {
    close(null);
  }

  @Override
  public void close(Handler> handler) {
    closeAll(true, handler);
  }

  private void closeAll(boolean pauseBuffer, Handler> handler) {
    closeRS(pauseBuffer, res -> {
      // close the statement
      close(st, stClosed, handler);
    });
  }

  @Override
  public SQLRowStream resultSetClosedHandler(Handler handler) {
    this.rsClosedHandler = handler;
    return this;
  }

  @Override
  public void moreResults() {
    if (more.compareAndSet(true, false)) {
      // pause streaming if rs is not complete
      pause();

      ctx.executeBlocking(this::getNextResultSet, statementsQueue, res -> {
        if (res.failed()) {
          if (exceptionHandler != null) {
            exceptionHandler.handle(res.cause());
          } else {
            log.debug(res.cause());
          }
        } else {
          if (more.get()) {
            resume();
          } else {
            if (endHandler != null) {
              endHandler.handle(null);
            }
          }
        }
      });
    }
  }

  private void getNextResultSet(Promise f) {
    try {
      // close if not already closed
      if (rsClosed.compareAndSet(false, true)) {
        rs.close();
      }
      // is there more rs data?
      if (st.getMoreResults()) {
        rs = st.getResultSet();
        metaData = rs.getMetaData();
        cols = metaData.getColumnCount();
        columns = null;
        // reset
        // paused.set(true);
        stClosed.set(false);
        rsClosed.set(false);
        more.set(true);
      }

      f.complete();
    } catch (SQLException e) {
      f.fail(e);
    }
  }

  private void close(AutoCloseable closeable, AtomicBoolean lock, Handler> handler) {
    if (lock.compareAndSet(false, true)) {
      ctx.executeBlocking(f -> {
        try {
          closeable.close();
          f.complete();
        } catch (Exception e) {
          f.fail(e);
        }
      }, statementsQueue, handler);
    } else {
      if (handler != null) {
        handler.handle(Future.succeededFuture());
      }
    }
  }
}




© 2015 - 2025 Weber Informatics LLC | Privacy Policy