feat: add crdb consistency guarantees (#728)

crdbx/staleness.go

@@ -0,0 +1,105 @@
+// Copyright © 2023 Ory Corp
+// SPDX-License-Identifier: Apache-2.0
+
+package crdbx
+
+import (
+	"net/http"
+
+	"github.com/ory/x/dbal"
+
+	"github.com/gobuffalo/pop/v6"
+
+	"github.com/ory/x/sqlcon"
+)
+
+// Control API consistency guarantees
+//
+// swagger:model consistencyRequestParameters
+type ConsistencyRequestParameters struct {
+	// Read Consistency Level (experimental)
+	//
+	// The read consistency level determines the consistency guarantee for reads and queries:
+	//
+	// - strong (slow): The read is guaranteed to return the most recent data committed at the start of the read.
+	// - eventual (very fast): The result will return data that is about 4.8 seconds old.
+	//
+	// The default consistency guarantee can be changed in the Ory Network Console or using the Ory CLI with
+	// `ory patch project --replace '/serve/default_consistency_level="strong"'`.
+	//
+	// This feature is fully functional only in Ory Network and currently experimental.
+	//
+	// required: false
+	// in: query
+	Consistency ConsistencyLevel `json:"consistency"`
+}
+
+// ConsistencyLevel is the consistency level.
+// swagger:enum ConsistencyLevel
+type ConsistencyLevel string
+
+const (
+	// ConsistencyLevelUnset is the unset / default consistency level.
+	ConsistencyLevelUnset ConsistencyLevel = ""
+	// ConsistencyLevelStrong is the strong consistency level.
+	ConsistencyLevelStrong ConsistencyLevel = "strong"
+	// ConsistencyLevelEventual is the eventual consistency level using follower read timestamps.
+	ConsistencyLevelEventual ConsistencyLevel = "eventual"
+)
+
+// ConsistencyLevelFromRequest extracts the consistency level from a request.
+func ConsistencyLevelFromRequest(r *http.Request) ConsistencyLevel {
+	return ConsistencyLevelFromString(r.URL.Query().Get("consistency"))
+}
+
+// ConsistencyLevelFromString converts a string to a ConsistencyLevel.
+// If the string is not recognized or unset, ConsistencyLevelStrong is returned.
+func ConsistencyLevelFromString(in string) ConsistencyLevel {
+	switch in {
+	case string(ConsistencyLevelStrong):
+		return ConsistencyLevelStrong
+	case string(ConsistencyLevelEventual):
+		return ConsistencyLevelEventual
+	case string(ConsistencyLevelUnset):
+		return ConsistencyLevelStrong
+	}
+	return ConsistencyLevelStrong
+}
+
+// SetTransactionConsistency sets the transaction consistency level for CockroachDB.
+func SetTransactionConsistency(c *pop.Connection, level ConsistencyLevel, fallback ConsistencyLevel) error {
+	q := getTransactionConsistencyQuery(c.Dialect.Name(), level, fallback)
+	if len(q) == 0 {
+		return nil
+	}
+
+	return sqlcon.HandleError(c.RawQuery(q).Exec())
+}
+
+const transactionFollowerReadTimestamp = "SET TRANSACTION AS OF SYSTEM TIME follower_read_timestamp()"
+
+func getTransactionConsistencyQuery(dialect string, level ConsistencyLevel, fallback ConsistencyLevel) string {
+	if dialect != dbal.DriverCockroachDB {
+		// Only CockroachDB supports this.
+		return ""
+	}
+
+	switch level {
+	case ConsistencyLevelStrong:
+		// Nothing to do
+		return ""
+	case ConsistencyLevelEventual:
+		// Jumps to end of function
+	case ConsistencyLevelUnset:
+		fallthrough
+	default:
+		if fallback != ConsistencyLevelEventual {
+			// Nothing to do
+			return ""
+		}
+
+		// Jumps to end of function
+	}
+
+	return transactionFollowerReadTimestamp
+}

crdbx/staleness_test.go

@@ -0,0 +1,73 @@
+// Copyright © 2023 Ory Corp
+// SPDX-License-Identifier: Apache-2.0
+
+package crdbx
+
+import (
+	"fmt"
+	"net/http"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/ory/x/urlx"
+)
+
+func TestConsistencyLevelFromString(t *testing.T) {
+	assert.Equal(t, ConsistencyLevelStrong, ConsistencyLevelFromString(""))
+	assert.Equal(t, ConsistencyLevelStrong, ConsistencyLevelFromString("strong"))
+	assert.Equal(t, ConsistencyLevelEventual, ConsistencyLevelFromString("eventual"))
+	assert.Equal(t, ConsistencyLevelStrong, ConsistencyLevelFromString("lol"))
+}
+
+func TestConsistencyLevelFromRequest(t *testing.T) {
+	assert.Equal(t, ConsistencyLevelStrong, ConsistencyLevelFromRequest(&http.Request{URL: urlx.ParseOrPanic("/?consistency=strong")}))
+	assert.Equal(t, ConsistencyLevelEventual, ConsistencyLevelFromRequest(&http.Request{URL: urlx.ParseOrPanic("/?consistency=eventual")}))
+	assert.Equal(t, ConsistencyLevelStrong, ConsistencyLevelFromRequest(&http.Request{URL: urlx.ParseOrPanic("/?consistency=asdf")}))
+
+}
+
+func TestGetTransactionConsistency(t *testing.T) {
+	for k, tc := range []struct {
+		in       ConsistencyLevel
+		fallback ConsistencyLevel
+		dialect  string
+		expected string
+	}{
+		{
+			in:       ConsistencyLevelUnset,
+			fallback: ConsistencyLevelStrong,
+			dialect:  "cockroach",
+			expected: "",
+		},
+		{
+			in:       ConsistencyLevelStrong,
+			fallback: ConsistencyLevelStrong,
+			dialect:  "cockroach",
+			expected: "",
+		},
+		{
+			in:       ConsistencyLevelStrong,
+			fallback: ConsistencyLevelEventual,
+			dialect:  "cockroach",
+			expected: "",
+		},
+		{
+			in:       ConsistencyLevelUnset,
+			fallback: ConsistencyLevelEventual,
+			dialect:  "cockroach",
+			expected: transactionFollowerReadTimestamp,
+		},
+		{
+			in:       ConsistencyLevelEventual,
+			fallback: ConsistencyLevelEventual,
+			dialect:  "cockroach",
+			expected: transactionFollowerReadTimestamp,
+		},
+	} {
+		t.Run(fmt.Sprintf("case=%d", k), func(t *testing.T) {
+			q := getTransactionConsistencyQuery(tc.dialect, tc.in, tc.fallback)
+			assert.EqualValues(t, tc.expected, q)
+		})
+	}
+}