[WIP] Table constraint

common/utils/src/main/resources/error/error-conditions.json

@@ -788,6 +788,20 @@
     },
     "sqlState" : "XX000"
   },
+  "CONSTRAINT_ALREADY_EXISTS" : {
+    "message" : [
+      "Constraint '<constraintName>' already exists. Please delete the old constraint first.",
+      "Old constraint:",
+      "<oldConstraint>"
+    ],
+    "sqlState" : "42710"
+  },
+  "CONSTRAINT_DOES_NOT_EXIST" : {
+    "message" : [
+      "Cannot drop nonexistent constraint <constraintName> from table <tableName>."
+    ],
+    "sqlState" : "42704"
+  },
   "CONVERSION_INVALID_INPUT" : {
     "message" : [
       "The value <str> (<fmt>) cannot be converted to <targetType> because it is malformed. Correct the value as per the syntax, or change its format. Use <suggestion> to tolerate malformed input and return NULL instead."
@@ -2303,6 +2317,29 @@
     ],
     "sqlState" : "22P03"
   },
+  "INVALID_CHECK_CONSTRAINT": {
+    "message" : [
+      "The check constraint expression is invalid."
+    ],
+    "subClass" : {
+      "INVALID_V2_PREDICATE" : {
+        "message" : [
+          "It cannot be converted to a data source V2 predicate."
+        ]
+      },
+      "NONDETERMINISTIC" : {
+        "message" : [
+          "It contains nondeterministic expression."
+        ]
+      },
+      "MISSING_NAME": {
+        "message": [
+          "The check constraint must have a name."
+        ]
+      }
+    },
+    "sqlState": "42621"
+  },
   "INVALID_COLUMN_NAME_AS_PATH" : {
     "message" : [
       "The datasource <datasource> cannot save the column <columnName> because its name contains some characters that are not allowed in file paths. Please, use an alias to rename it."
@@ -2328,6 +2365,12 @@
     },
     "sqlState" : "22022"
   },
+  "INVALID_CONSTRAINT_CHARACTERISTICS": {
+    "message": [
+      "Constraint characteristics [<characteristics>] are duplicated or conflict with each other."
+    ],
+    "sqlState": "42613"
+  },
   "INVALID_CORRUPT_RECORD_TYPE" : {
     "message" : [
       "The column <columnName> for corrupt records must have the nullable STRING type, but got <actualType>."
@@ -3884,6 +3927,12 @@
     ],
     "sqlState" : "42P20"
   },
+  "MULTIPLE_PRIMARY_KEYS" : {
+    "message" : [
+      "Multiple primary keys are defined. Please ensure that only one primary key is defined for the table."
+    ],
+    "sqlState" : "42K0E"
+  },
   "MULTIPLE_QUERY_RESULT_CLAUSES_WITH_PIPE_OPERATORS" : {
     "message" : [
       "<clause1> and <clause2> cannot coexist in the same SQL pipe operator using '|>'. Please separate the multiple result clauses into separate pipe operators and then retry the query again."
@@ -5452,6 +5501,12 @@
     },
     "sqlState" : "0A000"
   },
+  "UNSUPPORTED_CONSTRAINT_CHARACTERISTIC": {
+    "message": [
+      "Constraint characteristic '<characteristic>' is not supported for constraint type '<constraintType>'."
+    ],
+    "sqlState": "0A000"
+  },
   "UNSUPPORTED_DATASOURCE_FOR_DIRECT_QUERY" : {
     "message" : [
       "Unsupported data source type for direct query on files: <dataSourceType>"
@@ -5611,6 +5666,11 @@
           "Attach a comment to the namespace <namespace>."
         ]
       },
+      "CONSTRAINT_TYPE" : {
+        "message" : [
+          "Constraint <constraint>."
+        ]
+      },
       "CONTINUE_EXCEPTION_HANDLER" : {
         "message" : [
           "CONTINUE exception handler is not supported. Use EXIT handler."

core/src/test/scala/org/apache/spark/SparkFunSuite.scala

@@ -371,10 +371,17 @@ abstract class SparkFunSuite
           "Invalid objectType of a query context Actual:" + actual.toString)
         assert(actual.objectName() === expected.objectName,
           "Invalid objectName of a query context. Actual:" + actual.toString)
-        assert(actual.startIndex() === expected.startIndex,
-          "Invalid startIndex of a query context. Actual:" + actual.toString)
-        assert(actual.stopIndex() === expected.stopIndex,
-          "Invalid stopIndex of a query context. Actual:" + actual.toString)
+        // If startIndex and stopIndex are -1, it means we simply want to check the
+        // fragment of the query context. This should be the case when the fragment is
+        // distinguished within the query text.
+        if (expected.startIndex != -1) {
+          assert(actual.startIndex() === expected.startIndex,
+            "Invalid startIndex of a query context. Actual:" + actual.toString)
+        }
+        if (expected.stopIndex != -1) {
+          assert(actual.stopIndex() === expected.stopIndex,
+            "Invalid stopIndex of a query context. Actual:" + actual.toString)
+        }
         assert(actual.fragment() === expected.fragment,
           "Invalid fragment of a query context. Actual:" + actual.toString)
       } else if (actual.contextType() == QueryContextType.DataFrame) {
@@ -478,6 +485,12 @@ abstract class SparkFunSuite
       ExpectedContext("", "", start, stop, fragment)
     }
 
+    // Check the fragment only. This is only used when the fragment is distinguished within
+    // the query text
+    def apply(fragment: String): ExpectedContext = {
+      ExpectedContext("", "", -1, -1, fragment)
+    }
+
     def apply(
         objectType: String,
         objectName: String,

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

@@ -222,6 +222,7 @@ DROP: 'DROP';
 ELSE: 'ELSE';
 ELSEIF: 'ELSEIF';
 END: 'END';
+ENFORCED: 'ENFORCED';
 ESCAPE: 'ESCAPE';
 ESCAPED: 'ESCAPED';
 EVOLUTION: 'EVOLUTION';
@@ -290,6 +291,7 @@ ITEMS: 'ITEMS';
 ITERATE: 'ITERATE';
 JOIN: 'JOIN';
 JSON: 'JSON';
+KEY: 'KEY';
 KEYS: 'KEYS';
 LANGUAGE: 'LANGUAGE';
 LAST: 'LAST';
@@ -337,6 +339,8 @@ NOT: 'NOT';
 NULL: 'NULL';
 NULLS: 'NULLS';
 NUMERIC: 'NUMERIC';
+NORELY: 'NORELY';
+NOVALIDATE: 'NOVALIDATE';
 OF: 'OF';
 OFFSET: 'OFFSET';
 ON: 'ON';
@@ -376,6 +380,7 @@ RECURSIVE: 'RECURSIVE';
 REDUCE: 'REDUCE';
 REFERENCES: 'REFERENCES';
 REFRESH: 'REFRESH';
+RELY: 'RELY';
 RENAME: 'RENAME';
 REPAIR: 'REPAIR';
 REPEAT: 'REPEAT';
@@ -475,6 +480,7 @@ UPDATE: 'UPDATE';
 USE: 'USE';
 USER: 'USER';
 USING: 'USING';
+VALIDATE: 'VALIDATE';
 VALUE: 'VALUE';
 VALUES: 'VALUES';
 VARCHAR: 'VARCHAR';

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

@@ -198,7 +198,7 @@ statement
         (RESTRICT | CASCADE)?                                          #dropNamespace
     | SHOW namespaces ((FROM | IN) multipartIdentifier)?
         (LIKE? pattern=stringLit)?                                        #showNamespaces
-    | createTableHeader (LEFT_PAREN colDefinitionList RIGHT_PAREN)? tableProvider?
+    | createTableHeader (LEFT_PAREN tableElementList RIGHT_PAREN)? tableProvider?
         createTableClauses
         (AS? query)?                                                   #createTable
     | CREATE TABLE (IF errorCapturingNot EXISTS)? target=tableIdentifier
@@ -208,7 +208,7 @@ statement
         createFileFormat |
         locationSpec |
         (TBLPROPERTIES tableProps=propertyList))*                      #createTableLike
-    | replaceTableHeader (LEFT_PAREN colDefinitionList RIGHT_PAREN)? tableProvider?
+    | replaceTableHeader (LEFT_PAREN tableElementList RIGHT_PAREN)? tableProvider?
         createTableClauses
         (AS? query)?                                                   #replaceTable
     | ANALYZE TABLE identifierReference partitionSpec? COMPUTE STATISTICS
@@ -261,6 +261,10 @@ statement
     | ALTER TABLE identifierReference
         (clusterBySpec | CLUSTER BY NONE)                              #alterClusterBy
     | ALTER TABLE identifierReference collationSpec                    #alterTableCollation
+    | ALTER TABLE identifierReference ADD tableConstraintDefinition    #addTableConstraint
+    | ALTER TABLE identifierReference
+        DROP CONSTRAINT (IF EXISTS)? name=identifier
+        (RESTRICT | CASCADE)?                                     #dropTableConstraint
     | DROP TABLE (IF EXISTS)? identifierReference PURGE?               #dropTable
     | DROP VIEW (IF EXISTS)? identifierReference                       #dropView
     | CREATE (OR REPLACE)? (GLOBAL? TEMPORARY)?
@@ -1334,6 +1338,15 @@ colType
     : colName=errorCapturingIdentifier dataType (errorCapturingNot NULL)? commentSpec?
     ;
 
+tableElementList
+    : tableElement (COMMA tableElement)*
+    ;
+
+tableElement
+    : tableConstraintDefinition
+    | colDefinition
+    ;
+
 colDefinitionList
     : colDefinition (COMMA colDefinition)*
     ;
@@ -1347,6 +1360,7 @@ colDefinitionOption
     | defaultExpression
     | generationExpression
     | commentSpec
+    | columnConstraintDefinition
     ;
 
 generationExpression
@@ -1516,6 +1530,62 @@ number
     | MINUS? BIGDECIMAL_LITERAL       #bigDecimalLiteral
     ;
 
+columnConstraintDefinition
+    : (CONSTRAINT name=errorCapturingIdentifier)? columnConstraint constraintCharacteristic*
+    ;
+
+columnConstraint
+    : checkConstraint
+    | uniqueSpec
+    | referenceSpec
+    ;
+
+tableConstraintDefinition
+    : (CONSTRAINT name=errorCapturingIdentifier)? tableConstraint constraintCharacteristic*
+    ;
+
+tableConstraint
+    : checkConstraint
+    | uniqueConstraint
+    | foreignKeyConstraint
+    ;
+
+checkConstraint
+    : CHECK LEFT_PAREN (expr=booleanExpression) RIGHT_PAREN
+    ;
+
+uniqueSpec
+    : UNIQUE
+    | PRIMARY KEY
+    ;
+
+uniqueConstraint
+    : uniqueSpec identifierList
+    ;
+
+referenceSpec
+    : REFERENCES multipartIdentifier (parentColumns=identifierList)?
+    ;
+
+foreignKeyConstraint
+    : FOREIGN KEY identifierList referenceSpec
+    ;
+
+constraintCharacteristic
+    : enforcedCharacteristic
+    | relyCharacteristic
+    ;
+
+enforcedCharacteristic
+    : ENFORCED
+    | NOT ENFORCED
+    ;
+
+relyCharacteristic
+    : RELY
+    | NORELY
+    ;
+
 alterColumnSpecList
     : alterColumnSpec (COMMA alterColumnSpec)*
     ;
@@ -1673,6 +1743,7 @@ ansiNonReserved
     | DOUBLE
     | DROP
     | ELSEIF
+    | ENFORCED
     | ESCAPED
     | EVOLUTION
     | EXCHANGE
@@ -1761,6 +1832,8 @@ ansiNonReserved
     | NANOSECONDS
     | NO
     | NONE
+    | NORELY
+    | NOVALIDATE
     | NULLS
     | NUMERIC
     | OF
@@ -1792,6 +1865,7 @@ ansiNonReserved
     | RECOVER
     | REDUCE
     | REFRESH
+    | RELY
     | RENAME
     | REPAIR
     | REPEAT
@@ -1875,6 +1949,7 @@ ansiNonReserved
     | UNTIL
     | UPDATE
     | USE
+    | VALIDATE
     | VALUE
     | VALUES
     | VARCHAR

sql/api/src/main/scala/org/apache/spark/sql/errors/QueryParsingErrors.scala

@@ -791,4 +791,24 @@ private[sql] object QueryParsingErrors extends DataTypeErrorsBase {
   def clusterByWithBucketing(ctx: ParserRuleContext): Throwable = {
     new ParseException(errorClass = "SPECIFY_CLUSTER_BY_WITH_BUCKETING_IS_NOT_ALLOWED", ctx)
   }
+
+  def invalidConstraintCharacteristics(
+      ctx: ParserRuleContext,
+      characteristics: String): Throwable = {
+    new ParseException(
+      errorClass = "INVALID_CONSTRAINT_CHARACTERISTICS",
+      messageParameters = Map("characteristics" -> characteristics),
+      ctx)
+  }
+
+  def constraintNotSupportedError(ctx: ParserRuleContext, constraint: String): Throwable = {
+    new ParseException(
+      errorClass = "UNSUPPORTED_FEATURE.CONSTRAINT_TYPE",
+      messageParameters = Map("constraint" -> constraint),
+      ctx)
+  }
+
+  def multiplePrimaryKeysError(ctx: ParserRuleContext): Throwable = {
+    new ParseException(errorClass = "MULTIPLE_PRIMARY_KEYS", ctx)
+  }
 }

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

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.connector.catalog;
 
 import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.connector.catalog.constraints.Constraint;
 import org.apache.spark.sql.connector.expressions.Transform;
 import org.apache.spark.sql.types.StructType;
 
@@ -83,4 +84,9 @@ default Map<String, String> properties() {
    * Returns the set of capabilities for this table.
    */
   Set<TableCapability> capabilities();
+
+  /**
+   * Returns the constraints for this table.
+   */
+  default Constraint[] constraints() { return new Constraint[0]; }
 }