diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java index 845d8adb2b3..c6c17a67f5d 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/sql/avatica/MSQDruidMeta.java +++ b/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 getEngineProperties() + { + return ImmutableMap.builder() + .putAll(super.getEngineProperties()) + .put("isMSQ", "true") + .put("isNative", "false") + .build(); + } + } diff --git a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.exec.MSQQuidemTest/msq2.iq b/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.exec.MSQQuidemTest/msq2.iq index f36f8e04460..f83a3747719 100644 --- a/extensions-core/multi-stage-query/src/test/quidem/org.apache.druid.msq.exec.MSQQuidemTest/msq2.iq +++ b/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 diff --git a/sql/src/main/java/org/apache/druid/sql/avatica/DruidMeta.java b/sql/src/main/java/org/apache/druid/sql/avatica/DruidMeta.java index e271a6f60ab..3143fba5d84 100644 --- a/sql/src/main/java/org/apache/druid/sql/avatica/DruidMeta.java +++ b/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 getEngineProperties() + { + return ImmutableMap.of("isNative", "true"); + } } diff --git a/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java b/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java index 30eca2782e7..c25f9275152 100644 --- a/sql/src/test/java/org/apache/druid/quidem/DruidAvaticaTestDriver.java +++ b/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 diff --git a/sql/src/test/java/org/apache/druid/quidem/DruidConnectionExtras.java b/sql/src/test/java/org/apache/druid/quidem/DruidConnectionExtras.java index 176411a6bb7..c7538775c06 100644 --- a/sql/src/test/java/org/apache/druid/quidem/DruidConnectionExtras.java +++ b/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 getEngineProperties(); + class DruidConnectionExtrasImpl implements DruidConnectionExtras { private final ObjectMapper objectMapper; private final DruidHookDispatcher druidHookDispatcher; + private final Map engineProperties; - public DruidConnectionExtrasImpl(ObjectMapper objectMapper, DruidHookDispatcher druidHookDispatcher) + public DruidConnectionExtrasImpl(ObjectMapper objectMapper, + DruidHookDispatcher druidHookDispatcher, + Map engineProperties) { this.objectMapper = objectMapper; this.druidHookDispatcher = druidHookDispatcher; + this.engineProperties = engineProperties; } @Override @@ -52,6 +59,12 @@ public interface DruidConnectionExtras { return druidHookDispatcher; } + + @Override + public Map 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!"); } + } diff --git a/sql/src/test/java/org/apache/druid/quidem/DruidQuidemConnectionFactory.java b/sql/src/test/java/org/apache/druid/quidem/DruidQuidemConnectionFactory.java index b9a7963f4d8..1c38702537c 100644 --- a/sql/src/test/java/org/apache/druid/quidem/DruidQuidemConnectionFactory.java +++ b/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 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 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); + } }