Alter node client passing: OpenSearchPluginModule -> QueryService -> CalcitePlanContext

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java

@@ -91,4 +91,8 @@ public static CalcitePlanContext create(
       FrameworkConfig config, QueryType queryType, NodeClient client) {
     return new CalcitePlanContext(config, queryType, client);
   }
+
+  public static CalcitePlanContext create(FrameworkConfig config, QueryType queryType) {
+    return new CalcitePlanContext(config, queryType, null);
+  }
 }

core/src/main/java/org/opensearch/sql/datasource/DataSourceService.java

@@ -9,7 +9,6 @@
 import java.util.Set;
 import org.opensearch.sql.datasource.model.DataSource;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
-import org.opensearch.transport.client.node.NodeClient;
 
 /** DataSource Service manage {@link DataSource}. */
 public interface DataSourceService {
@@ -88,13 +87,4 @@ public interface DataSourceService {
    */
   DataSourceMetadata verifyDataSourceAccessAndGetRawMetadata(
       String dataSourceName, RequestContext context);
-
-  /**
-   * Returns a NodeClient instance for executing requests against the OpenSearch cluster.
-   *
-   * @return NodeClient instance or throw an exception if not implemented.
-   */
-  default NodeClient getNodeClient() {
-    throw new UnsupportedOperationException("Not implemented");
-  }
 }

core/src/main/java/org/opensearch/sql/executor/QueryService.java

@@ -44,6 +44,7 @@
 import org.opensearch.sql.planner.Planner;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.physical.PhysicalPlan;
+import org.opensearch.transport.client.node.NodeClient;
 
 /** The low level interface of core engine. */
 @RequiredArgsConstructor
@@ -59,6 +60,7 @@ public class QueryService {
 
   private DataSourceService dataSourceService;
   private Settings settings;
+  private NodeClient nodeClient;
 
   /** Execute the {@link UnresolvedPlan}, using {@link ResponseListener} to get response.<br> */
   public void execute(
@@ -94,8 +96,7 @@ public void executeWithCalcite(
           (PrivilegedAction<Void>)
               () -> {
                 CalcitePlanContext context =
-                    CalcitePlanContext.create(
-                        buildFrameworkConfig(), queryType, dataSourceService.getNodeClient());
+                    CalcitePlanContext.create(buildFrameworkConfig(), queryType, nodeClient);
                 RelNode relNode = analyze(plan, context);
                 RelNode optimized = optimize(relNode);
                 RelNode calcitePlan = convertToCalcitePlan(optimized);
@@ -127,8 +128,7 @@ public void explainWithCalcite(
           (PrivilegedAction<Void>)
               () -> {
                 CalcitePlanContext context =
-                    CalcitePlanContext.create(
-                        buildFrameworkConfig(), queryType, dataSourceService.getNodeClient());
+                    CalcitePlanContext.create(buildFrameworkConfig(), queryType, nodeClient);
                 RelNode relNode = analyze(plan, context);
                 RelNode optimized = optimize(relNode);
                 RelNode calcitePlan = convertToCalcitePlan(optimized);

datasources/src/main/java/org/opensearch/sql/datasources/service/DataSourceServiceImpl.java

@@ -43,31 +43,14 @@ public class DataSourceServiceImpl implements DataSourceService {
 
   private final DataSourceUserAuthorizationHelper dataSourceUserAuthorizationHelper;
 
-  /** Certain functions (e.g., GEOIP) depend on NodeClient for performing RPC calls. **/
-  @Getter
-  private final NodeClient nodeClient;
-
   /** Construct from the set of {@link DataSourceFactory} at bootstrap time. */
-  public DataSourceServiceImpl(
-          Set<DataSourceFactory> dataSourceFactories,
-          DataSourceMetadataStorage dataSourceMetadataStorage,
-          DataSourceUserAuthorizationHelper dataSourceUserAuthorizationHelper) {
-    this(
-        dataSourceFactories,
-        dataSourceMetadataStorage,
-        dataSourceUserAuthorizationHelper,
-        null);
-  }
-
   public DataSourceServiceImpl(
       Set<DataSourceFactory> dataSourceFactories,
       DataSourceMetadataStorage dataSourceMetadataStorage,
-      DataSourceUserAuthorizationHelper dataSourceUserAuthorizationHelper,
-      NodeClient nodeClient) {
+      DataSourceUserAuthorizationHelper dataSourceUserAuthorizationHelper) {
     this.dataSourceMetadataStorage = dataSourceMetadataStorage;
     this.dataSourceUserAuthorizationHelper = dataSourceUserAuthorizationHelper;
     this.dataSourceLoaderCache = new DataSourceLoaderCacheImpl(dataSourceFactories);
-    this.nodeClient = nodeClient;
   }
 
   @Override

integ-test/src/test/java/org/opensearch/sql/calcite/standalone/CalcitePPLIntegTestCase.java

@@ -87,9 +87,7 @@ public void init() throws IOException {
                 .add(new OpenSearchDataSourceFactory(client, getSettings()))
                 .build(),
             getDataSourceMetadataStorage(),
-            getDataSourceUserRoleHelper(),
-            // Node client is not accessible and used in standalone tests, so we pass null
-            null);
+            getDataSourceUserRoleHelper());
     dataSourceService.createDataSource(defaultOpenSearchDataSourceMetadata());
 
     ModulesBuilder modules = new ModulesBuilder();
@@ -372,8 +370,9 @@ public QueryPlanFactory queryPlanFactory(ExecutionEngine executionEngine) {
           new Analyzer(
               new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository);
       Planner planner = new Planner(LogicalPlanOptimizer.create());
+      // NodeClient is not used in integration test, so we pass null
       QueryService queryService =
-          new QueryService(analyzer, executionEngine, planner, dataSourceService, settings);
+          new QueryService(analyzer, executionEngine, planner, dataSourceService, settings, null);
       return new QueryPlanFactory(queryService);
     }
   }

integ-test/src/test/java/org/opensearch/sql/ppl/StandaloneIT.java

@@ -84,9 +84,7 @@ public void init() throws Exception {
                 .add(new OpenSearchDataSourceFactory(client, defaultSettings()))
                 .build(),
             getDataSourceMetadataStorage(),
-            getDataSourceUserRoleHelper(),
-            // Node client is not accessible and used in standalone tests, so we pass null
-            null);
+            getDataSourceUserRoleHelper());
     dataSourceService.createDataSource(defaultOpenSearchDataSourceMetadata());
 
     ModulesBuilder modules = new ModulesBuilder();

plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java

@@ -326,8 +326,7 @@ private DataSourceServiceImpl createDataSourceService() {
             .add(new SecurityLakeDataSourceFactory(pluginSettings))
             .build(),
         dataSourceMetadataStorage,
-        dataSourceUserAuthorizationHelper,
-        client);
+        dataSourceUserAuthorizationHelper);
   }
 
   @Override

plugin/src/main/java/org/opensearch/sql/plugin/config/OpenSearchPluginModule.java

@@ -98,13 +98,22 @@ public SQLService sqlService(QueryManager queryManager, QueryPlanFactory queryPl
   /** {@link QueryPlanFactory}. */
   @Provides
   public QueryPlanFactory queryPlanFactory(
-      DataSourceService dataSourceService, ExecutionEngine executionEngine, Settings settings) {
+      DataSourceService dataSourceService,
+      ExecutionEngine executionEngine,
+      Settings settings,
+      OpenSearchClient client) {
     Analyzer analyzer =
         new Analyzer(
             new ExpressionAnalyzer(functionRepository), dataSourceService, functionRepository);
     Planner planner = new Planner(LogicalPlanOptimizer.create());
     QueryService queryService =
-        new QueryService(analyzer, executionEngine, planner, dataSourceService, settings);
+        new QueryService(
+            analyzer,
+            executionEngine,
+            planner,
+            dataSourceService,
+            settings,
+            client.getNodeClient());
     return new QueryPlanFactory(queryService);
   }
 }

ppl/build.gradle

@@ -54,7 +54,6 @@ dependencies {
     api project(':common')
     api project(':core')
     api project(':protocol')
-    implementation group: 'org.opensearch.client', name: 'opensearch-rest-high-level-client', version: "${opensearch_version}"
 
     testImplementation group: 'junit', name: 'junit', version: '4.13.2'
     testImplementation group: 'org.hamcrest', name: 'hamcrest-library', version: "${hamcrest_version}"

ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLAbstractTest.java

@@ -74,7 +74,7 @@ protected CalcitePlanContext createBuilderContext() {
   private CalcitePlanContext createBuilderContext(UnaryOperator<RelBuilder.Config> transform) {
     config.context(Contexts.of(transform.apply(RelBuilder.Config.DEFAULT)));
     // NodeClient is not used in this test, so we pass null
-    return CalcitePlanContext.create(config.build(), PPL, null);
+    return CalcitePlanContext.create(config.build(), PPL);
   }
 
   /** Get the root RelNode of the given PPL query */