Fix USE statement when schema does not exist

prestodb · Nov 27, 2024 · af08811 · af08811
1 parent 4912ca7
commit af08811
Show file tree

Hide file tree

Showing 7 changed files with 65 additions and 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
@@ -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<Use>
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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<? extends Type> 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
@@ -157,6 +157,12 @@ public void testSetSessionNativeWorkerSessionProperty()
     {
     }
 
+    @Test
+    public void testUse()
+    {
+        // USE statement is not supported
+    }
+
     private void checkOptimizerInfo(String explain, String optimizerType, List<String> optimizers)
     {
         checkOptimizerInfo(explain, optimizerType, optimizers, new ArrayList<>());