dialect/sql: added support for batch delete query

Author: LostImagin4tion

dialect/sql/builder.go

@@ -819,9 +819,10 @@ type DeleteBuilder struct {
 	schema string
 	where  *Predicate
 
-	// For YDB's DELETE FROM ... ON SELECT pattern
-	onSelect  *Selector
+	// YDB-specific:
+	onSelect  *Selector // DELETE FROM ... ON SELECT pattern
 	returning []string
+	isBatch   bool // use BATCH DELETE instead of DELETE
 }
 
 // Delete creates a builder for the `DELETE` statement.
@@ -839,6 +840,18 @@ type DeleteBuilder struct {
 //		)
 func Delete(table string) *DeleteBuilder { return &DeleteBuilder{table: table} }
 
+// BatchDelete creates a builder for the `BATCH DELETE FROM` statement (YDB-specific).
+// BATCH DELETE processes large tables in batches, minimizing lock invalidation risk.
+//
+// Note: BATCH DELETE is only supported in YDB dialect.
+//
+// BatchDelete("/local/my_table")
+//
+//	.Where(GT("Key1", 1))
+func BatchDelete(table string) *DeleteBuilder {
+	return &DeleteBuilder{table: table, isBatch: true}
+}
+
 // Schema sets the database name for the table whose row will be deleted.
 func (d *DeleteBuilder) Schema(name string) *DeleteBuilder {
 	d.schema = name
@@ -891,8 +904,32 @@ func (d *DeleteBuilder) Returning(columns ...string) *DeleteBuilder {
 
 // Query returns query representation of a `DELETE` statement.
 func (d *DeleteBuilder) Query() (string, []any) {
+	query, args, _ := d.QueryErr()
+	return query, args
+}
+
+func (d *DeleteBuilder) QueryErr() (string, []any, error) {
 	b := d.Builder.clone()
-	b.WriteString("DELETE FROM ")
+
+	// BATCH DELETE (YDB-specific)
+	if d.isBatch {
+		if !b.ydb() {
+			b.AddError(fmt.Errorf("BATCH DELETE: unsupported dialect: %q", b.dialect))
+			return "", nil, b.Err()
+		}
+		if len(d.returning) > 0 {
+			b.AddError(fmt.Errorf("BATCH DELETE: RETURNING clause is not supported"))
+			return "", nil, b.Err()
+		}
+		if d.onSelect != nil {
+			b.AddError(fmt.Errorf("BATCH DELETE: DELETE ON pattern is not supported"))
+			return "", nil, b.Err()
+		}
+		b.WriteString("BATCH DELETE FROM ")
+	} else {
+		b.WriteString("DELETE FROM ")
+	}
+
 	b.writeSchema(d.schema)
 	b.Ident(d.table)
 
@@ -907,7 +944,7 @@ func (d *DeleteBuilder) Query() (string, []any) {
 	}
 
 	joinReturning(d.returning, &b)
-	return b.String(), b.args
+	return b.String(), b.args, nil
 }
 
 // Predicate is a where predicate.
@@ -3772,6 +3809,18 @@ func (d *DialectBuilder) Delete(table string) *DeleteBuilder {
 	return b
 }
 
+// BatchDelete creates a DeleteBuilder for the BATCH DELETE statement with the configured dialect.
+// BATCH DELETE is only supported in YDB dialect.
+//
+//	Dialect(dialect.YDB).
+//		BatchDelete("users").
+//		Where(GT("Key1", 1))
+func (d *DialectBuilder) BatchDelete(table string) *DeleteBuilder {
+	b := BatchDelete(table)
+	b.SetDialect(d.dialect)
+	return b
+}
+
 // Select creates a Selector for the configured dialect.
 //
 //	Dialect(dialect.Postgres).

dialect/sql/builder_test.go

@@ -2297,40 +2297,67 @@ func TestBatchUpdate_YDB(t *testing.T) {
 		}, args)
 	})
 
-	t.Run("BATCH UPDATE without WHERE clause", func(t *testing.T) {
+	t.Run("BATCH UPDATE on non-YDB dialect should error", func(t *testing.T) {
+		builder := Dialect(dialect.MySQL).
+			BatchUpdate("users").
+			Set("status", "active").
+			Where(GT("created_at", "2024-01-01"))
+		
+		query, args, err := builder.QueryErr()
+		
+		require.Empty(t, query)
+		require.Empty(t, args)
+		require.Error(t, err)
+	})
+
+	t.Run("BATCH UPDATE with RETURNING should error", func(t *testing.T) {
+		builder := Dialect(dialect.YDB).
+			BatchUpdate("users").
+			Set("status", "active").
+			Returning("id", "status")
+		
+		query, args, err := builder.QueryErr()
+		
+		require.Empty(t, query)
+		require.Empty(t, args)
+		require.Error(t, err)
+	})
+
+	t.Run("BATCH UPDATE with UPDATE ON pattern should error", func(t *testing.T) {
 		d := Dialect(dialect.YDB)
-		query, args := d.BatchUpdate("users").
+		subquery := d.Select("id").From(Table("orders")).Where(EQ("status", "pending"))
+		
+		builder := d.BatchUpdate("users").
 			Set("status", "active").
-			Query()
+			On(subquery)
 
-		require.Equal(t, "BATCH UPDATE `users` SET `status` = $p0", query)
-		require.Equal(t, []any{
-			driver.NamedValue{Name: "p0", Value: "active"},
-		}, args)
+		query, args, err := builder.QueryErr()
+		
+		require.Empty(t, query)
+		require.Empty(t, args)
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "BATCH UPDATE: UPDATE ON pattern is not supported")
 	})
+}
 
-	t.Run("BATCH UPDATE with multiple SET clauses", func(t *testing.T) {
+func TestBatchDelete_YDB(t *testing.T) {
+	t.Run("Basic BATCH DELETE", func(t *testing.T) {
 		d := Dialect(dialect.YDB)
-		query, args := d.BatchUpdate("products").
-			Set("price", 100).
-			Set("stock", 50).
-			Set("updated_at", "2024-01-01").
-			Where(EQ("category", "electronics")).
+		query, args := d.BatchDelete("my_table").
+			Where(And(GT("Key1", 1), GTE("Key2", "One"))).
 			Query()
 
-		require.Contains(t, query, "BATCH UPDATE `products` SET")
-		require.Contains(t, query, "`price` = $p0")
-		require.Contains(t, query, "`stock` = $p1")
-		require.Contains(t, query, "`updated_at` = $p2")
-		require.Contains(t, query, "WHERE `category` = $p3")
-		require.Len(t, args, 4)
+		require.Equal(t, "BATCH DELETE FROM `my_table` WHERE `Key1` > $p0 AND `Key2` >= $p1", query)
+		require.Equal(t, []any{
+			driver.NamedValue{Name: "p0", Value: 1},
+			driver.NamedValue{Name: "p1", Value: "One"},
+		}, args)
 	})
 
-	t.Run("BATCH UPDATE on non-YDB dialect should error", func(t *testing.T) {
+	t.Run("BATCH DELETE on non-YDB dialect should error", func(t *testing.T) {
 		builder := Dialect(dialect.MySQL).
-			BatchUpdate("users").
-			Set("status", "active").
-			Where(GT("created_at", "2024-01-01"))
+			BatchDelete("users").
+			Where(GT("id", 100))
 
 		query, args, err := builder.QueryErr()
 
@@ -2339,17 +2366,32 @@ func TestBatchUpdate_YDB(t *testing.T) {
 		require.Error(t, err)
 	})
 
-	t.Run("BATCH UPDATE with RETURNING should error", func(t *testing.T) {
+	t.Run("BATCH DELETE with RETURNING should error", func(t *testing.T) {
 		builder := Dialect(dialect.YDB).
-			BatchUpdate("users").
-			Set("status", "active").
-			Returning("id", "status")
+			BatchDelete("users").
+			Where(GT("id", 100)).
+			Returning("id")
+		
+		query, args, err := builder.QueryErr()
+		
+		require.Empty(t, query)
+		require.Empty(t, args)
+		require.Error(t, err)
+	})
+
+	t.Run("BATCH DELETE with DELETE ON pattern should error", func(t *testing.T) {
+		d := Dialect(dialect.YDB)
+		subquery := d.Select("id").From(Table("users")).Where(EQ("status", "deleted"))
+		
+		builder := d.BatchDelete("users").
+			On(subquery)
 
 		query, args, err := builder.QueryErr()
 
 		require.Empty(t, query)
 		require.Empty(t, args)
 		require.Error(t, err)
+		require.Contains(t, err.Error(), "BATCH DELETE: DELETE ON pattern is not supported")
 	})
 }