From 57765a499b4bc328a543c0988e0e2ce64b7834d5 Mon Sep 17 00:00:00 2001
From: Jonathan Wei <jon-wei@users.noreply.github.com>
Date: Thu, 6 Feb 2020 18:35:26 -0800
Subject: [PATCH] Allow overriding default JoinableFactory in
 SpecificSegmentsQuerySegmentWalker (#9330)

---
 .../druid/sql/calcite/util/CalciteTests.java  |  3 +-
 .../SpecificSegmentsQuerySegmentWalker.java   | 35 ++++++++++++++-----
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java
index 51a505de686..0261b6a5884 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/CalciteTests.java
@@ -741,7 +741,8 @@ public class CalciteTests
 
     return new SpecificSegmentsQuerySegmentWalker(
         conglomerate,
-        INJECTOR.getInstance(LookupExtractorFactoryContainerProvider.class)
+        INJECTOR.getInstance(LookupExtractorFactoryContainerProvider.class),
+        null
     ).add(
         DataSegment.builder()
                    .dataSource(DATASOURCE1)
diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/util/SpecificSegmentsQuerySegmentWalker.java b/sql/src/test/java/org/apache/druid/sql/calcite/util/SpecificSegmentsQuerySegmentWalker.java
index 7fe3f782089..36db41c47bf 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/util/SpecificSegmentsQuerySegmentWalker.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/util/SpecificSegmentsQuerySegmentWalker.java
@@ -71,6 +71,7 @@ import org.apache.druid.timeline.partition.PartitionChunk;
 import org.apache.druid.timeline.partition.PartitionHolder;
 import org.joda.time.Interval;
 
+import javax.annotation.Nullable;
 import java.io.Closeable;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -102,19 +103,23 @@ public class SpecificSegmentsQuerySegmentWalker implements QuerySegmentWalker, C
 
   /**
    * Create an instance using the provided query runner factory conglomerate and lookup provider.
+   * If a JoinableFactory is provided, it will be used instead of the default.
    */
   public SpecificSegmentsQuerySegmentWalker(
       final QueryRunnerFactoryConglomerate conglomerate,
-      final LookupExtractorFactoryContainerProvider lookupProvider
+      final LookupExtractorFactoryContainerProvider lookupProvider,
+      @Nullable final JoinableFactory joinableFactory
   )
   {
     this.conglomerate = conglomerate;
-    this.joinableFactory = MapJoinableFactoryTest.fromMap(
-        ImmutableMap.<Class<? extends DataSource>, JoinableFactory>builder()
-            .put(InlineDataSource.class, new InlineJoinableFactory())
-            .put(LookupDataSource.class, new LookupJoinableFactory(lookupProvider))
-            .build()
-    );
+    this.joinableFactory = joinableFactory == null ?
+                           MapJoinableFactoryTest.fromMap(
+                               ImmutableMap.<Class<? extends DataSource>, JoinableFactory>builder()
+                                   .put(InlineDataSource.class, new InlineJoinableFactory())
+                                   .put(LookupDataSource.class, new LookupJoinableFactory(lookupProvider))
+                                   .build()
+                           ) : joinableFactory;
+
     this.walker = new ClientQuerySegmentWalker(
         new NoopServiceEmitter(),
         new DataServerLikeWalker(),
@@ -160,9 +165,20 @@ public class SpecificSegmentsQuerySegmentWalker implements QuerySegmentWalker, C
   }
 
   /**
-   * Create an instance without any lookups.
+   * Create an instance without any lookups, using the default JoinableFactory
    */
   public SpecificSegmentsQuerySegmentWalker(final QueryRunnerFactoryConglomerate conglomerate)
+  {
+    this(conglomerate, null);
+  }
+
+  /**
+   * Create an instance without any lookups, optionally allowing the default JoinableFactory to be overridden
+   */
+  public SpecificSegmentsQuerySegmentWalker(
+      final QueryRunnerFactoryConglomerate conglomerate,
+      @Nullable JoinableFactory joinableFactory
+  )
   {
     this(
         conglomerate,
@@ -179,7 +195,8 @@ public class SpecificSegmentsQuerySegmentWalker implements QuerySegmentWalker, C
           {
             return Optional.empty();
           }
-        }
+        },
+        joinableFactory
     );
   }