/*
 * Copyright 2017 requery.io
 *
 * 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 io.requery.sql;

import io.requery.meta.Type;
import io.requery.query.Expression;
import io.requery.query.MutableTuple;
import io.requery.query.NamedExpression;
import io.requery.query.Result;
import io.requery.query.Tuple;
import io.requery.query.element.InsertType;
import io.requery.query.element.QueryElement;
import io.requery.query.element.QueryOperation;
import io.requery.sql.gen.DefaultOutput;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Set;

/**
 * Executes insert operation returning the set of generated keys as a result, if no generated keys
 * then a single value representing the affected row count.
 *
 * @author Nikhil Purushe
 */
class InsertReturningOperation extends PreparedQueryOperation implements
    QueryOperation<Result<Tuple>> {

    private final Set<? extends Expression<?>> selection;

    InsertReturningOperation(RuntimeConfiguration configuration,
                             final Set<? extends Expression<?>> selection) {
        super(configuration, new GeneratedResultReader() {
            @Override
            public void read(int index, ResultSet results) throws SQLException {
                // skipped, not used
            }

            @Override
            public String[] generatedColumns() {
                String[] keys = new String[selection.size()];
                int index = 0;
                for (Expression<?> expression : selection) {
                    keys[index] = expression.getName();
                    index++;
                }
                return keys;
            }
        });
        this.selection = selection;
    }

    @Override
    public Result<Tuple> evaluate(final QueryElement<Result<Tuple>> query) {
        DefaultOutput generator = new DefaultOutput(configuration, query);
        String sql = generator.toSql();
        BoundParameters parameters = generator.parameters();
        int count;
        PreparedStatement statement = null;
        TransactionProvider transactionProvider = configuration.getTransactionProvider();
        Set<Type<?>> types = query.entityTypes();
        try (TransactionScope scope = new TransactionScope(transactionProvider, types)) {
            Connection connection = configuration.getConnection();
            StatementListener listener = configuration.getStatementListener();
            if (query.insertType() == InsertType.SELECT) {
                statement = connection.prepareStatement(sql, Statement.NO_GENERATED_KEYS);
            } else {
                statement = prepare(sql, connection);
            }
            mapParameters(statement, parameters);
            listener.beforeExecuteUpdate(statement, sql, parameters);
            count = statement.executeUpdate();
            listener.afterExecuteUpdate(statement, count);
            scope.commit();
            if (selection == null || selection.isEmpty() || query.insertType() == InsertType.SELECT) {
                connection.close();
                MutableTuple tuple = new MutableTuple(1);
                tuple.set(0, NamedExpression.ofInteger("count"), count);
                return new SingleResult<Tuple>(tuple);
            } else {
                ResultSet results = statement.getGeneratedKeys();
                return new GeneratedKeyResult(configuration, selection, connection, results, count);
            }
        } catch (Exception e) {
            throw StatementExecutionException.closing(statement, e, sql);
        }
    }
}

@@ -613,16 +613,55 @@
 y.sql;%0A%0A
+import io.requery.TransactionListener;%0A
 import i
@@ -2690,98 +2690,8 @@
 ll;%0A
-        TransactionProvider transactionProvider = configuration.getTransactionProvider();%0A
     
@@ -2748,13 +2748,8 @@
     
-try (
 Tran
@@ -2759,71 +2759,150 @@
 tion
-Scope scope = new TransactionScope(transactionProvider, typ
+Listener transactionListener = new CompositeTransactionListener(%0A                configuration.getTransactionListenerFactori
 es
+(
 ))
+;%0A        try
  %7B%0A 
@@ -3327,32 +3327,33 @@
 t, parameters);%0A
+%0A
             list
@@ -3406,16 +3406,70 @@
 eters);%0A
+            transactionListener.beforeCommit(types);%0A%0A
         
@@ -3503,24 +3503,25 @@
 teUpdate();%0A
+%0A
             
@@ -3583,24 +3583,49 @@
     
-scope.c
+transactionListener.afterC
 ommit(
+types
 );%0A
+%0A
     

Fix connection closed exception