[SPARK-32402][SQL] Implement ALTER TABLE in JDBC Table Catalog

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala

@@ -34,6 +34,7 @@ import org.apache.spark.sql.catalyst.encoders.RowEncoder
 import org.apache.spark.sql.catalyst.expressions.SpecificInternalRow
 import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
 import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, DateTimeUtils, GenericArrayData}
+import org.apache.spark.sql.connector.catalog.TableChange
 import org.apache.spark.sql.execution.datasources.jdbc.connection.ConnectionProvider
 import org.apache.spark.sql.jdbc.{JdbcDialect, JdbcDialects, JdbcType}
 import org.apache.spark.sql.types._
@@ -94,13 +95,7 @@ object JdbcUtils extends Logging {
    * Drops a table from the JDBC database.
    */
   def dropTable(conn: Connection, table: String, options: JDBCOptions): Unit = {
-    val statement = conn.createStatement
-    try {
-      statement.setQueryTimeout(options.queryTimeout)
-      statement.executeUpdate(s"DROP TABLE $table")
-    } finally {
-      statement.close()
-    }
+    executeStatement(conn, options, s"DROP TABLE $table")
   }
 
   /**
@@ -184,7 +179,7 @@ object JdbcUtils extends Logging {
     }
   }
 
-  private def getJdbcType(dt: DataType, dialect: JdbcDialect): JdbcType = {
+  private[sql] def getJdbcType(dt: DataType, dialect: JdbcDialect): JdbcType = {
     dialect.getJDBCType(dt).orElse(getCommonJDBCType(dt)).getOrElse(
       throw new IllegalArgumentException(s"Can't get JDBC type for ${dt.catalogString}"))
   }
@@ -882,13 +877,7 @@ object JdbcUtils extends Logging {
     // table_options or partition_options.
     // E.g., "CREATE TABLE t (name string) ENGINE=InnoDB DEFAULT CHARSET=utf8"
     val sql = s"CREATE TABLE $tableName ($strSchema) $createTableOptions"
-    val statement = conn.createStatement
-    try {
-      statement.setQueryTimeout(options.queryTimeout)
-      statement.executeUpdate(sql)
-    } finally {
-      statement.close()
-    }
+    executeStatement(conn, options, sql)
   }
 
   /**
@@ -900,10 +889,41 @@ object JdbcUtils extends Logging {
       newTable: String,
       options: JDBCOptions): Unit = {
     val dialect = JdbcDialects.get(options.url)
+    executeStatement(conn, options, dialect.renameTable(oldTable, newTable))
+  }
+
+  /**
+   * Update a table from the JDBC database.
+   */
+  def alterTable(
+      conn: Connection,
+      tableName: String,
+      changes: Seq[TableChange],
+      options: JDBCOptions): Unit = {
+    val dialect = JdbcDialects.get(options.url)
+    conn.setAutoCommit(false)
     val statement = conn.createStatement
     try {
       statement.setQueryTimeout(options.queryTimeout)
-      statement.executeUpdate(dialect.renameTable(oldTable, newTable))
+      for (sql <- dialect.alterTable(tableName, changes)) {
+        statement.executeUpdate(sql)
+      }
+      conn.commit()
+    } catch {
+      case e: SQLException =>
+        if (conn != null) conn.rollback()
+        throw e
+    } finally {
+      statement.close()
+      conn.setAutoCommit(true)
+    }
+  }
+
+  private def executeStatement(conn: Connection, options: JDBCOptions, sql: String): Unit = {
+    val statement = conn.createStatement
+    try {
+      statement.setQueryTimeout(options.queryTimeout)
+      statement.executeUpdate(sql)
     } finally {
       statement.close()
     }

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCTableCatalog.scala

@@ -129,11 +129,12 @@ class JDBCTableCatalog extends TableCatalog with Logging {
     JDBCTable(ident, schema, writeOptions)
   }
 
-  // TODO (SPARK-32402): Implement ALTER TABLE in JDBC Table Catalog
   override def alterTable(ident: Identifier, changes: TableChange*): Table = {
-    // scalastyle:off throwerror
-    throw new NotImplementedError()
-    // scalastyle:on throwerror
+    checkNamespace(ident.namespace())
+    withConnection { conn =>
+      JdbcUtils.alterTable(conn, getTableName(ident), changes, options)
+      loadTable(ident)
+    }
   }
 
   private def checkNamespace(namespace: Array[String]): Unit = {

sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala

@@ -19,9 +19,14 @@ package org.apache.spark.sql.jdbc
 
 import java.sql.{Connection, Date, Timestamp}
 
+import scala.collection.mutable.ArrayBuilder
+
 import org.apache.commons.lang3.StringUtils
 
 import org.apache.spark.annotation.{DeveloperApi, Since}
+import org.apache.spark.sql.connector.catalog.TableChange
+import org.apache.spark.sql.connector.catalog.TableChange._
+import org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils
 import org.apache.spark.sql.types._
 
 /**
@@ -184,15 +189,65 @@ abstract class JdbcDialect extends Serializable {
   /**
    * Rename an existing table.
    *
-   * TODO (SPARK-32382): Override this method in the dialects that don't support such syntax.
-   *
    * @param oldTable The existing table.
    * @param newTable New name of the table.
    * @return The SQL statement to use for renaming the table.
    */
   def renameTable(oldTable: String, newTable: String): String = {
     s"ALTER TABLE $oldTable RENAME TO $newTable"
   }
+
+  /**
+   * Alter an existing table.
+   *
+   * @param tableName The name of the table to be altered.
+   * @param changes Changes to apply to the table.
+   * @return The SQL statements to use for altering the table.
+   */
+  def alterTable(tableName: String, changes: Seq[TableChange]): Array[String] = {
+    val updateClause = ArrayBuilder.make[String]
+    for (change <- changes) {
+      change match {
+        case add: AddColumn =>
+          add.fieldNames match {
+            case Array(name) =>
+              val dataType = JdbcUtils.getJdbcType(add.dataType(), this).databaseTypeDefinition
+              updateClause += s"ALTER TABLE $tableName ADD COLUMN $name $dataType"
+          }
+        case rename: RenameColumn =>
+          rename.fieldNames match {
+            case Array(name) =>
+              updateClause += s"ALTER TABLE $tableName RENAME COLUMN $name TO ${rename.newName}"
+          }
+        case delete: DeleteColumn =>
+          delete.fieldNames match {
+            case Array(name) =>
+              updateClause += s"ALTER TABLE $tableName DROP COLUMN $name"
+          }
+        case update: UpdateColumnType =>
+          update.fieldNames match {
+            case Array(name) =>
+              val dataType = JdbcUtils.getJdbcType(update.newDataType(), this)
+                .databaseTypeDefinition
+              updateClause += s"ALTER TABLE $tableName ALTER COLUMN $name $dataType"
+          }
+        case update: UpdateColumnNullability =>
+          update.fieldNames match {
+            case Array(name) =>
+              if (update.nullable()) {
+                updateClause += s"ALTER TABLE $tableName ALTER COLUMN $name SET NULL"
+              } else {
+                updateClause += s"ALTER TABLE $tableName ALTER COLUMN $name SET NOT NULL"
+              }
+          }
+        // scalastyle:off throwerror
+        case _ => throw new NotImplementedError(s"JDBC alterTable has Unsupported" +
+          s" TableChange $change")
+        // scalastyle:on throwerror
+      }
+    }
+    updateClause.result()
+  }
 }
 
 /**

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/jdbc/JDBCTableCatalogSuite.scala

@@ -19,6 +19,8 @@ package org.apache.spark.sql.execution.datasources.v2.jdbc
 import java.sql.{Connection, DriverManager}
 import java.util.Properties
 
+import scala.collection.mutable.ArrayBuilder
+
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.{QueryTest, Row}
 import org.apache.spark.sql.test.SharedSparkSession
@@ -106,4 +108,70 @@ class JDBCTableCatalogSuite extends QueryTest with SharedSparkSession {
         Seq(Row("test", "people"), Row("test", "new_table")))
     }
   }
+
+  test("alter table ... add column") {
+    withTable("h2.test.alt_table") {
+      withConnection { conn =>
+        sql("CREATE TABLE h2.test.alt_table (ID INTEGER) USING _")
+      }
+      assert(checkColumnExistence("h2.test.alt_table", Array("ID")))
+      sql("ALTER TABLE h2.test.alt_table ADD COLUMNS (C1 INTEGER, C2 STRING)")
+      assert(checkColumnExistence("h2.test.alt_table", Array("ID", "C1", "C2")))
+      sql("ALTER TABLE h2.test.alt_table ADD COLUMNS (C3 DOUBLE)")
+      assert(checkColumnExistence("h2.test.alt_table", Array("ID", "C1", "C2", "C3")))
+    }
+  }
+
+  test("alter table ... rename column") {
+    withTable("h2.test.alt_table") {
+      withConnection { conn =>
+        sql("CREATE TABLE h2.test.alt_table (ID INTEGER) USING _")
+      }
+      assert(checkColumnExistence("h2.test.alt_table", Array("ID")))
+      sql("ALTER TABLE h2.test.alt_table RENAME COLUMN ID TO C")
+      assert(checkColumnExistence("h2.test.alt_table", Array("C")))
+    }
+  }
+
+  test("alter table ... drop column") {
+    withTable("h2.test.alt_table") {
+      withConnection { conn =>
+        sql("CREATE TABLE h2.test.alt_table (C1 INTEGER, C2 INTEGER) USING _")
+      }
+      assert(checkColumnExistence("h2.test.alt_table", Array("C1", "C2")))
+      sql("ALTER TABLE h2.test.alt_table DROP COLUMN C1")
+      assert(checkColumnExistence("h2.test.alt_table", Array("C2")))
+    }
+  }
+
+  test("alter table ... update column type") {
+    withTable("h2.test.alt_table") {
+      withConnection { conn =>
+        sql("CREATE TABLE h2.test.alt_table (ID INTEGER) USING _")
+      }
+      sql("ALTER TABLE h2.test.alt_table ALTER COLUMN id TYPE DOUBLE")
+      assert(sql(s"DESCRIBE TABLE h2.test.alt_table").select("data_type").first()
+        === Row("double"))
+    }
+  }
+
+  test("alter table ... update column comment not supported") {
+    withTable("h2.test.alt_table") {
+      withConnection { conn =>
+        sql("CREATE TABLE h2.test.alt_table (ID INTEGER) USING _")
+      }
+      val thrown = intercept[scala.NotImplementedError] {
+        sql("ALTER TABLE h2.test.alt_table ALTER COLUMN ID COMMENT 'test'")
+      }
+      assert(thrown.getMessage.contains("JDBC alterTable has Unsupported TableChange"))
+    }
+  }
+
+  private def checkColumnExistence(tableName: String, columns: Array[String]): Boolean = {
+    val rows = ArrayBuilder.make[Row]
+    for (column <- columns) {
+      rows += Row(column)
+    }
+    sql(s"DESCRIBE TABLE $tableName").select("col_name").take(columns.length) === rows.result()
+  }
 }