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

com.github.davidmoten.rx.jdbc.QueryUpdateOnSubscribe Maven / Gradle / Ivy

There is a newer version: 0.7.19
Show newest version
package com.github.davidmoten.rx.jdbc;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.List;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.github.davidmoten.rx.jdbc.exceptions.SQLRuntimeException;

import rx.Observable;
import rx.Observable.OnSubscribe;
import rx.Subscriber;
import rx.Subscription;
import rx.exceptions.Exceptions;
import rx.functions.Action0;
import rx.subscriptions.Subscriptions;

/**
 * Executes the update query.
 */
final class QueryUpdateOnSubscribe implements OnSubscribe {

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

    static final String BEGIN_TRANSACTION = "begin";

    /**
     * Special sql command that brings about a rollback.
     */
    static final String ROLLBACK = "rollback";

    /**
     * Special sql command that brings about a commit.
     */
    static final String COMMIT = "commit";

    /**
     * Returns an Observable of the results of pushing one set of parameters
     * through a select query.
     * 
     * @param params
     *            one set of parameters to be run with the query
     * @return
     */
    static  Observable create(QueryUpdate query, List parameters) {
        return Observable.create(new QueryUpdateOnSubscribe(query, parameters));
    }

    /**
     * The query to be executed.
     */
    private final QueryUpdate query;

    /**
     * The parameters to run the query against (may be a subset of the query
     * parameters specified in the query because the query may be run multiple
     * times with multiple sets of parameters).
     */
    private final List parameters;

    /**
     * Constructor.
     * 
     * @param query
     * @param parameters
     */
    private QueryUpdateOnSubscribe(QueryUpdate query, List parameters) {
        this.query = query;
        this.parameters = parameters;
    }

    @Override
    public void call(Subscriber subscriber) {
        final State state = new State();
        try {
            if (isBeginTransaction())
                performBeginTransaction(subscriber);
            else {
                query.context().setupBatching();
                getConnection(state);
                subscriber.add(createUnsubscriptionAction(state));
                if (isCommit())
                    performCommit(subscriber, state);
                else if (isRollback())
                    performRollback(subscriber, state);
                else
                    performUpdate(subscriber, state);
            }
        } catch (Throwable e) {
            query.context().endTransactionObserve();
            query.context().endTransactionSubscribe();
            try {
                close(state);
            } finally {
                handleException(e, subscriber);
            }
        }
    }

    private Subscription createUnsubscriptionAction(final State state) {
        return Subscriptions.create(new Action0() {
            @Override
            public void call() {
                close(state);
            }
        });
    }

    private boolean isBeginTransaction() {
        return query.sql().equals(BEGIN_TRANSACTION);
    }

    @SuppressWarnings("unchecked")
    private void performBeginTransaction(Subscriber subscriber) {
        query.context().beginTransactionObserve();
        debug("beginTransaction emitting 1");
        subscriber.onNext((T) Integer.valueOf(1));
        debug("emitted 1");
        complete(subscriber);
    }

    /**
     * Gets the current connection.
     */
    private void getConnection(State state) {
        state.con = query.context().connectionProvider().get();
        debug("getting connection");
        debug("cp={}", query.context().connectionProvider());
    }

    /**
     * Returns true if and only if the sql statement is a commit command.
     * 
     * @return if is commit
     */
    private boolean isCommit() {
        return query.sql().equals(COMMIT);
    }

    /**
     * Returns true if and only if the sql statement is a rollback command.
     * 
     * @return if is rollback
     */
    private boolean isRollback() {
        return query.sql().equals(ROLLBACK);
    }

    /**
     * Commits the current transaction. Throws {@link RuntimeException} if
     * connection is in autoCommit mode.
     * 
     * @param subscriber
     * @param state
     */
    @SuppressWarnings("unchecked")
    private void performCommit(Subscriber subscriber, State state) {
        getConnection(state);
        query.context().endTransactionObserve();
        if (subscriber.isUnsubscribed())
            return;

        debug("committing");
        Conditions.checkTrue(!Util.isAutoCommit(state.con));
        Util.commit(state.con);
        // must close before onNext so that connection is released and is
        // available to a query that might process the onNext
        close(state);

        if (subscriber.isUnsubscribed())
            return;

        subscriber.onNext((T) Integer.valueOf(1));
        debug("committed");
        complete(subscriber);
    }

    /**
     * Rolls back the current transaction. Throws {@link RuntimeException} if
     * connection is in autoCommit mode.
     * 
     * @param subscriber
     * @param state
     */
    @SuppressWarnings("unchecked")
    private void performRollback(Subscriber subscriber, State state) {
        debug("rolling back");
        query.context().endTransactionObserve();
        Conditions.checkTrue(!Util.isAutoCommit(state.con));
        Util.rollback(state.con);
        // must close before onNext so that connection is released and is
        // available to a query that might process the onNext
        close(state);
        subscriber.onNext((T) Integer.valueOf(0));
        debug("rolled back");
        complete(subscriber);
    }

    /**
     * Executes the prepared statement.
     * 
     * @param subscriber
     * 
     * @throws SQLException
     */
    @SuppressWarnings("unchecked")
    private void performUpdate(final Subscriber subscriber, State state)
            throws SQLException {
        if (subscriber.isUnsubscribed()) {
            return;
        }
        if (query.context().batchSize() > 1 && !query.context().isTransactionOpen()) {
            throw new SQLRuntimeException("batching can only be performed within a transaction");
        }
        int keysOption;
        if (query.returnGeneratedKeys()) {
            keysOption = Statement.RETURN_GENERATED_KEYS;
        } else {
            keysOption = Statement.NO_GENERATED_KEYS;
        }
        state.ps = state.con.prepareStatement(query.sql(), keysOption);
        Util.setParameters(state.ps, parameters, query.names());

        if (subscriber.isUnsubscribed())
            return;

        int count;
        try {
            debug("executing sql={}, parameters {}", query.sql(), parameters);
            if (state.ps instanceof PreparedStatementBatch
                    && parameters instanceof ArrayListFinal) {
                count = state.ps.executeUpdate();
                count += ((PreparedStatementBatch) state.ps).executeBatchRemaining();
            } else {
                count = state.ps.executeUpdate();
            }
            debug("executed ps={}", state.ps);
            if (query.returnGeneratedKeys()) {
                debug("getting generated keys");
                ResultSet rs = state.ps.getGeneratedKeys();
                debug("returned generated key result set {}", rs);
                state.rs = rs;
                Observable params = Observable.just(new Parameter(state));
                Observable depends = Observable.empty();
                Observable o = new QuerySelect(QuerySelect.RETURN_GENERATED_KEYS, params,
                        depends, query.context(), query.context().resultSetTransform())
                                .execute(query.returnGeneratedKeysFunction());
                Subscriber sub = createSubscriber(subscriber);
                o.unsafeSubscribe(sub);
            }
        } catch (SQLException e) {
            throw new SQLException("failed to execute sql=" + query.sql(), e);
        }
        if (!query.returnGeneratedKeys()) {
            // must close before onNext so that connection is released and is
            // available to a query that might process the onNext
            close(state);
            if (subscriber.isUnsubscribed())
                return;
            debug("onNext");
            subscriber.onNext((T) (Integer) count);
            complete(subscriber);
        }
    }

    private Subscriber createSubscriber(final Subscriber subscriber) {
        return new Subscriber(subscriber) {

            @Override
            public void onCompleted() {
                complete(subscriber);
            }

            @Override
            public void onError(Throwable e) {
                subscriber.onError(e);
            }

            @Override
            public void onNext(T t) {
                subscriber.onNext(t);
            }
        };
    }

    /**
     * Notify observer that sequence is complete.
     * 
     * @param subscriber
     * @param state
     */
    private void complete(Subscriber subscriber) {
        if (!subscriber.isUnsubscribed()) {
            debug("onCompleted");
            subscriber.onCompleted();
        } else
            debug("unsubscribed");
    }

    /**
     * Notify observer of an error.
     * 
     * @param e
     * @param subscriber
     */
    private void handleException(Throwable e, Subscriber subscriber) {
        debug("onError: ", e.getMessage());
        Exceptions.throwOrReport(e, subscriber);
    }

    /**
     * Cancels a running PreparedStatement, closing it and the current
     * Connection but only if auto commit mode.
     */
    private void close(State state) {
        // ensure close happens once only to avoid race conditions
        if (state.closed.compareAndSet(false, true)) {
            Util.closeQuietly(state.ps);
            if (isCommit() || isRollback())
                Util.closeQuietly(state.con);
            else
                Util.closeQuietlyIfAutoCommit(state.con);
        }
    }

    private static void debug(String message, Object... objects) {
        log.debug(message, objects);
    }

}