/*
 * 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.trino.tests.utils;

import io.trino.tempto.query.QueryExecutionException;
import io.trino.tempto.query.QueryExecutor;
import io.trino.tempto.query.QueryResult;
import net.jodah.failsafe.Failsafe;

import java.sql.Connection;

import static io.trino.tempto.context.ThreadLocalTestContextHolder.testContext;
import static io.trino.tests.hive.HiveProductTest.ERROR_COMMITTING_WRITE_TO_HIVE_RETRY_POLICY;

public final class QueryExecutors
{
    public static QueryExecutor onPresto()
    {
        return testContext().getDependency(QueryExecutor.class, "presto");
    }

    public static QueryExecutor connectToPresto(String prestoConfig)
    {
        return testContext().getDependency(QueryExecutor.class, prestoConfig);
    }

    public static QueryExecutor onHive()
    {
        return new QueryExecutor()
        {
            private final QueryExecutor delegate = testContext().getDependency(QueryExecutor.class, "hive");

            @Override
            public QueryResult executeQuery(String sql, QueryParam... params)
                    throws QueryExecutionException
            {
                return Failsafe.with(ERROR_COMMITTING_WRITE_TO_HIVE_RETRY_POLICY)
                        .get(() -> delegate.executeQuery(sql, params));
            }

            @Override
            public Connection getConnection()
            {
                return delegate.getConnection();
            }

            @Override
            public void close()
            {
                delegate.close();
            }
        };
    }

    public static QueryExecutor onSqlServer()
    {
        return testContext().getDependency(QueryExecutor.class, "sqlserver");
    }

    public static QueryExecutor onMySql()
    {
        return testContext().getDependency(QueryExecutor.class, "mysql");
    }

    public static QueryExecutor onSpark()
    {
        return testContext().getDependency(QueryExecutor.class, "spark");
    }

    private QueryExecutors() {}
}

@@ -1072,57 +1072,24 @@
 urn 
-testContext().getDependency(QueryExecutor.class, 
+connectToPresto(
 %22pre
@@ -1181,141 +1181,8 @@
   %7B%0A
-        return testContext().getDependency(QueryExecutor.class, prestoConfig);%0A    %7D%0A%0A    public static QueryExecutor onHive()%0A    %7B%0A
     
@@ -1326,14 +1326,20 @@
 ss, 
-%22hive%22
+prestoConfig
 );%0A%0A
@@ -1944,32 +1944,159 @@
       %7D;%0A    %7D%0A%0A
+    public static QueryExecutor onHive()%0A    %7B%0A        return testContext().getDependency(QueryExecutor.class, %22hive%22);%0A    %7D%0A%0A
     public stati

Move 'Error committing write to Hive' retry to Trino query executor