From af088116ea4f6fbf55969950ca5f052a4af5f29f Mon Sep 17 00:00:00 2001 From: hantangwangd Date: Tue, 26 Nov 2024 17:29:44 +0800 Subject: [PATCH] Fix `USE` statement when schema does not exist --- .../facebook/presto/execution/UseTask.java | 18 ++++++++++++++- .../presto/metadata/MetadataUtil.java | 3 +-- .../presto/spark/TestPrestoQueries.java | 7 ++++++ .../presto/tests/AbstractTestQueries.java | 22 +++++++++++++++++++ .../presto/tests/TestLocalQueries.java | 6 +++++ .../tests/TestQueryPlanDeterminism.java | 6 +++++ .../tests/TestVerboseOptimizerInfo.java | 6 +++++ 7 files changed, 65 insertions(+), 3 deletions(-) diff --git a/presto-main/src/main/java/com/facebook/presto/execution/UseTask.java b/presto-main/src/main/java/com/facebook/presto/execution/UseTask.java index f677165a97f64..03a582bcec592 100644 --- a/presto-main/src/main/java/com/facebook/presto/execution/UseTask.java +++ b/presto-main/src/main/java/com/facebook/presto/execution/UseTask.java @@ -14,10 +14,12 @@ package com.facebook.presto.execution; import com.facebook.presto.Session; +import com.facebook.presto.common.CatalogSchemaName; import com.facebook.presto.metadata.Metadata; import com.facebook.presto.spi.security.AccessControl; import com.facebook.presto.sql.analyzer.SemanticException; import com.facebook.presto.sql.tree.Expression; +import com.facebook.presto.sql.tree.Identifier; import com.facebook.presto.sql.tree.Use; import com.facebook.presto.transaction.TransactionManager; import com.google.common.util.concurrent.ListenableFuture; @@ -26,7 +28,9 @@ import static com.facebook.presto.metadata.MetadataUtil.getConnectorIdOrThrow; import static com.facebook.presto.sql.analyzer.SemanticErrorCode.CATALOG_NOT_SPECIFIED; +import static com.facebook.presto.sql.analyzer.SemanticErrorCode.MISSING_SCHEMA; import static com.google.common.util.concurrent.Futures.immediateFuture; +import static java.lang.String.format; public class UseTask implements SessionTransactionControlTask @@ -52,7 +56,7 @@ public ListenableFuture execute( checkAndSetCatalog(statement, metadata, stateMachine, session); - stateMachine.setSetSchema(statement.getSchema().getValueLowerCase()); + checkAndSetSchema(statement, metadata, stateMachine, session); return immediateFuture(null); } @@ -72,4 +76,16 @@ private void checkAndSetCatalog(Use statement, Metadata metadata, QueryStateMach stateMachine.setSetCatalog(catalog); } } + + private void checkAndSetSchema(Use statement, Metadata metadata, QueryStateMachine stateMachine, Session session) + { + String catalog = statement.getCatalog() + .map(Identifier::getValueLowerCase) + .orElseGet(() -> session.getCatalog().map(String::toLowerCase).get()); + String schema = statement.getSchema().getValueLowerCase(); + if (!metadata.getMetadataResolver(session).schemaExists(new CatalogSchemaName(catalog, schema))) { + throw new SemanticException(MISSING_SCHEMA, format("Schema does not exist: %s.%s", catalog, schema)); + } + stateMachine.setSetSchema(schema); + } } diff --git a/presto-main/src/main/java/com/facebook/presto/metadata/MetadataUtil.java b/presto-main/src/main/java/com/facebook/presto/metadata/MetadataUtil.java index 6e13a37eccbf5..7433057b3a30a 100644 --- a/presto-main/src/main/java/com/facebook/presto/metadata/MetadataUtil.java +++ b/presto-main/src/main/java/com/facebook/presto/metadata/MetadataUtil.java @@ -41,7 +41,6 @@ import java.util.List; import java.util.Optional; -import static com.facebook.presto.spi.StandardErrorCode.NOT_FOUND; import static com.facebook.presto.spi.StandardErrorCode.SYNTAX_ERROR; import static com.facebook.presto.spi.security.PrincipalType.ROLE; import static com.facebook.presto.spi.security.PrincipalType.USER; @@ -94,7 +93,7 @@ public static SchemaTableName toSchemaTableName(QualifiedObjectName qualifiedObj public static ConnectorId getConnectorIdOrThrow(Session session, Metadata metadata, String catalogName) { return metadata.getCatalogHandle(session, catalogName) - .orElseThrow(() -> new PrestoException(NOT_FOUND, "Catalog does not exist: " + catalogName)); + .orElseThrow(() -> new SemanticException(MISSING_CATALOG, "Catalog does not exist: " + catalogName)); } public static ConnectorId getConnectorIdOrThrow(Session session, Metadata metadata, String catalogName, Statement statement, String errorMsg) diff --git a/presto-spark-base/src/test/java/com/facebook/presto/spark/TestPrestoQueries.java b/presto-spark-base/src/test/java/com/facebook/presto/spark/TestPrestoQueries.java index 4da201f77a9dc..17c9c957cca20 100644 --- a/presto-spark-base/src/test/java/com/facebook/presto/spark/TestPrestoQueries.java +++ b/presto-spark-base/src/test/java/com/facebook/presto/spark/TestPrestoQueries.java @@ -15,6 +15,7 @@ import com.facebook.presto.testing.QueryRunner; import com.facebook.presto.tests.AbstractTestQueries; +import org.testng.annotations.Test; public class TestPrestoQueries extends AbstractTestQueries @@ -193,4 +194,10 @@ public void testSetSessionNativeWorkerSessionProperty() { // prepared statement is not supported by Presto on Spark } + + @Test + public void testUse() + { + // USE statement is not supported + } } diff --git a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java index 787d43f37279f..080c530875095 100644 --- a/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java +++ b/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java @@ -2788,6 +2788,28 @@ public void testShowSchemasFrom() assertTrue(result.getOnlyColumnAsSet().containsAll(ImmutableSet.of(getSession().getSchema().get(), INFORMATION_SCHEMA))); } + @Test + public void testUse() + { + Session sessionWithDefaultCatalogAndSchema = getSession(); + String catalog = sessionWithDefaultCatalogAndSchema.getCatalog().get(); + String schema = sessionWithDefaultCatalogAndSchema.getSchema().get(); + + assertQueryFails(sessionWithDefaultCatalogAndSchema, "USE non_exist_schema", format("Schema does not exist: %s.non_exist_schema", catalog)); + assertQueryFails(sessionWithDefaultCatalogAndSchema, "USE non_exist_catalog.any_schema", "Catalog does not exist: non_exist_catalog"); + assertQueryFails(sessionWithDefaultCatalogAndSchema, format("USE %s.non_exist_schema", catalog), format("Schema does not exist: %s.non_exist_schema", catalog)); + assertUpdate(sessionWithDefaultCatalogAndSchema, format("USE %s.%s", catalog, schema)); + + Session sessionWithoutDefaultCatalogAndSchema = Session.builder(getSession()) + .setCatalog(null) + .setSchema(null) + .build(); + assertQueryFails(sessionWithoutDefaultCatalogAndSchema, "USE any_schema", ".* Catalog must be specified when session catalog is not set"); + assertQueryFails(sessionWithoutDefaultCatalogAndSchema, "USE non_exist_catalog.any_schema", "Catalog does not exist: non_exist_catalog"); + assertQueryFails(sessionWithoutDefaultCatalogAndSchema, format("USE %s.non_exist_schema", catalog), format("Schema does not exist: %s.non_exist_schema", catalog)); + assertUpdate(sessionWithoutDefaultCatalogAndSchema, format("USE %s.%s", catalog, schema)); + } + @Test public void testShowSchemasLike() { diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java index 63d4fc2e5838a..b381b01f5ada3 100644 --- a/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java +++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestLocalQueries.java @@ -117,6 +117,12 @@ public void testDecimal() assertQuery("SELECT 0.1", "SELECT CAST('0.1' AS DECIMAL)"); } + @Test + public void testUse() + { + // USE statement is not supported + } + @Test public void testIOExplain() { diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java index 9f51fc2130404..98fdee5fe7ebf 100644 --- a/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java +++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestQueryPlanDeterminism.java @@ -216,6 +216,12 @@ protected void assertTableColumnNames(String tableName, String... columnNames) { } + @Test + public void testUse() + { + // USE statement is not supported + } + @Override protected MaterializedResult computeExpected(@Language("SQL") String sql, List resultTypes) { diff --git a/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java b/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java index 2cbcaf93babc8..3cb9f26f1b8fd 100644 --- a/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java +++ b/presto-tests/src/test/java/com/facebook/presto/tests/TestVerboseOptimizerInfo.java @@ -157,6 +157,12 @@ public void testSetSessionNativeWorkerSessionProperty() { } + @Test + public void testUse() + { + // USE statement is not supported + } + private void checkOptimizerInfo(String explain, String optimizerType, List optimizers) { checkOptimizerInfo(explain, optimizerType, optimizers, new ArrayList<>());