connection supplies properties approach

extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java

@@ -160,4 +160,14 @@ public class MSQDruidMeta extends DruidMeta
 
   }
 
+  @Override
+  public Map<String, String> getEngineProperties()
+  {
+    return ImmutableMap.<String, String>builder()
+        .putAll(super.getEngineProperties())
+        .put("isMSQ", "true")
+        .put("isNative", "false")
+        .build();
+  }
+
 }

extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.exec.MSQQuidemTest/msq2.iq

@@ -1,5 +1,3 @@
-!set isNative false
-!set isMSQ true
 !set plannerStrategy DECOUPLED
 !if (isMSQ) {
 !use druidtest://?componentSupplier=DrillWindowQueryMSQComponentSupplier

sql/src/main/java/org/apache/druid/sql/avatica/DruidMeta.java

@@ -933,4 +933,9 @@ public class DruidMeta extends MetaImpl
   {
     return Calcites.escapeStringLiteral(toEscape) + " ESCAPE '\\'";
   }
+
+  public Map<String, String> getEngineProperties()
+  {
+    return ImmutableMap.of("isNative", "true");
+  }
 }

sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java

@@ -120,9 +120,16 @@ public class DruidAvaticaTestDriver implements Driver
 
     @Provides
     @LazySingleton
-    public DruidConnectionExtras getConnectionExtras(ObjectMapper objectMapper, DruidHookDispatcher druidHookDispatcher)
+    public DruidConnectionExtras getConnectionExtras(
+        ObjectMapper objectMapper,
+        DruidHookDispatcher druidHookDispatcher,
+        DruidMeta meta)
     {
-      return new DruidConnectionExtras.DruidConnectionExtrasImpl(objectMapper, druidHookDispatcher);
+      return new DruidConnectionExtras.DruidConnectionExtrasImpl(
+          objectMapper,
+          druidHookDispatcher,
+          meta.getEngineProperties()
+      );
     }
 
     @Provides
@@ -144,7 +151,6 @@ public class DruidAvaticaTestDriver implements Driver
     {
       closer.close();
     }
-
   }
 
   static class AvaticaJettyServer implements Closeable

sql/src/test/java/org/apache/druid/quidem/DruidConnectionExtras.java

@@ -23,6 +23,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.druid.sql.hook.DruidHookDispatcher;
 
 import java.sql.Connection;
+import java.util.Map;
 
 public interface DruidConnectionExtras
 {
@@ -30,15 +31,21 @@ public interface DruidConnectionExtras
 
   DruidHookDispatcher getDruidHookDispatcher();
 
+  Map<String, String> getEngineProperties();
+
   class DruidConnectionExtrasImpl implements DruidConnectionExtras
   {
     private final ObjectMapper objectMapper;
     private final DruidHookDispatcher druidHookDispatcher;
+    private final Map<String, String> engineProperties;
 
-    public DruidConnectionExtrasImpl(ObjectMapper objectMapper, DruidHookDispatcher druidHookDispatcher)
+    public DruidConnectionExtrasImpl(ObjectMapper objectMapper,
+        DruidHookDispatcher druidHookDispatcher,
+        Map<String, String> engineProperties)
     {
       this.objectMapper = objectMapper;
       this.druidHookDispatcher = druidHookDispatcher;
+      this.engineProperties = engineProperties;
     }
 
     @Override
@@ -52,6 +59,12 @@ public interface DruidConnectionExtras
     {
       return druidHookDispatcher;
     }
+
+    @Override
+    public Map<String, String> getEngineProperties()
+    {
+      return engineProperties;
+    }
   }
 
   static DruidConnectionExtras unwrapOrThrow(Connection connection)
@@ -61,4 +74,5 @@ public interface DruidConnectionExtras
     }
     throw new UnsupportedOperationException("Expected DruidConnectionExtras to be implemented by connection!");
   }
+
 }

sql/src/test/java/org/apache/druid/quidem/DruidQuidemConnectionFactory.java

@@ -21,13 +21,16 @@ package org.apache.druid.quidem;
 
 import net.hydromatic.quidem.Quidem.ConnectionFactory;
 import net.hydromatic.quidem.Quidem.PropertyHandler;
+
 import java.sql.Connection;
 import java.sql.DriverManager;
+import java.util.Map;
 import java.util.Properties;
 
 public class DruidQuidemConnectionFactory implements ConnectionFactory, PropertyHandler
 {
   private Properties props = new Properties();
+  private Map<String, String> engineProperties;
 
   public DruidQuidemConnectionFactory()
   {
@@ -39,14 +42,33 @@ public class DruidQuidemConnectionFactory implements ConnectionFactory, Property
   public Connection connect(String name, boolean reference) throws Exception
   {
     if (name.startsWith("druidtest://")) {
-      return DriverManager.getConnection(name, props);
+      Connection connection = DriverManager.getConnection(name, props);
+      engineProperties = unwrapEngineProperties(connection);
+      return connection;
     }
     throw new RuntimeException("unknown connection '" + name + "'");
   }
 
+  private Map<String, String> unwrapEngineProperties(Connection connection)
+  {
+    if(connection instanceof DruidConnectionExtras) {
+      DruidConnectionExtras extras = ((DruidConnectionExtras) connection);
+      return extras.getEngineProperties();
+    }
+    return null;
+  }
+
   @Override
   public void onSet(String key, Object value)
   {
     props.setProperty(key, value.toString());
   }
+
+  public Object getEnv(String env)
+  {
+    if (engineProperties == null) {
+      return null;
+    }
+    return engineProperties.get(env);
+  }
 }