/*
 * 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.plugin.memory;

import com.google.common.collect.ImmutableMap;
import io.airlift.log.Logger;
import io.airlift.log.Logging;
import io.trino.Session;
import io.trino.plugin.tpch.TpchPlugin;
import io.trino.testing.DistributedQueryRunner;
import io.trino.tpch.TpchTable;

import java.util.Map;

import static io.airlift.testing.Closeables.closeAllSuppress;
import static io.trino.plugin.tpch.TpchMetadata.TINY_SCHEMA_NAME;
import static io.trino.testing.QueryAssertions.copyTpchTables;
import static io.trino.testing.TestingSession.testSessionBuilder;

public final class MemoryQueryRunner
{
    private static final String CATALOG = "memory";

    private MemoryQueryRunner() {}

    public static DistributedQueryRunner createMemoryQueryRunner(
            Map<String, String> extraProperties,
            Iterable<TpchTable<?>> tables)
            throws Exception
    {
        Session session = testSessionBuilder()
                .setCatalog(CATALOG)
                .setSchema("default")
                .build();

        DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(session)
                .setExtraProperties(extraProperties)
                .build();

        try {
            queryRunner.installPlugin(new MemoryPlugin());
            queryRunner.createCatalog(CATALOG, "memory", ImmutableMap.of());

            queryRunner.installPlugin(new TpchPlugin());
            queryRunner.createCatalog("tpch", "tpch", ImmutableMap.of());

            copyTpchTables(queryRunner, "tpch", TINY_SCHEMA_NAME, session, tables);

            return queryRunner;
        }
        catch (Exception e) {
            closeAllSuppress(e, queryRunner);
            throw e;
        }
    }

    public static void main(String[] args)
            throws Exception
    {
        Logging.initialize();
        DistributedQueryRunner queryRunner = createMemoryQueryRunner(
                ImmutableMap.of("http-server.http.port", "8080"),
                TpchTable.getTables());
        Thread.sleep(10);
        Logger log = Logger.get(MemoryQueryRunner.class);
        log.info("======== SERVER STARTED ========");
        log.info("\n====\n%s\n====", queryRunner.getCoordinator().getBaseUrl());
    }
}

@@ -588,16 +588,64 @@
 emory;%0A%0A
+import com.google.common.collect.ImmutableList;%0A
 import c
@@ -890,16 +890,39 @@
 Table;%0A%0A
+import java.util.List;%0A
 import j
@@ -1192,16 +1192,64 @@
 Builder;
+%0Aimport static java.util.Objects.requireNonNull;
 %0A%0Apublic
@@ -1575,82 +1575,76 @@
     
-Session session = testSessionBuilder()%0A                .setCatalog(CATALOG
+return builder()%0A                .setExtraProperties(extraProperties
 )%0A  
@@ -1665,24 +1665,28 @@
 .set
-Schema(%22default%22
+InitialTables(tables
 )%0A  
@@ -1713,93 +1713,337 @@
 ();%0A
-%0A
     
+%7D%0A%0A
     
-DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(s
+public static Builder builder()%0A    %7B%0A        return new Builder();%0A    %7D%0A%0A    public static class Builder%0A            extends DistributedQueryRunner.Builder%3CBuilder%3E%0A    %7B%0A        private List%3CTpchTable%3C?%3E%3E initialTables = ImmutableList.of();%0A%0A        protected Builder()%0A        %7B%0A            super(createS
 ession
-)
+());
 %0A   
@@ -2043,32 +2043,35 @@
 ));%0A        
+%7D%0A%0A
         
 .setExtraPro
@@ -2062,61 +2062,398 @@
     
-.setExtraProperties(extraProperties)%0A                
+public Builder setInitialTables(Iterable%3CTpchTable%3C?%3E%3E initialTables)%0A        %7B%0A            this.initialTables = ImmutableList.copyOf(requireNonNull(initialTables, %22initialTables is null%22));%0A            return self();%0A        %7D%0A%0A        @Override%0A        public DistributedQueryRunner build()%0A                throws Exception%0A        %7B%0A            DistributedQueryRunner queryRunner = super
 .bui
@@ -2467,16 +2467,20 @@
         
+    
 try %7B%0A  
@@ -2481,32 +2481,36 @@
 y %7B%0A            
+  
+  
 queryRunner.inst
@@ -2532,32 +2532,36 @@
 emoryPlugin());%0A
+    
             quer
@@ -2626,32 +2626,36 @@
 );%0A%0A            
+  
+  
 queryRunner.inst
@@ -2675,32 +2675,36 @@
  TpchPlugin());%0A
+    
             quer
@@ -2774,16 +2774,20 @@
         
+  
+  
 copyTpch
@@ -2836,18 +2836,33 @@
 ME, 
-session, t
+createSession(), initialT
 able
@@ -2862,24 +2862,28 @@
 alTables);%0A%0A
+    
             
@@ -2906,26 +2906,34 @@
 er;%0A        
+  
-%7D%0A
+  %7D%0A    
         catc
@@ -2962,16 +2962,20 @@
         
+  
+  
 closeAll
@@ -3008,24 +3008,28 @@
             
+    
 throw e;%0A   
@@ -3025,16 +3025,251 @@
 hrow e;%0A
+            %7D%0A        %7D%0A%0A        private static Session createSession()%0A        %7B%0A            return testSessionBuilder()%0A                    .setCatalog(CATALOG)%0A                    .setSchema(%22default%22)%0A                    .build();%0A
         

Add builder to MemoryQueryRunner