[SPARK-28970][SQL] Implement USE CATALOG/NAMESPACE for Data Source V2

docs/sql-keywords.md

@@ -57,6 +57,7 @@ Below is a list of all the keywords in Spark SQL.
   <tr><td>CASCADE</td><td>non-reserved</td><td>non-reserved</td><td>reserved</td></tr>
   <tr><td>CASE</td><td>reserved</td><td>non-reserved</td><td>reserved</td></tr>
   <tr><td>CAST</td><td>reserved</td><td>non-reserved</td><td>reserved</td></tr>
+  <tr><td>CATALOG</td><td>reserved</td><td>reserved</td><td>non-reserved</td></tr>
   <tr><td>CHANGE</td><td>non-reserved</td><td>non-reserved</td><td>non-reserved</td></tr>
   <tr><td>CHECK</td><td>reserved</td><td>non-reserved</td><td>reserved</td></tr>
   <tr><td>CLEAR</td><td>non-reserved</td><td>non-reserved</td><td>non-reserved</td></tr>

sql/catalyst/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBase.g4

@@ -82,7 +82,9 @@ singleTableSchema
 statement
     : query                                                            #statementDefault
     | ctes? dmlStatementNoWith                                         #dmlStatement
-    | USE db=errorCapturingIdentifier                                  #use
+    | USE CATALOG catalog=errorCapturingIdentifier                     #useCatalog
+    | USE namespace=multipartIdentifier
+        (IN catalog=errorCapturingIdentifier)?                         #use
     | CREATE database (IF NOT EXISTS)? db=errorCapturingIdentifier
         ((COMMENT comment=STRING) |
          locationSpec |
@@ -1399,6 +1401,7 @@ CACHE: 'CACHE';
 CASCADE: 'CASCADE';
 CASE: 'CASE';
 CAST: 'CAST';
+CATALOG: 'CATALOG';
 CHANGE: 'CHANGE';
 CHECK: 'CHECK';
 CLEAR: 'CLEAR';

sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/BaseSessionCatalog.java

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.sql.connector.catalog;
+
+import org.apache.spark.annotation.Experimental;
+
+/**
+ * An interface that aggregates different catalog interfaces that can be supported
+ * by a session catalog.
+ */
+@Experimental
+public interface BaseSessionCatalog extends TableCatalog, SupportsNamespaces {
+}

sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/CatalogExtension.java

@@ -22,17 +22,17 @@
 
 /**
  * An API to extend the Spark built-in session catalog. Implementation can get the built-in session
- * catalog from {@link #setDelegateCatalog(TableCatalog)}, implement catalog functions with
+ * catalog from {@link #setDelegateCatalog(BaseSessionCatalog)}, implement catalog functions with
  * some custom logic and call the built-in session catalog at the end. For example, they can
  * implement {@code createTable}, do something else before calling {@code createTable} of the
  * built-in session catalog.
  */
 @Experimental
-public interface CatalogExtension extends TableCatalog {
+public interface CatalogExtension extends BaseSessionCatalog {
 
   /**
    * This will be called only once by Spark to pass in the Spark built-in session catalog, after
    * {@link #initialize(String, CaseInsensitiveStringMap)} is called.
    */
-  void setDelegateCatalog(TableCatalog delegate);
+  void setDelegateCatalog(BaseSessionCatalog delegate);
 }

sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/DelegatingCatalogExtension.java

@@ -20,6 +20,7 @@
 import java.util.Map;
 
 import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.catalyst.analysis.NamespaceAlreadyExistsException;
 import org.apache.spark.sql.catalyst.analysis.NoSuchNamespaceException;
 import org.apache.spark.sql.catalyst.analysis.NoSuchTableException;
 import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException;
@@ -36,9 +37,9 @@
 @Experimental
 public abstract class DelegatingCatalogExtension implements CatalogExtension {
 
-  private TableCatalog delegate;
+  private BaseSessionCatalog delegate;
 
-  public final void setDelegateCatalog(TableCatalog delegate) {
+  public final void setDelegateCatalog(BaseSessionCatalog delegate) {
     this.delegate = delegate;
   }
 
@@ -97,4 +98,49 @@ public void renameTable(
       Identifier newIdent) throws NoSuchTableException, TableAlreadyExistsException {
     delegate.renameTable(oldIdent, newIdent);
   }
+
+  @Override
+  public String[] defaultNamespace() {
+    return delegate.defaultNamespace();
+  }
+
+  @Override
+  public String[][] listNamespaces() throws NoSuchNamespaceException {
+    return delegate.listNamespaces();
+  }
+
+  @Override
+  public String[][] listNamespaces(String[] namespace) throws NoSuchNamespaceException {
+    return delegate.listNamespaces(namespace);
+  }
+
+  @Override
+  public boolean namespaceExists(String[] namespace) {
+    return delegate.namespaceExists(namespace);
+  }
+
+  @Override
+  public Map<String, String> loadNamespaceMetadata(
+      String[] namespace) throws NoSuchNamespaceException {
+    return delegate.loadNamespaceMetadata(namespace);
+  }
+
+  @Override
+  public void createNamespace(
+      String[] namespace,
+      Map<String, String> metadata) throws NamespaceAlreadyExistsException {
+    delegate.createNamespace(namespace, metadata);
+  }
+
+  @Override
+  public void alterNamespace(
+      String[] namespace,
+      NamespaceChange... changes) throws NoSuchNamespaceException {
+    delegate.alterNamespace(namespace, changes);
+  }
+
+  @Override
+  public boolean dropNamespace(String[] namespace) throws NoSuchNamespaceException {
+    return delegate.dropNamespace(namespace);
+  }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -38,7 +38,7 @@ import org.apache.spark.sql.catalyst.plans.logical.sql._
 import org.apache.spark.sql.catalyst.rules._
 import org.apache.spark.sql.catalyst.trees.TreeNodeRef
 import org.apache.spark.sql.catalyst.util.toPrettySQL
-import org.apache.spark.sql.connector.catalog.{CatalogManager, CatalogPlugin, Identifier, LookupCatalog, Table, TableCatalog, TableChange, V1Table}
+import org.apache.spark.sql.connector.catalog.{BaseSessionCatalog, CatalogManager, CatalogPlugin, Identifier, LookupCatalog, NamespaceChange, Table, TableChange, V1Table}
 import org.apache.spark.sql.connector.expressions.{FieldReference, IdentityTransform, Transform}
 import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
 import org.apache.spark.sql.internal.SQLConf
@@ -60,7 +60,7 @@ object SimpleAnalyzer extends Analyzer(
   },
   new SQLConf().copy(SQLConf.CASE_SENSITIVE -> true))
 
-object FakeV2SessionCatalog extends TableCatalog {
+object FakeV2SessionCatalog extends BaseSessionCatalog {
   private def fail() = throw new UnsupportedOperationException
   override def listTables(namespace: Array[String]): Array[Identifier] = fail()
   override def loadTable(ident: Identifier): Table = {
@@ -75,7 +75,17 @@ object FakeV2SessionCatalog extends TableCatalog {
   override def dropTable(ident: Identifier): Boolean = fail()
   override def renameTable(oldIdent: Identifier, newIdent: Identifier): Unit = fail()
   override def initialize(name: String, options: CaseInsensitiveStringMap): Unit = fail()
-  override def name(): String = fail()
+  override def name(): String = "fake_v2_session"
+  override val defaultNamespace: Array[String] = Array("default")
+  override def namespaceExists(namespace: Array[String]): Boolean = fail()
+  override def listNamespaces(): Array[Array[String]] = fail()
+  override def listNamespaces(namespace: Array[String]): Array[Array[String]] = fail()
+  override def loadNamespaceMetadata(namespace: Array[String]): util.Map[String, String] = fail()
+  override def createNamespace(
+      namespace: Array[String],
+      metadata: util.Map[String, String]): Unit = fail()
+  override def alterNamespace(namespace: Array[String], changes: NamespaceChange*): Unit = fail()
+  override def dropNamespace(namespace: Array[String]): Boolean = fail()
 }
 
 /**
@@ -119,7 +129,7 @@ object AnalysisContext {
  */
 class Analyzer(
     catalog: SessionCatalog,
-    v2SessionCatalog: TableCatalog,
+    v2SessionCatalog: BaseSessionCatalog,
     conf: SQLConf,
     maxIterations: Int)
   extends RuleExecutor[LogicalPlan] with CheckAnalysis with LookupCatalog {
@@ -129,7 +139,7 @@ class Analyzer(
     this(catalog, FakeV2SessionCatalog, conf, conf.optimizerMaxIterations)
   }
 
-  def this(catalog: SessionCatalog, v2SessionCatalog: TableCatalog, conf: SQLConf) = {
+  def this(catalog: SessionCatalog, v2SessionCatalog: BaseSessionCatalog, conf: SQLConf) = {
     this(catalog, v2SessionCatalog, conf, conf.optimizerMaxIterations)
   }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -36,7 +36,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.{First, Last}
 import org.apache.spark.sql.catalyst.parser.SqlBaseParser._
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
-import org.apache.spark.sql.catalyst.plans.logical.sql.{AlterTableAddColumnsStatement, AlterTableAlterColumnStatement, AlterTableDropColumnsStatement, AlterTableRenameColumnStatement, AlterTableSetLocationStatement, AlterTableSetPropertiesStatement, AlterTableUnsetPropertiesStatement, AlterViewSetPropertiesStatement, AlterViewUnsetPropertiesStatement, CreateTableAsSelectStatement, CreateTableStatement, DeleteFromStatement, DescribeColumnStatement, DescribeTableStatement, DropTableStatement, DropViewStatement, InsertIntoStatement, QualifiedColType, ReplaceTableAsSelectStatement, ReplaceTableStatement, ShowNamespacesStatement, ShowTablesStatement}
+import org.apache.spark.sql.catalyst.plans.logical.sql.{AlterTableAddColumnsStatement, AlterTableAlterColumnStatement, AlterTableDropColumnsStatement, AlterTableRenameColumnStatement, AlterTableSetLocationStatement, AlterTableSetPropertiesStatement, AlterTableUnsetPropertiesStatement, AlterViewSetPropertiesStatement, AlterViewUnsetPropertiesStatement, CreateTableAsSelectStatement, CreateTableStatement, DeleteFromStatement, DescribeColumnStatement, DescribeTableStatement, DropTableStatement, DropViewStatement, InsertIntoStatement, QualifiedColType, ReplaceTableAsSelectStatement, ReplaceTableStatement, ShowNamespacesStatement, ShowTablesStatement, UseCatalogAndNamespaceStatement}
 import org.apache.spark.sql.catalyst.util.DateTimeUtils.{getZoneId, stringToDate, stringToTimestamp}
 import org.apache.spark.sql.connector.expressions.{ApplyTransform, BucketTransform, DaysTransform, Expression => V2Expression, FieldReference, HoursTransform, IdentityTransform, LiteralValue, MonthsTransform, Transform, YearsTransform}
 import org.apache.spark.sql.internal.SQLConf
@@ -2421,6 +2421,22 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
       ctx.EXISTS != null)
   }
 
+  /**
+   * Create a [[UseCatalogAndNamespaceStatement]] logical plan.
+   */
+  override def visitUseCatalog(ctx: UseCatalogContext): LogicalPlan = withOrigin(ctx) {
+    UseCatalogAndNamespaceStatement(Some(ctx.catalog.getText), None)
+  }
+
+  /**
+   * Create a [[UseCatalogAndNamespaceStatement]] logical plan.
+   */
+  override def visitUse(ctx: UseContext): LogicalPlan = withOrigin(ctx) {
+    UseCatalogAndNamespaceStatement(
+      Option(ctx.catalog).map(_.getText),
+      Some(visitMultipartIdentifier(ctx.namespace)))
+  }
+
   /**
    * Create a [[ShowTablesStatement]] command.
    */

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala

@@ -25,7 +25,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression,
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning, RangePartitioning, RoundRobinPartitioning}
 import org.apache.spark.sql.catalyst.util.truncatedString
-import org.apache.spark.sql.connector.catalog.{Identifier, SupportsNamespaces, TableCatalog, TableChange}
+import org.apache.spark.sql.connector.catalog.{CatalogManager, Identifier, SupportsNamespaces, TableCatalog, TableChange}
 import org.apache.spark.sql.connector.catalog.TableChange.{AddColumn, ColumnChange}
 import org.apache.spark.sql.connector.expressions.Transform
 import org.apache.spark.sql.types._
@@ -638,6 +638,14 @@ case class ShowTables(
     AttributeReference("tableName", StringType, nullable = false)())
 }
 
+/**
+ * The logical plan of the USE/USE CATALOG command that works for v2 catalogs.
+ */
+case class UseCatalogAndNamespace(
+    catalogManager: CatalogManager,
+    catalogName: Option[String],
+    namespace: Option[Seq[String]]) extends Command
+
 /**
  * Insert query result into a directory.
  *

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/sql/UseCatalogAndNamespaceStatement.scala

@@ -0,0 +1,28 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.spark.sql.catalyst.plans.logical.sql
+
+/**
+ * A USE/USE CATALOG statement, as parsed from SQL.
+ */
+case class UseCatalogAndNamespaceStatement(
+    catalogName: Option[String],
+    namespace: Option[Seq[String]])
+    extends ParsedStatement {
+  require(catalogName.nonEmpty || namespace.nonEmpty)
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/CatalogManager.scala

@@ -28,7 +28,7 @@ import org.apache.spark.sql.internal.SQLConf
  * the caller to look up a catalog by name.
  */
 private[sql]
-class CatalogManager(conf: SQLConf, defaultSessionCatalog: TableCatalog) extends Logging {
+class CatalogManager(conf: SQLConf, defaultSessionCatalog: BaseSessionCatalog) extends Logging {
 
   private val catalogs = mutable.HashMap.empty[String, CatalogPlugin]
 
@@ -84,34 +84,31 @@ class CatalogManager(conf: SQLConf, defaultSessionCatalog: TableCatalog) extends
   private var _currentNamespace: Option[Array[String]] = None
 
   def currentNamespace: Array[String] = synchronized {
-    _currentNamespace.getOrElse {
-      currentCatalog.map { catalogName =>
-        getDefaultNamespace(catalog(catalogName))
-      }.getOrElse(Array("default")) // The builtin catalog use "default" as the default database.
-    }
+    _currentNamespace.getOrElse(getDefaultNamespace(currentCatalog))
   }
 
   def setCurrentNamespace(namespace: Array[String]): Unit = synchronized {
     _currentNamespace = Some(namespace)
   }
 
-  private var _currentCatalog: Option[String] = None
+  private var _currentCatalogName: Option[String] = None
 
-  // Returns the name of current catalog. None means the current catalog is the builtin catalog.
-  def currentCatalog: Option[String] = synchronized {
-    _currentCatalog.orElse(conf.defaultV2Catalog)
+  def currentCatalog: CatalogPlugin = synchronized {
+    _currentCatalogName.map(catalogName => catalog(catalogName))
+      .orElse(defaultCatalog)
+      .getOrElse(v2SessionCatalog)
   }
 
   def setCurrentCatalog(catalogName: String): Unit = synchronized {
-    _currentCatalog = Some(catalogName)
+    _currentCatalogName = Some(catalogName)
     _currentNamespace = None
   }
 
   // Clear all the registered catalogs. Only used in tests.
   private[sql] def reset(): Unit = synchronized {
     catalogs.clear()
     _currentNamespace = None
-    _currentCatalog = None
+    _currentCatalogName = None
   }
 }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/connector/catalog/LookupCatalog.scala

@@ -38,6 +38,11 @@ private[sql] trait LookupCatalog extends Logging {
    */
   def defaultCatalog: Option[CatalogPlugin] = catalogManager.defaultCatalog
 
+  /**
+   * Returns the current catalog set.
+   */
+  def currentCatalog: CatalogPlugin = catalogManager.currentCatalog
+
   /**
    * This catalog is a v2 catalog that delegates to the v1 session catalog. it is used when the
    * session catalog is responsible for an identifier, but the source requires the v2 catalog API.
@@ -80,14 +85,14 @@ private[sql] trait LookupCatalog extends Logging {
     }
   }
 
-  type CatalogNamespace = (Option[CatalogPlugin], Seq[String])
+  type DefaultCatalogAndNamespace = (Option[CatalogPlugin], Seq[String])
 
   /**
    * Extract catalog and namespace from a multi-part identifier with the default catalog if needed.
    * Catalog name takes precedence over namespaces.
    */
-  object CatalogNamespace {
-    def unapply(parts: Seq[String]): Some[CatalogNamespace] = parts match {
+  object DefaultCatalogAndNamespace {
+    def unapply(parts: Seq[String]): Some[DefaultCatalogAndNamespace] = parts match {
       case Seq(catalogName, tail @ _*) =>
         try {
           Some((Some(catalogManager.catalog(catalogName)), tail))
@@ -98,6 +103,24 @@ private[sql] trait LookupCatalog extends Logging {
     }
   }
 
+  type CurrentCatalogAndNamespace = (CatalogPlugin, Seq[String])
+
+  /**
+   * Extract catalog and namespace from a multi-part identifier with the current catalog if needed.
+   * Catalog name takes precedence over namespaces.
+   */
+  object CurrentCatalogAndNamespace {
+    def unapply(parts: Seq[String]): Some[CurrentCatalogAndNamespace] = parts match {
+      case Seq(catalogName, tail @ _*) =>
+        try {
+          Some((catalogManager.catalog(catalogName), tail))
+        } catch {
+          case _: CatalogNotFoundException =>
+            Some((currentCatalog, parts))
+        }
+    }
+  }
+
   /**
    * Extract legacy table identifier from a multi-part identifier.
    *