only check non-virtual columns from unique key

go/logic/inspect.go

@@ -589,6 +589,7 @@ func (this *Inspector) applyColumnTypes(databaseName, tableName string, columnsL
 		columnName := m.GetString("COLUMN_NAME")
 		columnType := m.GetString("COLUMN_TYPE")
 		columnOctetLength := m.GetUint("CHARACTER_OCTET_LENGTH")
+		extra := m.GetString("EXTRA")
 		for _, columnsList := range columnsLists {
 			column := columnsList.GetColumn(columnName)
 			if column == nil {
@@ -621,6 +622,9 @@ func (this *Inspector) applyColumnTypes(databaseName, tableName string, columnsL
 				column.Type = sql.BinaryColumnType
 				column.BinaryOctetLength = columnOctetLength
 			}
+			if strings.Contains(extra, " GENERATED") {
+				column.IsVirtual = true
+			}
 			if charset := m.GetString("CHARACTER_SET_NAME"); charset != "" {
 				column.Charset = charset
 			}

go/sql/builder.go

@@ -546,6 +546,12 @@ func NewDMLUpdateQueryBuilder(databaseName, tableName string, tableColumns, shar
 	if uniqueKeyColumns.Len() == 0 {
 		return nil, fmt.Errorf("no unique key columns found in NewDMLUpdateQueryBuilder")
 	}
+	// If unique key contains virtual columns, those column won't be in sharedColumns
+	// which only contains non-vritual columns
+	nonVirtualUniqueKeyColumns := uniqueKeyColumns.FilterBy(func(column Column) bool { return !column.IsVirtual })
+	if !nonVirtualUniqueKeyColumns.IsSubsetOf(sharedColumns) {
+		return nil, fmt.Errorf("unique key columns  is not a subset of shared columns in DMLUpdateQueryBuilder")
+	}
 	databaseName = EscapeName(databaseName)
 	tableName = EscapeName(tableName)
 	setClause, err := BuildSetPreparedClause(mappedSharedColumns)
@@ -580,11 +586,6 @@ func NewDMLUpdateQueryBuilder(databaseName, tableName string, tableColumns, shar
 // BuildQuery builds the arguments array for a DML event UPDATE query.
 // It returns the query string, the shared arguments array, and the unique key arguments array.
 func (b *DMLUpdateQueryBuilder) BuildQuery(valueArgs, whereArgs []interface{}) (string, []interface{}, []interface{}, error) {
-	// TODO: move this check back to `NewDMLUpdateQueryBuilder()`, needs fix on generated columns.
-	if !b.uniqueKeyColumns.IsSubsetOf(b.sharedColumns) {
-		return "", nil, nil, fmt.Errorf("unique key columns is not a subset of shared columns in DMLUpdateQueryBuilder")
-	}
-
 	sharedArgs := make([]interface{}, 0, b.sharedColumns.Len())
 	for _, column := range b.sharedColumns.Columns() {
 		tableOrdinal := b.tableColumns.Ordinals[column.Name]

go/sql/types.go

@@ -40,6 +40,7 @@ type CharacterSetConversion struct {
 type Column struct {
 	Name                 string
 	IsUnsigned           bool
+	IsVirtual            bool
 	Charset              string
 	Type                 ColumnType
 	EnumValues           string
@@ -244,6 +245,18 @@ func (this *ColumnList) IsSubsetOf(other *ColumnList) bool {
 	return true
 }
 
+func (this *ColumnList) FilterBy(f func(Column) bool) *ColumnList {
+	filteredList := &ColumnList{}
+	filteredList.Ordinals = this.Ordinals
+	filteredCols := []Column{}
+	for _, column := range this.columns {
+		if f(column) {
+			filteredCols = append(filteredCols, column)
+		}
+	}
+	return filteredList
+}
+
 func (this *ColumnList) Len() int {
 	return len(this.columns)
 }