org.babyfish.jimmer.sql.runtime.ConnectionManager Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of jimmer-sql Show documentation

A revolutionary ORM framework for both java and kotlin

There is a newer version: 0.9.19

package org.babyfish.jimmer.sql.runtime;

import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Objects;
import java.util.function.Function;

@FunctionalInterface
public interface ConnectionManager {

    ConnectionManager EXTERNAL_ONLY = new ConnectionManager() {
        @Override
        public  R execute(@Nullable Connection con, Function block) {
            Objects.requireNonNull(con, "External connection cannot be null");
            return block.apply(con);
        }
    };

    static ConnectionManager singleConnectionManager(Connection connection) {
        if (connection == null) {
            return EXTERNAL_ONLY;
        }
        return new ConnectionManager() {
            @Override
            public  R execute(@Nullable Connection con, Function block) {
                return block.apply(con == null ? connection : con);
            }
        };
    }

    static ConnectionManager simpleConnectionManager(DataSource dataSource) {
        return new ConnectionManager() {

            private final ThreadLocal local = new ThreadLocal<>();

            @Override
            public  R execute(@Nullable Connection con, Function block) {
                if (con == null) {
                    return execute(block);
                }
                return block.apply(con);
            }

            @Override
            public  R execute(Function block) {
                Connection con = local.get();
                if (con != null) {
                    return block.apply(con);
                }
                try {
                    con = dataSource.getConnection();
                    try {
                        local.set(con);
                        try {
                            return block.apply(con);
                        } finally {
                            local.remove();
                        }
                    } finally {
                        con.close();
                    }
                } catch (SQLException ex) {
                    throw new ExecutionException("Cannot open connection from datasource", ex);
                }
            }
        };
    }

     R execute(@Nullable Connection con, Function block);

    default  R execute(Function block) {
        return execute(null, block);
    }
}