buckelieg.fn.db.ScriptQuery Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of db-fn Show documentation
Functional style programming with plain JDBC
There is a newer version: 0.3.10
/*
 * Copyright 2016- Anatoly Kutyakov
 *
 * Licensed 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 buckelieg.fn.db;

import javax.annotation.Nonnull;
import javax.annotation.ParametersAreNonnullByDefault;
import javax.annotation.concurrent.NotThreadSafe;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.SQLWarning;
import java.sql.Statement;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;

import static buckelieg.fn.db.Utils.newSQLRuntimeException;
import static buckelieg.fn.db.Utils.parseScript;

@SuppressWarnings("unchecked")
@NotThreadSafe
@ParametersAreNonnullByDefault
final class ScriptQuery implements Script {

    private static final Consumer NOOP = e -> {
        // do nothing
    };
    private final String script;
    private final TrySupplier connectionSupplier;
    private ExecutorService conveyor;
    private int timeout;
    private boolean escaped = true;
    private boolean skipErrors = true;
    private boolean skipWarnings = true;
    private Consumer errorHandler = NOOP;

    /**
     * Creates script executor query
     *
     * @param connectionSupplier db connection provider
     * @param script             an arbitrary SQL script to execute
     * @throws IllegalArgumentException in case of illegal comment lines encountered
     */
    ScriptQuery(TrySupplier connectionSupplier, String script) {
        this.connectionSupplier = connectionSupplier;
        try {
            this.script = Utils.cutComments(Objects.requireNonNull(script, "Script string must be provided"));
        } catch (SQLException e) {
            throw new IllegalArgumentException(e);
        }
    }

    /**
     * Executes script. All comments are cut out.
     * Therefore all RDBMS-scpecific hints are ignored (like Oracle APPEND) etc.
     * Note that script commit must be done inside script itself.
     *
     * @return a time, taken by this script to complete in milliseconds
     * @throws SQLRuntimeException in case of any errors including {@link SQLWarning} (if corresponding option is set) OR (if timeout is set) - in case of execution run out of time.
     */
    @Nonnull
    @Override
    public Long execute() {
        switch (timeout) {
            case 0: {
                return doExecute();
            }
            default: {
                try {
                    conveyor = Executors.newSingleThreadExecutor(); // TODO implement executor that uses current thread
                    return conveyor.submit(this::doExecute).get(timeout, TimeUnit.SECONDS);
                } catch (Exception e) {
                    throw newSQLRuntimeException(e);
                }
            }
        }
    }

    private long doExecute() {
        long start = System.currentTimeMillis();
        try {
            Connection conn = connectionSupplier.get();
            for (String query : parseScript(script)) {
                try (Statement statement = conn.createStatement()) {
                    statement.setEscapeProcessing(escaped);
                    if (skipErrors) {
                        try {
                            statement.execute(query);
                        } catch (SQLException e) {
                            errorHandler.accept(e);
                        }
                    } else {
                        statement.execute(query);
                    }
                    Optional warning = Optional.ofNullable(statement.getWarnings());
                    if (!skipWarnings && warning.isPresent()) {
                        throw warning.get();
                    } else {
                        warning.ifPresent(errorHandler::accept);
                    }
                }
            }
        } catch (Exception e) {
            throw newSQLRuntimeException(e);
        } finally {
            close();
        }
        return System.currentTimeMillis() - start;
    }

    @Nonnull
    @Override
    public Script escaped(boolean escapeProcessing) {
        this.escaped = escapeProcessing;
        return this;
    }

    @Nonnull
    @Override
    public Script print(Consumer printer) {
        Objects.requireNonNull(printer, "Printer must be provided").accept(script);
        return this;
    }

    @Nonnull
    @Override
    public Script skipErrors(boolean skipErrors) {
        this.skipErrors = skipErrors;
        return this;
    }

    @Nonnull
    @Override
    public Script skipWarnings(boolean skipWarnings) {
        this.skipWarnings = skipWarnings;
        return this;
    }

    @Nonnull
    @Override
    public Script timeout(int timeout) {
        this.timeout = timeout < 0 ? 0 : timeout;
        return this;
    }

    @Nonnull
    @Override
    public Script errorHandler(Consumer handler) {
        this.errorHandler = Objects.requireNonNull(handler, "Error handler must be provided");
        return this;
    }

    @Override
    public String toString() {
        return script;
    }

    @Override
    public void close() {
        if (conveyor != null) conveyor.shutdownNow();
    }
}