TestGetGeneratedKeysPostgres.java

/*
 * 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 org.jdbi.v3.sqlobject;

import static org.assertj.core.api.Assertions.assertThat;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.time.OffsetDateTime;
import java.time.ZoneOffset;
import java.util.Arrays;
import java.util.List;

import org.jdbi.v3.core.mapper.RowMapper;
import org.jdbi.v3.core.rule.PgDatabaseRule;
import org.jdbi.v3.core.statement.StatementContext;
import org.jdbi.v3.sqlobject.config.RegisterRowMapper;
import org.jdbi.v3.sqlobject.customizer.Bind;
import org.jdbi.v3.sqlobject.statement.GetGeneratedKeys;
import org.jdbi.v3.sqlobject.statement.SqlBatch;
import org.jdbi.v3.sqlobject.statement.SqlQuery;
import org.jdbi.v3.sqlobject.statement.SqlUpdate;
import org.jdbi.v3.sqlobject.statement.UseRowMapper;
import org.junit.Rule;
import org.junit.Test;

public class TestGetGeneratedKeysPostgres
{
    @Rule
    public PgDatabaseRule dbRule = new PgDatabaseRule().withPlugin(new SqlObjectPlugin());

    @Test
    public void testFoo() throws Exception {
        dbRule.getJdbi().useExtension(DAO.class, dao -> {
            dao.createSequence();
            dao.createTable();

            long brian_id = dao.insert("Brian");
            long keith_id = dao.insert("Keith");

            assertThat(dao.findNameById(brian_id)).isEqualTo("Brian");
            assertThat(dao.findNameById(keith_id)).isEqualTo("Keith");
        });
    }

    @Test
    public void testBatch() throws Exception {
        dbRule.getJdbi().useExtension(DAO.class, dao -> {
            dao.createSequence();
            dao.createTable();

            int[] ids = dao.insert(Arrays.asList("Burt", "Macklin"));

            assertThat(dao.findNameById(ids[0])).isEqualTo("Burt");
            assertThat(dao.findNameById(ids[1])).isEqualTo("Macklin");
        });
    }

    @Test
    public void testUseRowMapperUpdate() throws Exception {
        dbRule.getJdbi().useExtension(UseRowMapperDao.class, dao -> {
            dao.createTable();

            IdCreateTime result = dao.insert("foo");

            assertThat(result.id).isEqualTo(1);
            assertThat(result.createdOn).isNotNull();
        });
    }

    @Test
    public void testUseRowMapperBatch() throws Exception {
        dbRule.getJdbi().useExtension(UseRowMapperDao.class, dao -> {
            dao.createTable();

            List<IdCreateTime> results = dao.insertBatch("foo", "bar");

            assertThat(results).extracting(ic -> ic.id).containsExactly(1L, 2L);
            assertThat(results).extracting(ic -> ic.createdOn).hasSize(2).doesNotContainNull();
        });
    }

    @Test
    public void testRegisterRowMapperUpdate() throws Exception {
        dbRule.getJdbi().useExtension(RegisterRowMapperDao.class, dao -> {
            dao.createTable();

            IdCreateTime result = dao.insert("foo");

            assertThat(result.id).isEqualTo(1);
            assertThat(result.createdOn).isNotNull();
        });
    }

    @Test
    public void testRegisterRowMapperBatch() throws Exception {
        dbRule.getJdbi().useExtension(RegisterRowMapperDao.class, dao -> {
            dao.createTable();

            List<IdCreateTime> results = dao.insertBatch("foo", "bar");

            assertThat(results).extracting(ic -> ic.id).containsExactly(1L, 2L);
            assertThat(results).extracting(ic -> ic.createdOn).hasSize(2).doesNotContainNull();
        });
    }

    public interface DAO {
        @SqlUpdate("create sequence id_sequence INCREMENT 1 START WITH 100")
        void createSequence();

        @SqlUpdate("create table if not exists something (name text, id int DEFAULT nextval('id_sequence'), CONSTRAINT something_id PRIMARY KEY ( id ))")
        void createTable();

        @SqlUpdate("insert into something (name, id) values (:name, nextval('id_sequence'))")
        @GetGeneratedKeys("id")
        long insert(@Bind("name") String name);

        @SqlBatch("insert into something (name) values (:name)")
        @GetGeneratedKeys("id")
        int[] insert(@Bind("name") List<String> names);

        @SqlQuery("select name from something where id = :id")
        String findNameById(@Bind long id);
    }

    public interface UseRowMapperDao {
        @SqlUpdate("create table something (id serial, name varchar(50), created_on timestamp default now())")
        void createTable();

        @SqlUpdate("insert into something(name) values (:name)")
        @GetGeneratedKeys({"id", "created_on"})
        @UseRowMapper(IdCreateTimeMapper.class)
        IdCreateTime insert(String name);

        @SqlBatch("insert into something(name) values (:name)")
        @GetGeneratedKeys({"id", "created_on"})
        @UseRowMapper(IdCreateTimeMapper.class)
        List<IdCreateTime> insertBatch(String... name);
    }

    public interface RegisterRowMapperDao {
        @SqlUpdate("create table something (id serial, name varchar(50), created_on timestamp default now())")
        void createTable();

        @SqlUpdate("insert into something(name) values (:name)")
        @GetGeneratedKeys({"id", "created_on"})
        @RegisterRowMapper(IdCreateTimeMapper.class)
        IdCreateTime insert(String name);

        @SqlBatch("insert into something(name) values (:name)")
        @GetGeneratedKeys({"id", "created_on"})
        @RegisterRowMapper(IdCreateTimeMapper.class)
        List<IdCreateTime> insertBatch(String... name);
    }

    public static class IdCreateTimeMapper implements RowMapper<IdCreateTime> {
        @Override
        public IdCreateTime map(ResultSet rs, StatementContext ctx) throws SQLException {
            return new IdCreateTime(
                    rs.getLong("id"),
                    rs.getTimestamp("created_on").toInstant().atOffset(ZoneOffset.UTC));
        }
    }

    public static class IdCreateTime {
        final long id;
        final OffsetDateTime createdOn;

        public IdCreateTime(long id, OffsetDateTime createdOn) {
            this.id = id;
            this.createdOn = createdOn;
        }
    }
}