-
Notifications
You must be signed in to change notification settings - Fork 208
Extract unified query context for shared config management #4933
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
dai-chen
merged 9 commits into
opensearch-project:main
from
dai-chen:add-unified-query-context
Dec 17, 2025
Merged
Changes from all commits
Commits
Show all changes
9 commits
Select commit
Hold shift + click to select a range
b0d5ad9
Extract UnifiedQueryPlanner.Builder to UnifiedQueryContext
dai-chen 89f3542
Remove backward compatibility code
dai-chen d9d852a
Refactor unified query context with setting setter
dai-chen 2e5da1a
Initialize unified query context with default system limits
dai-chen 5d4595c
Refactor setting map read
dai-chen d82c836
Update javadoc and rename queryType to language
dai-chen 2137ab6
Address AI comments
dai-chen 9c96594
Reuse context in test base class
dai-chen f650fb5
Remove session in doc
dai-chen File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
189 changes: 189 additions & 0 deletions
189
api/src/main/java/org/opensearch/sql/api/UnifiedQueryContext.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,189 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.sql.api; | ||
|
|
||
| import static org.opensearch.sql.common.setting.Settings.Key.PPL_JOIN_SUBSEARCH_MAXOUT; | ||
| import static org.opensearch.sql.common.setting.Settings.Key.PPL_SUBSEARCH_MAXOUT; | ||
| import static org.opensearch.sql.common.setting.Settings.Key.QUERY_SIZE_LIMIT; | ||
|
|
||
| import java.util.HashMap; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.Objects; | ||
| import lombok.Value; | ||
| import org.apache.calcite.jdbc.CalciteSchema; | ||
| import org.apache.calcite.plan.RelTraitDef; | ||
| import org.apache.calcite.rel.metadata.DefaultRelMetadataProvider; | ||
| import org.apache.calcite.schema.Schema; | ||
| import org.apache.calcite.schema.SchemaPlus; | ||
| import org.apache.calcite.sql.parser.SqlParser; | ||
| import org.apache.calcite.tools.FrameworkConfig; | ||
| import org.apache.calcite.tools.Frameworks; | ||
| import org.apache.calcite.tools.Programs; | ||
| import org.opensearch.sql.calcite.CalcitePlanContext; | ||
| import org.opensearch.sql.calcite.SysLimit; | ||
| import org.opensearch.sql.common.setting.Settings; | ||
| import org.opensearch.sql.executor.QueryType; | ||
|
|
||
| /** | ||
| * A reusable abstraction shared across unified query components (planner, compiler, etc.). This | ||
| * centralizes configuration for catalog schemas, query type, execution limits, and other settings, | ||
| * enabling consistent behavior across all unified query operations. | ||
| */ | ||
| @Value | ||
| public class UnifiedQueryContext { | ||
|
|
||
| /** CalcitePlanContext containing Calcite framework configuration and query type. */ | ||
| CalcitePlanContext planContext; | ||
|
|
||
| /** Settings containing execution limits and feature flags used by parsers and planners. */ | ||
| Settings settings; | ||
|
|
||
| /** Creates a new builder for UnifiedQueryContext. */ | ||
| public static Builder builder() { | ||
| return new Builder(); | ||
| } | ||
|
|
||
| /** Builder that constructs UnifiedQueryContext. */ | ||
| public static class Builder { | ||
| private QueryType queryType; | ||
| private final Map<String, Schema> catalogs = new HashMap<>(); | ||
| private String defaultNamespace; | ||
| private boolean cacheMetadata = false; | ||
|
|
||
| /** | ||
| * Setting values with defaults from SysLimit.DEFAULT. Only includes planning-required settings | ||
| * to avoid coupling with OpenSearchSettings. | ||
| */ | ||
| private final Map<Settings.Key, Object> settings = | ||
| new HashMap<Settings.Key, Object>( | ||
| Map.of( | ||
| QUERY_SIZE_LIMIT, SysLimit.DEFAULT.querySizeLimit(), | ||
| PPL_SUBSEARCH_MAXOUT, SysLimit.DEFAULT.subsearchLimit(), | ||
| PPL_JOIN_SUBSEARCH_MAXOUT, SysLimit.DEFAULT.joinSubsearchLimit())); | ||
|
|
||
| /** | ||
| * Sets the query language frontend to be used. | ||
| * | ||
| * @param queryType the {@link QueryType}, such as PPL | ||
| * @return this builder instance | ||
| */ | ||
| public Builder language(QueryType queryType) { | ||
| this.queryType = queryType; | ||
| return this; | ||
| } | ||
|
|
||
| /** | ||
| * Registers a catalog with the specified name and its associated schema. The schema can be a | ||
| * flat or nested structure (e.g., catalog → schema → table), depending on how data is | ||
| * organized. | ||
| * | ||
| * @param name the name of the catalog to register | ||
| * @param schema the schema representing the structure of the catalog | ||
| * @return this builder instance | ||
| */ | ||
| public Builder catalog(String name, Schema schema) { | ||
| catalogs.put(name, schema); | ||
| return this; | ||
| } | ||
|
|
||
| /** | ||
| * Sets the default namespace path for resolving unqualified table names. | ||
| * | ||
| * @param namespace dot-separated path (e.g., "spark_catalog.default" or "opensearch") | ||
| * @return this builder instance | ||
| */ | ||
| public Builder defaultNamespace(String namespace) { | ||
| this.defaultNamespace = namespace; | ||
| return this; | ||
| } | ||
|
|
||
| /** | ||
| * Enables or disables catalog metadata caching in the root schema. | ||
| * | ||
| * @param cache whether to enable metadata caching | ||
| * @return this builder instance | ||
| */ | ||
| public Builder cacheMetadata(boolean cache) { | ||
| this.cacheMetadata = cache; | ||
| return this; | ||
| } | ||
|
|
||
| /** | ||
| * Sets a specific setting value by name. | ||
| * | ||
| * @param name the setting key name (e.g., "plugins.query.size_limit") | ||
| * @param value the setting value | ||
| * @throws IllegalArgumentException if the setting name is not recognized | ||
| */ | ||
| public Builder setting(String name, Object value) { | ||
| Settings.Key key = | ||
| Settings.Key.of(name) | ||
| .orElseThrow(() -> new IllegalArgumentException("Unknown setting name: " + name)); | ||
| settings.put(key, value); | ||
| return this; | ||
| } | ||
|
|
||
| /** | ||
| * Builds a {@link UnifiedQueryContext} with the configuration. | ||
| * | ||
| * @return a new instance of {@link UnifiedQueryContext} | ||
| */ | ||
| public UnifiedQueryContext build() { | ||
| Objects.requireNonNull(queryType, "Must specify language before build"); | ||
|
|
||
| Settings settings = buildSettings(); | ||
| CalcitePlanContext planContext = | ||
| CalcitePlanContext.create( | ||
| buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType); | ||
| return new UnifiedQueryContext(planContext, settings); | ||
| } | ||
|
|
||
| private Settings buildSettings() { | ||
| return new Settings() { | ||
| @Override | ||
| @SuppressWarnings("unchecked") | ||
| public <T> T getSettingValue(Key key) { | ||
| return (T) settings.get(key); | ||
| } | ||
|
|
||
| @Override | ||
| public List<?> getSettings() { | ||
| return List.copyOf(settings.entrySet()); | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| @SuppressWarnings({"rawtypes"}) | ||
| private FrameworkConfig buildFrameworkConfig() { | ||
| SchemaPlus rootSchema = CalciteSchema.createRootSchema(true, cacheMetadata).plus(); | ||
| catalogs.forEach(rootSchema::add); | ||
|
|
||
| SchemaPlus defaultSchema = findSchemaByPath(rootSchema, defaultNamespace); | ||
| return Frameworks.newConfigBuilder() | ||
| .parserConfig(SqlParser.Config.DEFAULT) | ||
| .defaultSchema(defaultSchema) | ||
| .traitDefs((List<RelTraitDef>) null) | ||
| .programs(Programs.calc(DefaultRelMetadataProvider.INSTANCE)) | ||
| .build(); | ||
| } | ||
|
|
||
| private SchemaPlus findSchemaByPath(SchemaPlus rootSchema, String defaultPath) { | ||
| if (defaultPath == null) { | ||
| return rootSchema; | ||
| } | ||
|
|
||
| SchemaPlus current = rootSchema; | ||
| for (String part : defaultPath.split("\\.")) { | ||
| current = current.getSubSchema(part); | ||
| if (current == null) { | ||
| throw new IllegalArgumentException("Invalid default catalog path: " + defaultPath); | ||
| } | ||
| } | ||
| return current; | ||
| } | ||
| } | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.