From 308303f3fb6c15074e938dca767cb519749b8b9d Mon Sep 17 00:00:00 2001 From: Yuya Ebihara Date: Wed, 18 Sep 2024 12:17:22 +0900 Subject: [PATCH] Add support for WITH SESSION clause Co-Authored-By: Mateusz "Serafin" Gajewski <66972+wendigo@users.noreply.github.com> --- .../src/main/java/io/trino/Session.java | 11 +- .../dispatcher/LocalDispatchQueryFactory.java | 5 + .../io/trino/execution/QueryStateMachine.java | 9 + .../io/trino/server/CoordinatorModule.java | 3 + .../sql/SessionSpecificationEvaluator.java | 208 ++++++++++++++++++ .../dispatcher/TestLocalDispatchQuery.java | 1 + .../execution/BaseDataDefinitionTaskTest.java | 1 + .../java/io/trino/execution/TestCallTask.java | 1 + .../io/trino/execution/TestCommitTask.java | 1 + .../execution/TestCreateCatalogTask.java | 1 + .../trino/execution/TestDeallocateTask.java | 1 + .../trino/execution/TestDropCatalogTask.java | 1 + .../io/trino/execution/TestPrepareTask.java | 1 + .../execution/TestQueryStateMachine.java | 1 + .../trino/execution/TestResetSessionTask.java | 1 + .../io/trino/execution/TestRoleTasks.java | 1 + .../io/trino/execution/TestRollbackTask.java | 1 + .../io/trino/execution/TestSetPathTask.java | 1 + .../TestSetSessionAuthorizationTask.java | 1 + .../trino/execution/TestSetSessionTask.java | 1 + .../trino/execution/TestSetTimeZoneTask.java | 1 + .../execution/TestStartTransactionTask.java | 1 + .../sql/planner/TestingPlannerContext.java | 2 - .../sql/query/TestSessionSpecifications.java | 151 +++++++++++++ .../AbstractTestEngineOnlyQueries.java | 10 +- 25 files changed, 408 insertions(+), 8 deletions(-) create mode 100644 core/trino-main/src/main/java/io/trino/sql/SessionSpecificationEvaluator.java create mode 100644 core/trino-main/src/test/java/io/trino/sql/query/TestSessionSpecifications.java diff --git a/core/trino-main/src/main/java/io/trino/Session.java b/core/trino-main/src/main/java/io/trino/Session.java index 6aeaab9590484..1413b30155e98 100644 --- a/core/trino-main/src/main/java/io/trino/Session.java +++ b/core/trino-main/src/main/java/io/trino/Session.java @@ -55,9 +55,11 @@ import static com.google.common.base.MoreObjects.toStringHelper; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; +import static io.trino.SystemSessionProperties.TIME_ZONE_ID; import static io.trino.client.ProtocolHeaders.TRINO_HEADERS; import static io.trino.spi.StandardErrorCode.CATALOG_NOT_FOUND; import static io.trino.spi.StandardErrorCode.NOT_FOUND; +import static io.trino.spi.type.TimeZoneKey.getTimeZoneKey; import static io.trino.sql.SqlPath.EMPTY_PATH; import static io.trino.util.Failures.checkCondition; import static java.util.Objects.requireNonNull; @@ -416,6 +418,11 @@ public Session withDefaultProperties(Map systemPropertyDefaults, .putAll(catalogEntry.getValue()); } + return withProperties(systemProperties, catalogProperties); + } + + public Session withProperties(Map systemProperties, Map> catalogProperties) + { return new Session( queryId, querySpan, @@ -428,7 +435,9 @@ public Session withDefaultProperties(Map systemPropertyDefaults, schema, path, traceToken, - timeZoneKey, + Optional.ofNullable(systemProperties.get(TIME_ZONE_ID)) + .map(TimeZoneKey::getTimeZoneKey) + .orElse(timeZoneKey), locale, remoteUserAddress, userAgent, diff --git a/core/trino-main/src/main/java/io/trino/dispatcher/LocalDispatchQueryFactory.java b/core/trino-main/src/main/java/io/trino/dispatcher/LocalDispatchQueryFactory.java index 5205b1268a9ef..5a990c4f0a7f8 100644 --- a/core/trino-main/src/main/java/io/trino/dispatcher/LocalDispatchQueryFactory.java +++ b/core/trino-main/src/main/java/io/trino/dispatcher/LocalDispatchQueryFactory.java @@ -38,6 +38,7 @@ import io.trino.server.protocol.Slug; import io.trino.spi.TrinoException; import io.trino.spi.resourcegroups.ResourceGroupId; +import io.trino.sql.SessionSpecificationEvaluator; import io.trino.sql.tree.Statement; import io.trino.transaction.TransactionId; import io.trino.transaction.TransactionManager; @@ -59,6 +60,7 @@ public class LocalDispatchQueryFactory private final TransactionManager transactionManager; private final AccessControl accessControl; private final Metadata metadata; + private final SessionSpecificationEvaluator sessionSpecificationEvaluator; private final QueryMonitor queryMonitor; private final LocationFactory locationFactory; @@ -77,6 +79,7 @@ public LocalDispatchQueryFactory( QueryManager queryManager, QueryManagerConfig queryManagerConfig, TransactionManager transactionManager, + SessionSpecificationEvaluator sessionSpecificationEvaluator, AccessControl accessControl, Metadata metadata, QueryMonitor queryMonitor, @@ -92,6 +95,7 @@ public LocalDispatchQueryFactory( this.transactionManager = requireNonNull(transactionManager, "transactionManager is null"); this.accessControl = requireNonNull(accessControl, "accessControl is null"); this.metadata = requireNonNull(metadata, "metadata is null"); + this.sessionSpecificationEvaluator = requireNonNull(sessionSpecificationEvaluator, "sessionSpecificationEvaluator is null"); this.queryMonitor = requireNonNull(queryMonitor, "queryMonitor is null"); this.locationFactory = requireNonNull(locationFactory, "locationFactory is null"); this.executionFactories = requireNonNull(executionFactories, "executionFactories is null"); @@ -132,6 +136,7 @@ public DispatchQuery createDispatchQuery( planOptimizersStatsCollector, getQueryType(preparedQuery.getStatement()), faultTolerantExecutionExchangeEncryptionEnabled, + Optional.of(sessionSpecificationEvaluator.getSessionSpecificationApplier(preparedQuery)), version); // It is important that `queryCreatedEvent` is called here. Moving it past the `executor.submit` below diff --git a/core/trino-main/src/main/java/io/trino/execution/QueryStateMachine.java b/core/trino-main/src/main/java/io/trino/execution/QueryStateMachine.java index 32dd0dc665ff2..c3850d0a2f05c 100644 --- a/core/trino-main/src/main/java/io/trino/execution/QueryStateMachine.java +++ b/core/trino-main/src/main/java/io/trino/execution/QueryStateMachine.java @@ -53,6 +53,7 @@ import io.trino.spi.resourcegroups.ResourceGroupId; import io.trino.spi.security.SelectedRole; import io.trino.spi.type.Type; +import io.trino.sql.SessionSpecificationEvaluator.SessionSpecificationsApplier; import io.trino.sql.analyzer.Output; import io.trino.sql.planner.PlanFragment; import io.trino.tracing.TrinoAttributes; @@ -243,6 +244,7 @@ public static QueryStateMachine begin( PlanOptimizersStatsCollector queryStatsCollector, Optional queryType, boolean faultTolerantExecutionExchangeEncryptionEnabled, + Optional sessionSpecificationsApplier, NodeVersion version) { return beginWithTicker( @@ -262,6 +264,7 @@ public static QueryStateMachine begin( queryStatsCollector, queryType, faultTolerantExecutionExchangeEncryptionEnabled, + sessionSpecificationsApplier, version); } @@ -282,6 +285,7 @@ static QueryStateMachine beginWithTicker( PlanOptimizersStatsCollector queryStatsCollector, Optional queryType, boolean faultTolerantExecutionExchangeEncryptionEnabled, + Optional sessionSpecificationsApplier, NodeVersion version) { // if there is an existing transaction, activate it @@ -308,6 +312,11 @@ static QueryStateMachine beginWithTicker( session = session.withExchangeEncryption(serializeAesEncryptionKey(createRandomAesEncryptionKey())); } + // Apply WITH SESSION specifications which require transaction to be started to resolve catalog handles + if (sessionSpecificationsApplier.isPresent()) { + session = sessionSpecificationsApplier.orElseThrow().apply(session); + } + Span querySpan = session.getQuerySpan(); querySpan.setAttribute(TrinoAttributes.QUERY_TYPE, queryType.map(Enum::name).orElse("UNKNOWN")); diff --git a/core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java b/core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java index 1820818ace22f..e19b63748bfc1 100644 --- a/core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java +++ b/core/trino-main/src/main/java/io/trino/server/CoordinatorModule.java @@ -112,6 +112,7 @@ import io.trino.server.ui.WorkerResource; import io.trino.spi.VersionEmbedder; import io.trino.sql.PlannerContext; +import io.trino.sql.SessionSpecificationEvaluator; import io.trino.sql.analyzer.AnalyzerFactory; import io.trino.sql.analyzer.QueryExplainerFactory; import io.trino.sql.planner.OptimizerStatsMBeanExporter; @@ -210,6 +211,8 @@ protected void setup(Binder binder) // dispatcher binder.bind(DispatchManager.class).in(Scopes.SINGLETON); + // WITH SESSION interpreter + binder.bind(SessionSpecificationEvaluator.class).in(Scopes.SINGLETON); // export under the old name, for backwards compatibility newExporter(binder).export(DispatchManager.class).as(generator -> generator.generatedNameOf(QueryManager.class)); binder.bind(FailedDispatchQueryFactory.class).in(Scopes.SINGLETON); diff --git a/core/trino-main/src/main/java/io/trino/sql/SessionSpecificationEvaluator.java b/core/trino-main/src/main/java/io/trino/sql/SessionSpecificationEvaluator.java new file mode 100644 index 0000000000000..ac40c31dee142 --- /dev/null +++ b/core/trino-main/src/main/java/io/trino/sql/SessionSpecificationEvaluator.java @@ -0,0 +1,208 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.sql; + +import com.google.common.collect.HashBasedTable; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Table; +import com.google.inject.Inject; +import io.trino.Session; +import io.trino.execution.QueryPreparer.PreparedQuery; +import io.trino.metadata.SessionPropertyManager; +import io.trino.security.AccessControl; +import io.trino.security.SecurityContext; +import io.trino.spi.TrinoException; +import io.trino.spi.connector.CatalogHandle; +import io.trino.spi.session.PropertyMetadata; +import io.trino.spi.type.Type; +import io.trino.sql.tree.Expression; +import io.trino.sql.tree.NodeRef; +import io.trino.sql.tree.Parameter; +import io.trino.sql.tree.QualifiedName; +import io.trino.sql.tree.Query; +import io.trino.sql.tree.SessionSpecification; + +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.function.Function; + +import static com.google.common.base.Preconditions.checkState; +import static io.trino.execution.ParameterExtractor.bindParameters; +import static io.trino.metadata.MetadataUtil.getRequiredCatalogHandle; +import static io.trino.metadata.SessionPropertyManager.evaluatePropertyValue; +import static io.trino.metadata.SessionPropertyManager.serializeSessionProperty; +import static io.trino.spi.StandardErrorCode.CATALOG_NOT_FOUND; +import static io.trino.spi.StandardErrorCode.INVALID_SESSION_PROPERTY; +import static io.trino.sql.analyzer.SemanticExceptions.semanticException; +import static java.lang.String.format; +import static java.util.Objects.requireNonNull; + +public class SessionSpecificationEvaluator +{ + private final PlannerContext plannerContext; + private final AccessControl accessControl; + private final SessionPropertyManager sessionPropertyManager; + + @Inject + public SessionSpecificationEvaluator(PlannerContext plannerContext, AccessControl accessControl, SessionPropertyManager sessionPropertyManager) + { + this.plannerContext = requireNonNull(plannerContext, "plannerContext is null"); + this.accessControl = requireNonNull(accessControl, "accessControl is null"); + this.sessionPropertyManager = requireNonNull(sessionPropertyManager, "sessionPropertyManager is null"); + } + + public SessionSpecificationsApplier getSessionSpecificationApplier(PreparedQuery preparedQuery) + { + if (!(preparedQuery.getStatement() instanceof Query queryStatement)) { + return session -> session; + } + return session -> prepareSession(session, queryStatement.getSessionProperties(), bindParameters(preparedQuery.getStatement(), preparedQuery.getParameters())); + } + + public Session prepareSession(Session session, List specifications, Map, Expression> parameters) + { + ResolvedSessionSpecifications resolvedSessionSpecifications = resolve(session, parameters, specifications); + return overrideProperties(session, resolvedSessionSpecifications); + } + + public ResolvedSessionSpecifications resolve(Session session, Map, Expression> parameters, List specifications) + { + ImmutableMap.Builder sessionProperties = ImmutableMap.builder(); + Table catalogProperties = HashBasedTable.create(); + Set seenPropertyNames = new HashSet<>(); + + for (SessionSpecification specification : specifications) { + List nameParts = specification.getName().getParts(); + + if (!seenPropertyNames.add(specification.getName())) { + throw semanticException(INVALID_SESSION_PROPERTY, specification, "Session property %s already set", specification.getName()); + } + + if (nameParts.size() == 1) { + Optional> systemSessionPropertyMetadata = sessionPropertyManager.getSystemSessionPropertyMetadata(nameParts.getFirst()); + if (systemSessionPropertyMetadata.isEmpty()) { + throw semanticException(INVALID_SESSION_PROPERTY, specification, "Session property %s does not exist", specification.getName()); + } + sessionProperties.put(nameParts.getFirst(), toSessionValue(session, parameters, specification, systemSessionPropertyMetadata.get())); + } + + if (nameParts.size() == 2) { + CatalogHandle catalogHandle = getRequiredCatalogHandle(plannerContext.getMetadata(), session, specification, nameParts.getFirst()); + Optional> connectorSessionPropertyMetadata = sessionPropertyManager.getConnectorSessionPropertyMetadata(catalogHandle, nameParts.getLast()); + if (connectorSessionPropertyMetadata.isEmpty()) { + throw semanticException(INVALID_SESSION_PROPERTY, specification, "Session property %s does not exist", specification.getName()); + } + catalogProperties.put(nameParts.get(0), nameParts.get(1), toSessionValue(session, parameters, specification, connectorSessionPropertyMetadata.get())); + } + } + + return new ResolvedSessionSpecifications(sessionProperties.buildOrThrow(), catalogProperties.rowMap()); + } + + public Session overrideProperties(Session session, ResolvedSessionSpecifications resolvedSessionSpecifications) + { + requireNonNull(resolvedSessionSpecifications, "resolvedSessionSpecifications is null"); + + validateSystemProperties(session, resolvedSessionSpecifications.systemProperties()); + + // Catalog session properties were already evaluated so we need to evaluate overrides + if (session.getTransactionId().isPresent()) { + validateCatalogProperties(session, resolvedSessionSpecifications.catalogProperties()); + } + + // NOTE: properties are validated before calling overrideProperties + Map systemProperties = new HashMap<>(); + systemProperties.putAll(session.getSystemProperties()); + systemProperties.putAll(resolvedSessionSpecifications.systemProperties()); + + Map> catalogProperties = new HashMap<>(session.getCatalogProperties()); + for (Map.Entry> catalogEntry : resolvedSessionSpecifications.catalogProperties().entrySet()) { + catalogProperties.computeIfAbsent(catalogEntry.getKey(), id -> new HashMap<>()) + .putAll(catalogEntry.getValue()); + } + + return session.withProperties(systemProperties, catalogProperties); + } + + private String toSessionValue(Session session, Map, Expression> parameters, SessionSpecification specification, PropertyMetadata propertyMetadata) + { + Type type = propertyMetadata.getSqlType(); + Object objectValue; + + try { + objectValue = evaluatePropertyValue(specification.getValue(), type, session, plannerContext, accessControl, parameters); + } + catch (TrinoException e) { + throw new TrinoException( + INVALID_SESSION_PROPERTY, + format("Unable to set session property '%s' to '%s': %s", specification.getName(), specification.getValue(), e.getRawMessage())); + } + + String value = serializeSessionProperty(type, objectValue); + // verify the SQL value can be decoded by the property + try { + propertyMetadata.decode(objectValue); + } + catch (RuntimeException e) { + throw semanticException(INVALID_SESSION_PROPERTY, specification, "%s", e.getMessage()); + } + + return value; + } + + private void validateSystemProperties(Session session, Map systemProperties) + { + for (Map.Entry property : systemProperties.entrySet()) { + // verify permissions + accessControl.checkCanSetSystemSessionProperty(session.getIdentity(), session.getQueryId(), property.getKey()); + // validate session property value + sessionPropertyManager.validateSystemSessionProperty(property.getKey(), property.getValue()); + } + } + + private void validateCatalogProperties(Session session, Map> catalogsProperties) + { + checkState(session.getTransactionId().isPresent(), "Not in transaction"); + for (Map.Entry> catalogProperties : catalogsProperties.entrySet()) { + CatalogHandle catalogHandle = plannerContext.getMetadata().getCatalogHandle(session, catalogProperties.getKey()) + .orElseThrow(() -> new TrinoException(CATALOG_NOT_FOUND, "Catalog '%s' not found".formatted(catalogProperties.getKey()))); + + for (Map.Entry catalogProperty : catalogProperties.getValue().entrySet()) { + // verify permissions + accessControl.checkCanSetCatalogSessionProperty(new SecurityContext(session.getRequiredTransactionId(), session.getIdentity(), session.getQueryId(), session.getStart()), catalogProperties.getKey(), catalogProperty.getKey()); + // validate catalog session property value + sessionPropertyManager.validateCatalogSessionProperty(catalogProperties.getKey(), catalogHandle, catalogProperty.getKey(), catalogProperty.getValue()); + } + } + } + + public record ResolvedSessionSpecifications(Map systemProperties, Map> catalogProperties) + { + public ResolvedSessionSpecifications + { + systemProperties = ImmutableMap.copyOf(requireNonNull(systemProperties, "systemProperties is null")); + catalogProperties = ImmutableMap.copyOf(requireNonNull(catalogProperties, "catalogProperties is null")); + } + } + + @FunctionalInterface + public interface SessionSpecificationsApplier + extends Function + { + } +} diff --git a/core/trino-main/src/test/java/io/trino/dispatcher/TestLocalDispatchQuery.java b/core/trino-main/src/test/java/io/trino/dispatcher/TestLocalDispatchQuery.java index d08ccccce6226..8ad863e045ee4 100644 --- a/core/trino-main/src/test/java/io/trino/dispatcher/TestLocalDispatchQuery.java +++ b/core/trino-main/src/test/java/io/trino/dispatcher/TestLocalDispatchQuery.java @@ -122,6 +122,7 @@ public void testSubmittedForDispatchedQuery() createPlanOptimizersStatsCollector(), Optional.of(QueryType.DATA_DEFINITION), true, + Optional.empty(), new NodeVersion("test")); QueryMonitor queryMonitor = new QueryMonitor( JsonCodec.jsonCodec(StageInfo.class), diff --git a/core/trino-main/src/test/java/io/trino/execution/BaseDataDefinitionTaskTest.java b/core/trino-main/src/test/java/io/trino/execution/BaseDataDefinitionTaskTest.java index 80403b075ec6d..62679eb64bec3 100644 --- a/core/trino-main/src/test/java/io/trino/execution/BaseDataDefinitionTaskTest.java +++ b/core/trino-main/src/test/java/io/trino/execution/BaseDataDefinitionTaskTest.java @@ -243,6 +243,7 @@ private static QueryStateMachine stateMachine(TransactionManager transactionMana createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); } diff --git a/core/trino-main/src/test/java/io/trino/execution/TestCallTask.java b/core/trino-main/src/test/java/io/trino/execution/TestCallTask.java index 683cdf651a5df..d182af5509f75 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestCallTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestCallTask.java @@ -173,6 +173,7 @@ private QueryStateMachine stateMachine(TransactionManager transactionManager, Me createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); } diff --git a/core/trino-main/src/test/java/io/trino/execution/TestCommitTask.java b/core/trino-main/src/test/java/io/trino/execution/TestCommitTask.java index 12413dc1c381c..73750b17300df 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestCommitTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestCommitTask.java @@ -147,6 +147,7 @@ private QueryStateMachine createQueryStateMachine(String query, Session session, createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); } diff --git a/core/trino-main/src/test/java/io/trino/execution/TestCreateCatalogTask.java b/core/trino-main/src/test/java/io/trino/execution/TestCreateCatalogTask.java index 5fbf889d2d278..310d5dd910dd2 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestCreateCatalogTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestCreateCatalogTask.java @@ -84,6 +84,7 @@ public void setUp() createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); this.queryRunner = queryRunner; diff --git a/core/trino-main/src/test/java/io/trino/execution/TestDeallocateTask.java b/core/trino-main/src/test/java/io/trino/execution/TestDeallocateTask.java index a8bc91cedd206..5e64ad00d8d5a 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestDeallocateTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestDeallocateTask.java @@ -116,6 +116,7 @@ private Set executeDeallocate(String statementName, String sqlString, Se createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); Deallocate deallocate = new Deallocate(new NodeLocation(1, 1), new Identifier(statementName)); new DeallocateTask().execute(deallocate, stateMachine, emptyList(), WarningCollector.NOOP); diff --git a/core/trino-main/src/test/java/io/trino/execution/TestDropCatalogTask.java b/core/trino-main/src/test/java/io/trino/execution/TestDropCatalogTask.java index fc84d5a368b6c..73dc2065ea0f2 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestDropCatalogTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestDropCatalogTask.java @@ -129,6 +129,7 @@ private QueryStateMachine createNewQuery() createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); } } diff --git a/core/trino-main/src/test/java/io/trino/execution/TestPrepareTask.java b/core/trino-main/src/test/java/io/trino/execution/TestPrepareTask.java index 535bf717c8d8f..5504dc0896cb1 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestPrepareTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestPrepareTask.java @@ -139,6 +139,7 @@ private Map executePrepare(String statementName, Statement state createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); Prepare prepare = new Prepare(identifier(statementName), statement); new PrepareTask(new SqlParser()).execute(prepare, stateMachine, emptyList(), WarningCollector.NOOP); diff --git a/core/trino-main/src/test/java/io/trino/execution/TestQueryStateMachine.java b/core/trino-main/src/test/java/io/trino/execution/TestQueryStateMachine.java index 75221472da370..27f3e949319b3 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestQueryStateMachine.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestQueryStateMachine.java @@ -867,6 +867,7 @@ public QueryStateMachine build() createPlanOptimizersStatsCollector(), QUERY_TYPE, false, + Optional.empty(), new NodeVersion("test")); stateMachine.setInputs(INPUTS); stateMachine.setOutput(OUTPUT); diff --git a/core/trino-main/src/test/java/io/trino/execution/TestResetSessionTask.java b/core/trino-main/src/test/java/io/trino/execution/TestResetSessionTask.java index e145ebd366382..6b1d24ab086a0 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestResetSessionTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestResetSessionTask.java @@ -120,6 +120,7 @@ public void test() createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); getFutureValue(new ResetSessionTask(metadata, sessionPropertyManager).execute( diff --git a/core/trino-main/src/test/java/io/trino/execution/TestRoleTasks.java b/core/trino-main/src/test/java/io/trino/execution/TestRoleTasks.java index 310df2570f3ed..3114f84c27573 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestRoleTasks.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestRoleTasks.java @@ -176,6 +176,7 @@ protected QueryStateMachine execute(String statement, Data createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); task.execute((T) parser.createStatement(statement), stateMachine, ImmutableList.of(), WarningCollector.NOOP); return stateMachine; diff --git a/core/trino-main/src/test/java/io/trino/execution/TestRollbackTask.java b/core/trino-main/src/test/java/io/trino/execution/TestRollbackTask.java index 9250365bd59a1..72ed96c2ee943 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestRollbackTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestRollbackTask.java @@ -137,6 +137,7 @@ private QueryStateMachine createQueryStateMachine(String query, Session session, createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); } diff --git a/core/trino-main/src/test/java/io/trino/execution/TestSetPathTask.java b/core/trino-main/src/test/java/io/trino/execution/TestSetPathTask.java index 117efdd72cb41..425ba852faa5a 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestSetPathTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestSetPathTask.java @@ -124,6 +124,7 @@ private QueryStateMachine createQueryStateMachine(String query) createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); } diff --git a/core/trino-main/src/test/java/io/trino/execution/TestSetSessionAuthorizationTask.java b/core/trino-main/src/test/java/io/trino/execution/TestSetSessionAuthorizationTask.java index ced3605669405..0ac902af6ebb0 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestSetSessionAuthorizationTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestSetSessionAuthorizationTask.java @@ -124,6 +124,7 @@ private QueryStateMachine createStateMachine(Optional transaction createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); return stateMachine; } diff --git a/core/trino-main/src/test/java/io/trino/execution/TestSetSessionTask.java b/core/trino-main/src/test/java/io/trino/execution/TestSetSessionTask.java index 26a66a6995102..2d5038a2ffe28 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestSetSessionTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestSetSessionTask.java @@ -208,6 +208,7 @@ private void testSetSessionWithParameters(String property, Expression expression createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); getFutureValue(new SetSessionTask(plannerContext, accessControl, sessionPropertyManager).execute(new SetSession(new NodeLocation(1, 1), qualifiedPropName, expression), stateMachine, parameters, WarningCollector.NOOP)); diff --git a/core/trino-main/src/test/java/io/trino/execution/TestSetTimeZoneTask.java b/core/trino-main/src/test/java/io/trino/execution/TestSetTimeZoneTask.java index 726cfa4436559..c5401a5d99857 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestSetTimeZoneTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestSetTimeZoneTask.java @@ -266,6 +266,7 @@ private QueryStateMachine createQueryStateMachine(String query) createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); } diff --git a/core/trino-main/src/test/java/io/trino/execution/TestStartTransactionTask.java b/core/trino-main/src/test/java/io/trino/execution/TestStartTransactionTask.java index 04676e288dbb5..acf582ff2314f 100644 --- a/core/trino-main/src/test/java/io/trino/execution/TestStartTransactionTask.java +++ b/core/trino-main/src/test/java/io/trino/execution/TestStartTransactionTask.java @@ -268,6 +268,7 @@ private QueryStateMachine createQueryStateMachine(String query, Session session, createPlanOptimizersStatsCollector(), Optional.empty(), true, + Optional.empty(), new NodeVersion("test")); } diff --git a/core/trino-main/src/test/java/io/trino/sql/planner/TestingPlannerContext.java b/core/trino-main/src/test/java/io/trino/sql/planner/TestingPlannerContext.java index aff33e988cd9e..45983c22e82df 100644 --- a/core/trino-main/src/test/java/io/trino/sql/planner/TestingPlannerContext.java +++ b/core/trino-main/src/test/java/io/trino/sql/planner/TestingPlannerContext.java @@ -77,14 +77,12 @@ private Builder() {} public Builder withMetadata(Metadata metadata) { checkState(this.metadata == null, "metadata already set"); - checkState(this.transactionManager == null, "transactionManager already set"); this.metadata = requireNonNull(metadata, "metadata is null"); return this; } public Builder withTransactionManager(TransactionManager transactionManager) { - checkState(this.metadata == null, "metadata already set"); checkState(this.transactionManager == null, "transactionManager already set"); this.transactionManager = requireNonNull(transactionManager, "transactionManager is null"); return this; diff --git a/core/trino-main/src/test/java/io/trino/sql/query/TestSessionSpecifications.java b/core/trino-main/src/test/java/io/trino/sql/query/TestSessionSpecifications.java new file mode 100644 index 0000000000000..e4558e51f9f1f --- /dev/null +++ b/core/trino-main/src/test/java/io/trino/sql/query/TestSessionSpecifications.java @@ -0,0 +1,151 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.sql.query; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import io.trino.FeaturesConfig; +import io.trino.Session; +import io.trino.SystemSessionProperties; +import io.trino.metadata.AbstractMockMetadata; +import io.trino.metadata.MetadataManager; +import io.trino.metadata.ResolvedFunction; +import io.trino.metadata.SessionPropertyManager; +import io.trino.metadata.TypeRegistry; +import io.trino.security.AllowAllAccessControl; +import io.trino.spi.connector.CatalogHandle; +import io.trino.spi.function.OperatorType; +import io.trino.spi.session.PropertyMetadata; +import io.trino.spi.type.Type; +import io.trino.spi.type.TypeManager; +import io.trino.spi.type.TypeOperators; +import io.trino.sql.PlannerContext; +import io.trino.sql.SessionSpecificationEvaluator; +import io.trino.sql.parser.SqlParser; +import io.trino.sql.tree.SessionSpecification; +import io.trino.transaction.TestingTransactionManager; +import io.trino.transaction.TransactionManager; +import io.trino.type.InternalTypeManager; +import org.intellij.lang.annotations.Language; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.parallel.Execution; + +import java.util.Map; +import java.util.Optional; + +import static io.trino.sql.planner.TestingPlannerContext.plannerContextBuilder; +import static io.trino.testing.TestingSession.testSession; +import static io.trino.testing.TransactionBuilder.transaction; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS; +import static org.junit.jupiter.api.parallel.ExecutionMode.CONCURRENT; + +@TestInstance(PER_CLASS) +@Execution(CONCURRENT) +final class TestSessionSpecifications +{ + private static final SqlParser SQL_PARSER = new SqlParser(); + private static final SessionPropertyManager SESSION_PROPERTY_MANAGER = new SessionPropertyManager(ImmutableSet.of(new SystemSessionProperties()), catalogHandle -> Map.of( + "catalog_property", PropertyMetadata.stringProperty("catalog_property", "Test catalog property", "", false))); + + @Test + void testParseSystemSessionProperty() + { + assertThatThrownBy(() -> analyze("SESSION invalid_key = 'invalid_value'")) + .hasMessageContaining("line 1:1: Session property invalid_key does not exist"); + + assertThatThrownBy(() -> analyze("SESSION optimize_hash_generation = 'invalid_value'")) + .hasMessageContaining("Unable to set session property 'optimize_hash_generation' to ''invalid_value'': Cannot cast type varchar(13) to boolean"); + + assertThat(analyze("SESSION optimize_hash_generation = true").getSystemProperties()) + .isEqualTo(Map.of("optimize_hash_generation", "true")); + + assertThat(analyze("SESSION optimize_hash_generation = CAST('true' AS boolean)").getSystemProperties()) + .isEqualTo(Map.of("optimize_hash_generation", "true")); + + assertThatThrownBy(() -> analyze("SESSION optimize_hash_generation = true", "SESSION optimize_hash_generation = false")) + .hasMessageContaining("line 1:1: Session property optimize_hash_generation already set"); + } + + @Test + void testCatalogSessionProperty() + { + assertThatThrownBy(() -> analyze("SESSION test.invalid_key = 'invalid_value'")) + .hasMessageContaining("line 1:1: Session property test.invalid_key does not exist"); + + assertThatThrownBy(() -> analyze("SESSION test.catalog_property = true")) + .hasMessageContaining("Unable to set session property 'test.catalog_property' to 'true': Cannot cast type boolean to varchar"); + + assertThat(analyze("SESSION test.catalog_property = 'true'").getCatalogProperties("test")) + .isEqualTo(Map.of("catalog_property", "true")); + + assertThat(analyze("SESSION test.catalog_property = CAST(true AS varchar)").getCatalogProperties("test")) + .isEqualTo(Map.of("catalog_property", "true")); + + assertThatThrownBy(() -> analyze("SESSION test.catalog_property = 'true'", "SESSION test.catalog_property = 'false'").getCatalogProperties("test")) + .hasMessageContaining("line 1:1: Session property test.catalog_property already set"); + } + + private static Session analyze(@Language("SQL") String... statements) + { + ImmutableList.Builder sessionSpecifications = ImmutableList.builder(); + for (String statement : statements) { + sessionSpecifications.add(SQL_PARSER.createSessionSpecification(statement)); + } + + TransactionManager transactionManager = new TestingTransactionManager(); + PlannerContext plannerContext = plannerContextBuilder() + .withMetadata(new MockMetadata()) + .withTransactionManager(transactionManager) + .build(); + + return transaction(transactionManager, plannerContext.getMetadata(), new AllowAllAccessControl()) + .execute(testSession(), transactionSession -> { + SessionSpecificationEvaluator evaluator = new SessionSpecificationEvaluator(plannerContext, new AllowAllAccessControl(), SESSION_PROPERTY_MANAGER); + return evaluator.prepareSession(transactionSession, sessionSpecifications.build(), Map.of()); + }); + } + + private static class MockMetadata + extends AbstractMockMetadata + { + private final MetadataManager delegate; + + public MockMetadata() + { + FeaturesConfig featuresConfig = new FeaturesConfig(); + TypeOperators typeOperators = new TypeOperators(); + + TypeRegistry typeRegistry = new TypeRegistry(typeOperators, featuresConfig); + TypeManager typeManager = new InternalTypeManager(typeRegistry); + this.delegate = MetadataManager.testMetadataManagerBuilder() + .withTypeManager(typeManager) + .build(); + } + + @Override + public ResolvedFunction getCoercion(OperatorType operatorType, Type fromType, Type toType) + { + return delegate.getCoercion(operatorType, fromType, toType); + } + + @Override + public Optional getCatalogHandle(Session session, String catalogName) + { + return Optional.of(CatalogHandle.fromId(catalogName + ":NORMAL:v1")); + } + } +} diff --git a/testing/trino-testing/src/main/java/io/trino/testing/AbstractTestEngineOnlyQueries.java b/testing/trino-testing/src/main/java/io/trino/testing/AbstractTestEngineOnlyQueries.java index b793fed4c0bba..43d4e8fd2d3bd 100644 --- a/testing/trino-testing/src/main/java/io/trino/testing/AbstractTestEngineOnlyQueries.java +++ b/testing/trino-testing/src/main/java/io/trino/testing/AbstractTestEngineOnlyQueries.java @@ -6622,13 +6622,13 @@ public void testColumnNames() public void testInlineSession() { assertThat(query("WITH SESSION time_zone_id = 'Europe/Wonderland' SELECT current_timezone()")) - .failure() - .hasMessageContaining("Time zone not supported: Europe/Wonderland"); + .failure() + .hasMessageContaining("Time zone not supported: Europe/Wonderland"); assertThat(query("WITH SESSION time_zone_id = 'Europe/Warsaw' SELECT current_timezone()")) - .succeeds() - .skippingTypesCheck() - .matches("VALUES 'Europe/Warsaw'"); + .succeeds() + .skippingTypesCheck() + .matches("VALUES 'Europe/Warsaw'"); } @Test