Add AS OF SYSTEM TIME support for CockroachDB temporal queries

chainable_api.go

@@ -4,6 +4,7 @@
 	"fmt"
 	"regexp"
 	"strings"
+	"time"
 
 	"gorm.io/gorm/clause"
 	"gorm.io/gorm/utils"
@@ -85,6 +86,56 @@
 	return
 }
 
+// AsOfSystemTime sets the system time for CockroachDB temporal queries
+// This allows querying data as it existed at a specific point in time
+//
+//	// Query data as it existed 1 hour ago
+//	db.AsOfSystemTime("-1h").Find(&users)
+//	// Query data as it existed at a specific timestamp
+//	db.AsOfSystemTime(time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC)).Find(&users)
+//
+// Note: This feature is only supported by CockroachDB
+func (db *DB) AsOfSystemTime(timestamp interface{}) (tx *DB) {
+	tx = db.getInstance()
+
+	var asOfClause *clause.AsOfSystemTime
+
+	switch v := timestamp.(type) {
+	case string:
+		// Raw SQL expression like "-1h"
+		asOfClause = &clause.AsOfSystemTime{Raw: v}
+	case time.Time:
+		// Specific timestamp
+		asOfClause = &clause.AsOfSystemTime{Timestamp: v}
+	default:
+		tx.AddError(fmt.Errorf("unsupported timestamp type for AS OF SYSTEM TIME: %T", timestamp))
+		return tx
+	}
+
+	// Get or create the FROM clause
+	var fromClause clause.From
+	if v, ok := tx.Statement.Clauses["FROM"].Expression.(clause.From); ok {
+		fromClause = v
+	}
+
+	// Set the AS OF SYSTEM TIME
+	fromClause.AsOfSystemTime = asOfClause
+	tx.Statement.AddClause(fromClause)
+
+	return tx
+}
+
+// AsOfSystemTimeNow sets the system time to be as close to now as possible
+// This allows us to get the most recent data without the database engine attempting retries.
+//
+//	// Query data as it exists at the time of the query
+//	db.AsOfSystemTimeNow().Find(&users)
+//
+// Note: This feature is only supported by CockroachDB
+func (db *DB) AsOfSystemTimeNow() (tx *DB) {
+	return db.AsOfSystemTime("'-1µs'")
+}
+
 // Distinct specify distinct fields that you want querying
 //
 //	// Select distinct names of users

clause/as_of_system_time.go

@@ -0,0 +1,34 @@
+package clause
+
+import (
+	"fmt"
+	"strings"
+	"time"
+)
+
+// AsOfSystemTime represents CockroachDB's "AS OF SYSTEM TIME" clause
+// This allows querying data as it existed at a specific point in time
+type AsOfSystemTime struct {
+	Timestamp time.Time
+	Raw       string // For raw SQL expressions like "AS OF SYSTEM TIME '-1h'"
+}
+
+// Name returns the clause name
+func (a AsOfSystemTime) Name() string {
+	return "AS OF SYSTEM TIME"
+}
+
+// Build builds the "AS OF SYSTEM TIME" clause
+func (a AsOfSystemTime) Build(builder Builder) {
+	if a.Raw != "" {
+		sanitizedRaw := strings.ReplaceAll(a.Raw, ";", "")
+		builder.WriteString(fmt.Sprintf("AS OF SYSTEM TIME %s", sanitizedRaw))
+	} else if !a.Timestamp.IsZero() {
+		builder.WriteString(fmt.Sprintf("AS OF SYSTEM TIME '%s'", a.Timestamp.Format("2006-01-02 15:04:05.000000")))
+	}
+}
+
+// MergeClause merges the "AS OF SYSTEM TIME" clause
+func (a AsOfSystemTime) MergeClause(clause *Clause) {
+	clause.Expression = a
+}

clause/as_of_system_time_test.go

@@ -0,0 +1,83 @@
+package clause_test
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"gorm.io/gorm/clause"
+)
+
+func TestAsOfSystemTime(t *testing.T) {
+	results := []struct {
+		Clauses []clause.Interface
+		Result  string
+		Vars    []interface{}
+	}{
+		{
+			[]clause.Interface{
+				clause.Select{},
+				clause.From{
+					Tables:         []clause.Table{{Name: "users"}},
+					AsOfSystemTime: &clause.AsOfSystemTime{Timestamp: time.Date(2023, 1, 1, 12, 0, 0, 0, time.UTC)},
+				},
+			},
+			"SELECT * FROM `users` AS OF SYSTEM TIME '2023-01-01 12:00:00.000000'",
+			nil,
+		},
+		{
+			[]clause.Interface{
+				clause.Select{},
+				clause.From{
+					Tables:         []clause.Table{{Name: "users"}},
+					AsOfSystemTime: &clause.AsOfSystemTime{Raw: "'-1h'"},
+				},
+			},
+			"SELECT * FROM `users` AS OF SYSTEM TIME '-1h'",
+			nil,
+		},
+		{
+			[]clause.Interface{
+				clause.Select{},
+				clause.From{
+					Tables: []clause.Table{{Name: "users"}},
+					Joins: []clause.Join{
+						{
+							Type:  clause.InnerJoin,
+							Table: clause.Table{Name: "companies"},
+							ON: clause.Where{
+								[]clause.Expression{clause.Eq{clause.Column{Table: "companies", Name: "id"}, clause.Column{Table: "users", Name: "company_id"}}},
+							},
+						},
+					},
+					AsOfSystemTime: &clause.AsOfSystemTime{Raw: "'-1h'"},
+				},
+			},
+			"SELECT * FROM `users` INNER JOIN `companies` ON `companies`.`id` = `users`.`company_id` AS OF SYSTEM TIME '-1h'",
+			nil,
+		},
+	}
+
+	for idx, result := range results {
+		t.Run(fmt.Sprintf("case #%v", idx), func(t *testing.T) {
+			checkBuildClauses(t, result.Clauses, result.Result, result.Vars)
+		})
+	}
+}
+
+func TestAsOfSystemTimeName(t *testing.T) {
+	asOfClause := clause.AsOfSystemTime{}
+	if asOfClause.Name() != "AS OF SYSTEM TIME" {
+		t.Errorf("expected name 'AS OF SYSTEM TIME', got %q", asOfClause.Name())
+	}
+}
+
+func TestAsOfSystemTimeMergeClause(t *testing.T) {
+	asOfClause := clause.AsOfSystemTime{Timestamp: time.Now()}
+	var c clause.Clause
+	asOfClause.MergeClause(&c)
+
+	if c.Expression != asOfClause {
+		t.Error("expected expression to be set to the clause")
+	}
+}

clause/from.go

@@ -2,8 +2,9 @@
 
 // From from clause
 type From struct {
-	Tables []Table
-	Joins  []Join
+	Tables         []Table
+	Joins          []Join
+	AsOfSystemTime *AsOfSystemTime // CockroachDB specific: AS OF SYSTEM TIME
 }
 
 // Name from clause name
@@ -29,6 +30,12 @@
 		builder.WriteByte(' ')
 		join.Build(builder)
 	}
+
+	// Add AS OF SYSTEM TIME clause if specified
+	if from.AsOfSystemTime != nil {
+		builder.WriteByte(' ')
+		from.AsOfSystemTime.Build(builder)
+	}
 }
 
 // MergeClause merge from clause