Make Redirect Table Features Droppable

spark/src/main/scala/org/apache/spark/sql/delta/PreDowngradeTableFeatureCommand.scala

@@ -650,3 +650,45 @@ case class CheckpointProtectionPreDowngradeCommand(table: DeltaTableV2)
     false
   }
 }
+
+case class RedirectWriterOnlyPreDowngradeCommand(table: DeltaTableV2)
+  extends PreDowngradeTableFeatureCommand
+    with DeltaLogging {
+  /**
+   * We disable the feature by removing [[DeltaConfigs.REDIRECT_WRITER_ONLY]].
+   *
+   * @return True if the property is removed. False otherwise.
+   */
+  override def removeFeatureTracesIfNeeded(spark: SparkSession): Boolean = {
+    // Make sure feature data/metadata exist before proceeding.
+    if (RedirectWriterOnlyFeature.validateRemoval(table.initialSnapshot)) {
+      return false
+    }
+
+    val properties = Seq(DeltaConfigs.REDIRECT_WRITER_ONLY.key)
+    AlterTableUnsetPropertiesDeltaCommand(
+      table, properties, ifExists = false, fromDropFeatureCommand = true).run(spark)
+    true
+  }
+}
+
+case class RedirectReaderWriterPreDowngradeCommand(table: DeltaTableV2)
+  extends PreDowngradeTableFeatureCommand
+    with DeltaLogging {
+  /**
+   * We disable the feature by removing [[DeltaConfigs.REDIRECT_READER_WRITER]].
+   *
+   * @return True if the property is removed. False otherwise.
+   */
+  override def removeFeatureTracesIfNeeded(spark: SparkSession): Boolean = {
+    // Make sure feature data/metadata exist before proceeding.
+    if (RedirectReaderWriterFeature.validateRemoval(table.initialSnapshot)) {
+      return false
+    }
+
+    val properties = Seq(DeltaConfigs.REDIRECT_READER_WRITER.key)
+    AlterTableUnsetPropertiesDeltaCommand(
+      table, properties, ifExists = false, fromDropFeatureCommand = true).run(spark)
+    true
+  }
+}

spark/src/main/scala/org/apache/spark/sql/delta/TableFeature.scala

@@ -603,25 +603,58 @@ object TimestampNTZTableFeature extends ReaderWriterFeature(name = "timestampNtz
 
 object RedirectReaderWriterFeature
   extends ReaderWriterFeature(name = "redirectReaderWriter-preview")
-  with FeatureAutomaticallyEnabledByMetadata {
+  with FeatureAutomaticallyEnabledByMetadata with RemovableFeature {
   override def metadataRequiresFeatureToBeEnabled(
     protocol: Protocol,
     metadata: Metadata,
     spark: SparkSession
   ): Boolean = RedirectReaderWriter.isFeatureSet(metadata)
 
   override def automaticallyUpdateProtocolOfExistingTables: Boolean = true
+
+  override def preDowngradeCommand(table: DeltaTableV2): PreDowngradeTableFeatureCommand =
+    RedirectReaderWriterPreDowngradeCommand(table)
+
+  /**
+   * [[RedirectReaderWriterPreDowngradeCommand]] will try to remove
+   * [[DeltaConfigs.REDIRECT_READER_WRITER]],
+   * we check that here to make sure there is no concurrent txn that re-enables redirection.
+   */
+  override def validateRemoval(snapshot: Snapshot): Boolean =
+    !RedirectReaderWriter.isFeatureSet(snapshot.metadata)
+
+  // There is no action that is associated with this feature.
+  override def actionUsesFeature(action: Action): Boolean = false
+
+  // There is no action associated with this feature, so we don't need to truncate history to remove
+  // the traces of it. Note that the table properties for this feature will be left in the history
+  // but legacy clients who don't understand this feature will simply ignore them.
+  override def requiresHistoryProtection: Boolean = false
 }
 
 object RedirectWriterOnlyFeature extends WriterFeature(name = "redirectWriterOnly-preview")
-  with FeatureAutomaticallyEnabledByMetadata {
+  with FeatureAutomaticallyEnabledByMetadata with RemovableFeature {
   override def metadataRequiresFeatureToBeEnabled(
     protocol: Protocol,
     metadata: Metadata,
     spark: SparkSession
   ): Boolean = RedirectWriterOnly.isFeatureSet(metadata)
 
   override def automaticallyUpdateProtocolOfExistingTables: Boolean = true
+
+  override def preDowngradeCommand(table: DeltaTableV2): PreDowngradeTableFeatureCommand =
+    RedirectWriterOnlyPreDowngradeCommand(table)
+
+  /**
+   * [[RedirectWriterOnlyPreDowngradeCommand]] will try to remove
+   * [[DeltaConfigs.REDIRECT_WRITER_ONLY]],
+   * we check that here to make sure there is no concurrent txn that re-enables redirection.
+   */
+  override def validateRemoval(snapshot: Snapshot): Boolean =
+    !RedirectWriterOnly.isFeatureSet(snapshot.metadata)
+
+  // Writer features should directly return false, as it is only used for reader+writer features.
+  override def actionUsesFeature(action: Action): Boolean = false
 }
 
 trait BinaryVariantTableFeature {

spark/src/test/scala/org/apache/spark/sql/delta/DeltaProtocolVersionSuite.scala

@@ -32,6 +32,7 @@ import org.apache.spark.sql.delta.catalog.DeltaTableV2
 import org.apache.spark.sql.delta.commands.{AlterTableDropFeatureDeltaCommand, AlterTableSetPropertiesDeltaCommand, AlterTableUnsetPropertiesDeltaCommand}
 import org.apache.spark.sql.delta.commands.DeletionVectorUtils
 import org.apache.spark.sql.delta.coordinatedcommits._
+import org.apache.spark.sql.delta.redirect.{PathBasedRedirectSpec, RedirectReaderWriter, RedirectWriterOnly, TableRedirect}
 import org.apache.spark.sql.delta.sources.DeltaSQLConf
 import org.apache.spark.sql.delta.test.DeltaSQLCommandTest
 import org.apache.spark.sql.delta.test.DeltaTestImplicits._
@@ -4155,6 +4156,67 @@ trait DeltaProtocolVersionSuiteBase extends QueryTest
   }
   // ---- End Coordinated Commits Drop Feature Tests ----
 
+  private def testRedirectFeature(
+      redirectFeature: TableFeature,
+      tableRedirect: TableRedirect,
+      enableFastDrop: Boolean,
+      unsetTableProperty: Boolean): Unit = {
+    withSQLConf(DeltaSQLConf.FAST_DROP_FEATURE_ENABLED.key -> enableFastDrop.toString) {
+      test(s"drop ${redirectFeature.name} with fast drop - " +
+        s"enableFastDrop=$enableFastDrop, unsetTableProperty=$unsetTableProperty") {
+        withTempDir { dir =>
+          spark.sql(s"CREATE TABLE delta.`${dir.getCanonicalPath}` (id bigint) USING delta")
+          val deltaLog = DeltaLog.forTable(spark, dir)
+
+          val redirectSpec = new PathBasedRedirectSpec("targetPath")
+          tableRedirect.add(
+            deltaLog,
+            catalogTableOpt = None,
+            PathBasedRedirectSpec.REDIRECT_TYPE,
+            redirectSpec)
+
+          if (unsetTableProperty) {
+            sql(s"ALTER TABLE delta.`${dir.getCanonicalPath}` UNSET TBLPROPERTIES " +
+              s"('${tableRedirect.config.key}')")
+          }
+
+          val featureName = redirectFeature.name
+          // Both RedirectReaderWriterFeature and RedirectWriterOnlyFeature can be immediately
+          // dropped as they don't require history truncation. This is because there is no
+          // associated action with the features.
+          AlterTableDropFeatureDeltaCommand(DeltaTableV2(spark, deltaLog.dataPath), featureName)
+            .run(spark)
+
+          val snapshot = deltaLog.update()
+          // Writer feature is removed from the writer features set.
+          assert(!snapshot.protocol.writerFeatureNames.contains(featureName))
+          // Reader feature is removed from the reader features set.
+          assert(!snapshot.protocol.readerFeatureNames.contains(featureName))
+          assert(tableRedirect.config.fromMetaData(snapshot.metadata).isEmpty)
+
+          assertPropertiesAndShowTblProperties(deltaLog)
+
+          // Running the command again should throw an exception.
+          val e = intercept[DeltaTableFeatureException] {
+            AlterTableDropFeatureDeltaCommand(DeltaTableV2(spark, deltaLog.dataPath), featureName)
+              .run(spark)
+          }
+          assert(e.getErrorClass == "DELTA_FEATURE_DROP_FEATURE_NOT_PRESENT")
+        }
+      }
+    }
+  }
+
+  BOOLEAN_DOMAIN.foreach { unsetTableProperty =>
+    BOOLEAN_DOMAIN.foreach { enableFastDrop =>
+      // Test both writer-only and reader writer redirect feature.
+      testRedirectFeature(
+        RedirectWriterOnlyFeature, RedirectWriterOnly, enableFastDrop, unsetTableProperty)
+      testRedirectFeature(
+        RedirectReaderWriterFeature, RedirectReaderWriter, enableFastDrop, unsetTableProperty)
+    }
+  }
+
   // Create a table for testing that has an unsupported feature.
   private def withTestTableWithUnsupportedWriterFeature(
       emptyTable: Boolean)(testCode: String => Unit): Unit = {