Postgres backend (#29025)

* Postgres backend * Use batch to mark read transactions as read only * Allow the backend_version creation to fail * clarify setupAndMigrate * simplify AzureBeforeConnect * Allow idempotent retries for reads * rm go.sum && go mod tidy * clarify error messages * clarify TELEPORT_PGBK_TEST_PARAMS_JSON * Describe why we always use the real clock * move SetupAndMigrate to common * update revision comment * clarify batches * actually use the context in conn.Close * give a reason why we silence logs in the change feed * order expiration deletes by expiry time * reduce defaultChangeFeedBatchSize down to 1000 we seem to be handling outrageous load fine even with 1000, and 10k (if actually hit) all but guarantees that we go over the default work_mem, forcing postgres to spill to disk; it's a tunable, anyway * parse expires as timestamptz in pollChangeFeed * implement GetName * blindly try to grant REPLICATION * default_transaction_isolation in RuntimeParams * Always use NewDefaultAzureCredential * retry unique_violation and exclusion_violation * fix pgevents.TestConfig * extract isSerializationErrorCode * polish the default_transaction_isolation check * Fix typo in Clock doc comment Co-authored-by: Zac Bergquist <[email protected]> * expires is a Timestamptz --------- Co-authored-by: Zac Bergquist <[email protected]>
2026-10-11 22:49:54 +00:00 · 2023-07-26 15:44:01 +00:00
parent f89d645ac3
commit 11b38ef117
12 changed files with 946 additions and 63 deletions
@@ -0,0 +1,276 @@
+// Copyright 2023 Gravitational, Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pgbk
+
+import (
+	"context"
+	"encoding/hex"
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/gravitational/trace"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgtype/zeronull"
+	"github.com/sirupsen/logrus"
+
+	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/lib/backend"
+	pgcommon "github.com/gravitational/teleport/lib/backend/pgbk/common"
+	"github.com/gravitational/teleport/lib/defaults"
+)
+
+func (b *Backend) backgroundExpiry(ctx context.Context) {
+	defer b.log.Info("Exited expiry loop.")
+
+	for ctx.Err() == nil {
+		// "DELETE FROM kv WHERE expires <= now()" but more complicated: logical
+		// decoding can become really really slow if a transaction is big enough
+		// to spill on disk - max_changes_in_memory (4096) changes before
+		// Postgres 13, or logical_decoding_work_mem (64MiB) bytes of total size
+		// in Postgres 13 and later; thankfully, we can just limit our
+		// transactions to a small-ish number of affected rows (1000 seems to
+		// work ok) as we don't need atomicity for this; we run a tight loop
+		// here because it could be possible to have more than ExpiryBatchSize
+		// new items expire every ExpiryInterval, so we could end up not ever
+		// catching up
+		for i := 0; i < backend.DefaultRangeLimit/b.cfg.ExpiryBatchSize; i++ {
+			t0 := time.Now()
+			// TODO(espadolini): try getting keys in a read-only deferrable
+			// transaction and deleting them later to reduce potential
+			// serialization issues
+			deleted, err := pgcommon.RetryIdempotent(ctx, b.log, func() (int64, error) {
+				// LIMIT without ORDER BY might get executed poorly because the
+				// planner doesn't have any idea of how many rows will be chosen
+				// or skipped, and it's not necessary but it's a nice touch that
+				// we'll be deleting expired items in expiration order
+				tag, err := b.pool.Exec(ctx,
+					"DELETE FROM kv WHERE key IN (SELECT key FROM kv"+
+						" WHERE expires IS NOT NULL AND expires <= now()"+
+						" ORDER BY expires LIMIT $1 FOR UPDATE)",
+					b.cfg.ExpiryBatchSize,
+				)
+				if err != nil {
+					return 0, trace.Wrap(err)
+				}
+				return tag.RowsAffected(), nil
+			})
+			if err != nil {
+				b.log.WithError(err).Error("Failed to delete expired items.")
+				break
+			}
+
+			if deleted > 0 {
+				b.log.WithFields(logrus.Fields{
+					"deleted": deleted,
+					"elapsed": time.Since(t0).String(),
+				}).Debug("Deleted expired items.")
+			}
+
+			if deleted < int64(b.cfg.ExpiryBatchSize) {
+				break
+			}
+		}
+
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(time.Duration(b.cfg.ExpiryInterval)):
+		}
+	}
+}
+
+func (b *Backend) backgroundChangeFeed(ctx context.Context) {
+	defer b.log.Info("Exited change feed loop.")
+	defer b.buf.Close()
+
+	for ctx.Err() == nil {
+		b.log.Info("Starting change feed stream.")
+		err := b.runChangeFeed(ctx)
+		if ctx.Err() != nil {
+			break
+		}
+		b.log.WithError(err).Error("Change feed stream lost.")
+
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(defaults.HighResPollingPeriod):
+		}
+	}
+}
+
+// runChangeFeed will connect to the database, start a change feed and emit
+// events. Assumes that b.buf is not initialized but not closed, and will reset
+// it before returning.
+func (b *Backend) runChangeFeed(ctx context.Context) error {
+	// we manually copy the pool configuration and connect because we don't want
+	// to hit a connection limit or mess with the connection pool stats; we need
+	// a separate, long-running connection here anyway.
+	poolConfig := b.pool.Config()
+	if poolConfig.BeforeConnect != nil {
+		if err := poolConfig.BeforeConnect(ctx, poolConfig.ConnConfig); err != nil {
+			return trace.Wrap(err)
+		}
+	}
+	conn, err := pgx.ConnectConfig(ctx, poolConfig.ConnConfig)
+	if err != nil {
+		return trace.Wrap(err)
+	}
+	defer func() {
+		ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
+		defer cancel()
+		if err := conn.Close(ctx); err != nil && ctx.Err() != nil {
+			b.log.WithError(err).Warn("Error closing change feed connection.")
+		}
+	}()
+
+	// reading from a replication slot adds to the postgres log at "log" level
+	// (right below "fatal") for every poll, and we poll every second here, so
+	// we try to silence the logs for this connection; this can fail because of
+	// permission issues, which would delete the temporary slot (it's deleted on
+	// any error), so we have to do it before that
+	if _, err := conn.Exec(ctx, "SET log_min_messages TO fatal", pgx.QueryExecModeExec); err != nil {
+		b.log.WithError(err).Debug("Failed to silence log messages for change feed session.")
+	}
+
+	// this can be useful if we're some sort of admin but we haven't gotten the
+	// REPLICATION attribute yet
+	// HACK(espadolini): ALTER ROLE CURRENT_USER REPLICATION just crashes postgres on Azure
+	if _, err := conn.Exec(ctx,
+		fmt.Sprintf("ALTER ROLE \"%v\" REPLICATION", poolConfig.ConnConfig.User),
+		pgx.QueryExecModeExec,
+	); err != nil {
+		b.log.WithError(err).Debug("Failed to enable replication for the current user.")
+	}
+
+	u := uuid.New()
+	slotName := hex.EncodeToString(u[:])
+
+	b.log.WithField("slot_name", slotName).Info("Setting up change feed.")
+	if _, err := conn.Exec(ctx,
+		"SELECT * FROM pg_create_logical_replication_slot($1, 'wal2json', true)",
+		pgx.QueryExecModeExec, slotName,
+	); err != nil {
+		return trace.Wrap(err)
+	}
+
+	b.log.WithField("slot_name", slotName).Info("Change feed started.")
+	b.buf.SetInit()
+	defer b.buf.Reset()
+
+	for ctx.Err() == nil {
+		events, err := b.pollChangeFeed(ctx, conn, slotName)
+		if err != nil {
+			return trace.Wrap(err)
+		}
+
+		// tight loop if we hit the batch size
+		if events >= int64(b.cfg.ChangeFeedBatchSize) {
+			continue
+		}
+
+		select {
+		case <-ctx.Done():
+			return trace.Wrap(ctx.Err())
+		case <-time.After(time.Duration(b.cfg.ChangeFeedPollInterval)):
+		}
+	}
+	return trace.Wrap(err)
+}
+
+// pollChangeFeed will poll the change feed and emit any fetched events, if any.
+// It returns the count of received/emitted events.
+func (b *Backend) pollChangeFeed(ctx context.Context, conn *pgx.Conn, slotName string) (int64, error) {
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	defer cancel()
+
+	t0 := time.Now()
+
+	// TODO(espadolini): it might be better to do the JSON deserialization
+	// (potentially with additional checks for the schema) on the auth side
+	rows, _ := conn.Query(ctx,
+		`WITH jdata AS (
+  SELECT
+    data::jsonb AS data
+  FROM pg_logical_slot_get_changes($1, NULL, $2,
+    'format-version', '2', 'add-tables', 'public.kv', 'include-transaction', 'false')
+)
+SELECT
+  data->>'action',
+  decode(COALESCE(data->'columns'->0->>'value', data->'identity'->0->>'value'), 'hex'),
+  decode(data->'columns'->1->>'value', 'hex'),
+  (data->'columns'->2->>'value')::timestamptz,
+  (data->'columns'->3->>'value')::uuid
+FROM jdata`, slotName, b.cfg.ChangeFeedBatchSize)
+
+	var action string
+	var key []byte
+	var value []byte
+	var expires zeronull.Timestamptz
+	var revision zeronull.UUID
+	tag, err := pgx.ForEachRow(rows, []any{&action, &key, &value, &expires, &revision}, func() error {
+		switch action {
+		case "I", "U":
+			b.buf.Emit(backend.Event{
+				Type: types.OpPut,
+				Item: backend.Item{
+					Key:     key,
+					Value:   value,
+					Expires: time.Time(expires),
+				},
+			})
+			return nil
+		case "D":
+			b.buf.Emit(backend.Event{
+				Type: types.OpDelete,
+				Item: backend.Item{
+					Key: key,
+				},
+			})
+			return nil
+		case "M":
+			b.log.Debug("Received WAL message.")
+			return nil
+		case "B", "C":
+			b.log.Debug("Received transaction message in change feed (should not happen).")
+			return nil
+		case "T":
+			// it could be possible to just reset the event buffer and
+			// continue from the next row but it's not worth the effort
+			// compared to just killing this connection and reconnecting,
+			// and this should never actually happen anyway - deleting
+			// everything from the backend would leave Teleport in a very
+			// broken state
+			return trace.BadParameter("received truncate WAL message, can't continue")
+		default:
+			return trace.BadParameter("received unknown WAL message %q", action)
+		}
+	})
+	if err != nil {
+		return 0, trace.Wrap(err)
+	}
+
+	events := tag.RowsAffected()
+
+	if events > 0 {
+		b.log.WithFields(logrus.Fields{
+			"events":  events,
+			"elapsed": time.Since(t0).String(),
+		}).Debug("Fetched change feed events.")
+	}
+
+	return events, nil
+}
@@ -18,7 +18,6 @@ import (
 	"context"
 	"time"

-	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/gravitational/trace"
@@ -28,27 +27,11 @@ import (

 // AzureBeforeConnect will return a pgx BeforeConnect function suitable for
 // Azure AD authentication. The returned function will set the password of the
-// pgx.ConnConfig to a token for the relevant scope. If a client ID is provided,
-// authentication will only be attempted as the managed identity with said ID
-// rather than with all the default credentials.
-func AzureBeforeConnect(clientID string, log logrus.FieldLogger) (func(ctx context.Context, config *pgx.ConnConfig) error, error) {
-	var cred azcore.TokenCredential
-	if clientID != "" {
-		log.WithField("azure_client_id", clientID).Debug("Using Azure AD authentication with managed identity.")
-		c, err := azidentity.NewManagedIdentityCredential(&azidentity.ManagedIdentityCredentialOptions{
-			ID: azidentity.ClientID(clientID),
-		})
-		if err != nil {
-			return nil, trace.Wrap(err, "creating Azure managed identity credentials")
-		}
-		cred = c
-	} else {
-		log.Debug("Using Azure AD authentication with default credentials.")
-		c, err := azidentity.NewDefaultAzureCredential(nil)
-		if err != nil {
-			return nil, trace.Wrap(err, "creating default Azure credentials")
-		}
-		cred = c
+// connection to a token for the relevant scope.
+func AzureBeforeConnect(log logrus.FieldLogger) (func(ctx context.Context, config *pgx.ConnConfig) error, error) {
+	cred, err := azidentity.NewDefaultAzureCredential(nil)
+	if err != nil {
+		return nil, trace.Wrap(err, "creating Azure credentials")
 	}

 	beforeConnect := func(ctx context.Context, config *pgx.ConnConfig) error {
@@ -87,7 +87,9 @@ func TryEnsureDatabase(ctx context.Context, poolConfig *pgxpool.Config, log logr
 // Retry runs the closure potentially more than once, retrying quickly on
 // serialization or deadlock errors, and backing off more on other retryable
 // errors. It will not retry on network errors or other ambiguous errors after
-// any data has been sent.
+// any data has been sent. It will retry unique constraint violation and
+// exclusion constraint violations, so the closure should not rely on those for
+// normal behavior.
 func Retry[T any](ctx context.Context, log logrus.FieldLogger, f func() (T, error)) (T, error) {
 	const idempotent = false
 	v, err := retry(ctx, log, idempotent, f)
@@ -97,6 +99,8 @@ func Retry[T any](ctx context.Context, log logrus.FieldLogger, f func() (T, erro
 // RetryIdempotent runs the closure potentially more than once, retrying quickly
 // on serialization or deadlock errors, and backing off more on other errors. It
 // assumes that f is idempotent, so it will retry even in ambiguous situations.
+// It will retry unique constraint violation and exclusion constraint
+// violations, so the closure should not rely on those for normal behavior.
 func RetryIdempotent[T any](ctx context.Context, log logrus.FieldLogger, f func() (T, error)) (T, error) {
 	const idempotent = true
 	v, err := retry(ctx, log, idempotent, f)
@@ -131,7 +135,7 @@ func retry[T any](ctx context.Context, log logrus.FieldLogger, isIdempotent bool
 		var pgErr *pgconn.PgError
 		_ = errors.As(err, &pgErr)

-		if pgErr != nil && (pgErr.Code == pgerrcode.SerializationFailure || pgErr.Code == pgerrcode.DeadlockDetected) {
+		if pgErr != nil && isSerializationErrorCode(pgErr.Code) {
 			log.WithError(err).
 				WithField("attempt", i).
 				Debug("Operation failed due to conflicts, retrying quickly.")
@@ -175,6 +179,23 @@ func retry[T any](ctx context.Context, log logrus.FieldLogger, isIdempotent bool
 	return zeroT, trace.LimitExceeded("too many retries, last error: %v", err)
 }

+// isSerializationErrorCode returns true if the error code is for a
+// serialization error; this also includes unique_violation and
+// exclusion_violation, which are sometimes returned as a result of
+// serialization failures (and thus can be meaningfully retried) but can also be
+// a result of actual logical/relational errors, which would then cause the same
+// error to be raised again.
+func isSerializationErrorCode(code string) bool {
+	// source:
+	// https://www.postgresql.org/docs/current/mvcc-serialization-failure-handling.html
+	switch code {
+	case pgerrcode.SerializationFailure, pgerrcode.DeadlockDetected, pgerrcode.UniqueViolation, pgerrcode.ExclusionViolation:
+		return true
+	default:
+		return false
+	}
+}
+
 // RetryTx runs a closure like [Retry] or [RetryIdempotent], wrapped in
 // [pgx.BeginTxFunc].
 func RetryTx(
@@ -0,0 +1,516 @@
+// Copyright 2023 Gravitational, Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pgbk
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"sync"
+	"time"
+
+	"github.com/gravitational/trace"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgtype"
+	"github.com/jackc/pgx/v5/pgtype/zeronull"
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/jonboulle/clockwork"
+	"github.com/sirupsen/logrus"
+
+	"github.com/gravitational/teleport/api/types"
+	"github.com/gravitational/teleport/api/utils"
+	"github.com/gravitational/teleport/lib/backend"
+	pgcommon "github.com/gravitational/teleport/lib/backend/pgbk/common"
+)
+
+const (
+	Name    = "postgresql"
+	AltName = "postgres"
+
+	// componentName is the component name used for logging.
+	componentName = "pgbk"
+)
+
+const (
+	defaultChangeFeedBatchSize = 1000
+	defaultChangeFeedInterval  = backend.DefaultPollStreamPeriod
+
+	defaultExpiryBatchSize = 1000
+	defaultExpiryInterval  = 30 * time.Second
+)
+
+// AuthMode determines if we should use some environment-specific authentication
+// mechanism or credentials.
+type AuthMode string
+
+const (
+	// StaticAuth uses the static credentials as defined in the connection
+	// string.
+	StaticAuth AuthMode = ""
+	// AzureADAuth gets a connection token from Azure and uses it as the
+	// password when connecting.
+	AzureADAuth AuthMode = "azure"
+)
+
+// Check returns an error if the AuthMode is invalid.
+func (a AuthMode) Check() error {
+	switch a {
+	case StaticAuth, AzureADAuth:
+		return nil
+	default:
+		return trace.BadParameter("invalid authentication mode %q, should be %q or %q", a, StaticAuth, AzureADAuth)
+	}
+}
+
+// Config is the configuration struct for [Backend]; outside of tests or custom
+// code, it's usually generated by converting the [backend.Params] from the
+// Teleport configuration file.
+type Config struct {
+	ConnString string `json:"conn_string"`
+
+	AuthMode AuthMode `json:"auth_mode"`
+
+	ChangeFeedPollInterval types.Duration `json:"change_feed_poll_interval"`
+	ChangeFeedBatchSize    int            `json:"change_feed_batch_size"`
+
+	DisableExpiry   bool           `json:"disable_expiry"`
+	ExpiryInterval  types.Duration `json:"expiry_interval"`
+	ExpiryBatchSize int            `json:"expiry_batch_size"`
+}
+
+func (c *Config) CheckAndSetDefaults() error {
+	if err := c.AuthMode.Check(); err != nil {
+		return trace.Wrap(err)
+	}
+
+	if c.ChangeFeedPollInterval < 0 {
+		return trace.BadParameter("change feed poll interval must be non-negative")
+	}
+	if c.ChangeFeedPollInterval == 0 {
+		c.ChangeFeedPollInterval = types.Duration(defaultChangeFeedInterval)
+	}
+	if c.ChangeFeedBatchSize < 0 {
+		return trace.BadParameter("change feed batch size must be non-negative")
+	}
+	if c.ChangeFeedBatchSize == 0 {
+		c.ChangeFeedBatchSize = defaultChangeFeedBatchSize
+	}
+
+	if c.ExpiryInterval < 0 {
+		return trace.BadParameter("expiry interval must be non-negative")
+	}
+	if c.ExpiryInterval == 0 {
+		c.ExpiryInterval = types.Duration(defaultExpiryInterval)
+	}
+	if c.ExpiryBatchSize < 0 {
+		return trace.BadParameter("expiry batch size must be non-negative")
+	}
+	if c.ExpiryBatchSize == 0 {
+		c.ExpiryBatchSize = defaultExpiryBatchSize
+	}
+
+	return nil
+}
+
+// NewFromParams starts and returns a [*Backend] with the given params
+// (generally read from the Teleport configuration file).
+func NewFromParams(ctx context.Context, params backend.Params) (*Backend, error) {
+	var cfg Config
+	if err := utils.ObjectToStruct(params, &cfg); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	bk, err := NewWithConfig(ctx, cfg)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	return bk, nil
+}
+
+// NewWithConfig starts and returns a [*Backend] with the given [Config].
+func NewWithConfig(ctx context.Context, cfg Config) (*Backend, error) {
+	if err := cfg.CheckAndSetDefaults(); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	poolConfig, err := pgxpool.ParseConfig(cfg.ConnString)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	log := logrus.WithField(trace.Component, componentName)
+
+	if cfg.AuthMode == AzureADAuth {
+		bc, err := pgcommon.AzureBeforeConnect(log)
+		if err != nil {
+			return nil, trace.Wrap(err)
+		}
+		poolConfig.BeforeConnect = bc
+	}
+
+	const defaultTxIsoParamName = "default_transaction_isolation"
+	if defaultTxIso := poolConfig.ConnConfig.RuntimeParams[defaultTxIsoParamName]; defaultTxIso != "" {
+		log.WithField(defaultTxIsoParamName, defaultTxIso).
+			Error("The " + defaultTxIsoParamName + " parameter was overridden in the connection string; proceeding with an unsupported configuration.")
+	} else {
+		poolConfig.ConnConfig.RuntimeParams[defaultTxIsoParamName] = "serializable"
+	}
+
+	log.Info("Setting up backend.")
+
+	pgcommon.TryEnsureDatabase(ctx, poolConfig, log)
+
+	pool, err := pgxpool.NewWithConfig(ctx, poolConfig)
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	if err := pgcommon.SetupAndMigrate(ctx, log, pool, "backend_version", schemas); err != nil {
+		pool.Close()
+		return nil, trace.Wrap(err)
+	}
+
+	ctx, cancel := context.WithCancel(ctx)
+	b := &Backend{
+		cfg:    cfg,
+		log:    log,
+		pool:   pool,
+		buf:    backend.NewCircularBuffer(),
+		cancel: cancel,
+	}
+
+	if !cfg.DisableExpiry {
+		b.wg.Add(1)
+		go func() {
+			defer b.wg.Done()
+			b.backgroundExpiry(ctx)
+		}()
+	}
+
+	b.wg.Add(1)
+	go func() {
+		defer b.wg.Done()
+		b.backgroundChangeFeed(ctx)
+	}()
+
+	return b, nil
+}
+
+// Backend is a PostgreSQL-backed [backend.Backend].
+type Backend struct {
+	cfg  Config
+	log  logrus.FieldLogger
+	pool *pgxpool.Pool
+	buf  *backend.CircularBuffer
+
+	cancel context.CancelFunc
+	wg     sync.WaitGroup
+}
+
+func (b *Backend) Close() error {
+	b.cancel()
+	b.wg.Wait()
+	b.buf.Close()
+	b.pool.Close()
+	return nil
+}
+
+var schemas = []string{
+	`CREATE TABLE kv (
+		key bytea NOT NULL,
+		value bytea NOT NULL,
+		expires timestamptz,
+		revision uuid NOT NULL,
+		CONSTRAINT kv_pkey PRIMARY KEY (key)
+	);
+	CREATE INDEX kv_expires_idx ON kv (expires) WHERE expires IS NOT NULL;`,
+}
+
+var _ backend.Backend = (*Backend)(nil)
+
+// GetName implements [backend.Backend].
+func (*Backend) GetName() string {
+	return Name
+}
+
+// Create implements [backend.Backend].
+func (b *Backend) Create(ctx context.Context, i backend.Item) (*backend.Lease, error) {
+	revision := newRevision()
+	created, err := pgcommon.Retry(ctx, b.log, func() (bool, error) {
+		tag, err := b.pool.Exec(ctx,
+			"INSERT INTO kv (key, value, expires, revision) VALUES ($1, $2, $3, $4)"+
+				" ON CONFLICT (key) DO UPDATE SET"+
+				" value = excluded.value, expires = excluded.expires, revision = excluded.revision"+
+				" WHERE kv.expires IS NOT NULL AND kv.expires <= now()",
+			i.Key, i.Value, zeronull.Timestamptz(i.Expires.UTC()), revision)
+		if err != nil {
+			return false, trace.Wrap(err)
+		}
+		return tag.RowsAffected() > 0, nil
+	})
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	if !created {
+		return nil, trace.AlreadyExists("key %q already exists", i.Key)
+	}
+	return newLease(i), nil
+}
+
+// Put implements [backend.Backend].
+func (b *Backend) Put(ctx context.Context, i backend.Item) (*backend.Lease, error) {
+	revision := newRevision()
+	if _, err := pgcommon.Retry(ctx, b.log, func() (struct{}, error) {
+		_, err := b.pool.Exec(ctx,
+			"INSERT INTO kv (key, value, expires, revision) VALUES ($1, $2, $3, $4)"+
+				" ON CONFLICT (key) DO UPDATE SET"+
+				" value = excluded.value, expires = excluded.expires, revision = excluded.revision",
+			i.Key, i.Value, zeronull.Timestamptz(i.Expires.UTC()), revision)
+		return struct{}{}, trace.Wrap(err)
+	}); err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	return newLease(i), nil
+}
+
+// CompareAndSwap implements [backend.Backend].
+func (b *Backend) CompareAndSwap(ctx context.Context, expected backend.Item, replaceWith backend.Item) (*backend.Lease, error) {
+	if !bytes.Equal(expected.Key, replaceWith.Key) {
+		return nil, trace.BadParameter("expected and replaceWith keys should match")
+	}
+
+	revision := newRevision()
+	swapped, err := pgcommon.Retry(ctx, b.log, func() (bool, error) {
+		tag, err := b.pool.Exec(ctx,
+			"UPDATE kv SET value = $1, expires = $2, revision = $3"+
+				" WHERE key = $4 AND value = $5 AND (expires IS NULL OR expires > now())",
+			replaceWith.Value, zeronull.Timestamptz(replaceWith.Expires.UTC()), revision,
+			replaceWith.Key, expected.Value)
+		if err != nil {
+			return false, trace.Wrap(err)
+		}
+		return tag.RowsAffected() > 0, nil
+	})
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	if !swapped {
+		return nil, trace.CompareFailed("key %q does not exist or does not match expected", replaceWith.Key)
+	}
+	return newLease(replaceWith), nil
+}
+
+// Update implements [backend.Backend].
+func (b *Backend) Update(ctx context.Context, i backend.Item) (*backend.Lease, error) {
+	revision := newRevision()
+	updated, err := pgcommon.Retry(ctx, b.log, func() (bool, error) {
+		tag, err := b.pool.Exec(ctx,
+			"UPDATE kv SET value = $1, expires = $2, revision = $3"+
+				" WHERE key = $4 AND (expires IS NULL OR expires > now())",
+			i.Value, zeronull.Timestamptz(i.Expires.UTC()), revision, i.Key)
+		if err != nil {
+			return false, trace.Wrap(err)
+		}
+		return tag.RowsAffected() > 0, nil
+	})
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	if !updated {
+		return nil, trace.NotFound("key %q does not exist", i.Key)
+	}
+	return newLease(i), nil
+}
+
+// Get implements [backend.Backend].
+func (b *Backend) Get(ctx context.Context, key []byte) (*backend.Item, error) {
+	item, err := pgcommon.RetryIdempotent(ctx, b.log, func() (*backend.Item, error) {
+		batch := new(pgx.Batch)
+		// batches run in an implicit transaction
+		batch.Queue("SET transaction_read_only TO on")
+
+		var item *backend.Item
+		batch.Queue("SELECT value, expires, revision FROM kv"+
+			" WHERE key = $1 AND (expires IS NULL OR expires > now())", key,
+		).QueryRow(func(row pgx.Row) error {
+			var value []byte
+			var expires zeronull.Timestamptz
+			var revision pgtype.UUID
+			if err := row.Scan(&value, &expires, &revision); err != nil {
+				if errors.Is(err, pgx.ErrNoRows) {
+					return nil
+				}
+				return trace.Wrap(err)
+			}
+
+			item = &backend.Item{
+				Key:     key,
+				Value:   value,
+				Expires: time.Time(expires).UTC(),
+				// revision isn't supported in backend.Item yet
+			}
+			return nil
+		})
+
+		if err := b.pool.SendBatch(ctx, batch).Close(); err != nil {
+			return nil, trace.Wrap(err)
+		}
+
+		return item, nil
+	})
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	if item == nil {
+		return nil, trace.NotFound("key %q does not exist", key)
+	}
+	return item, nil
+}
+
+// GetRange implements [backend.Backend].
+func (b *Backend) GetRange(ctx context.Context, startKey []byte, endKey []byte, limit int) (*backend.GetResult, error) {
+	if limit <= 0 {
+		limit = backend.DefaultRangeLimit
+	}
+
+	items, err := pgcommon.RetryIdempotent(ctx, b.log, func() ([]backend.Item, error) {
+		batch := new(pgx.Batch)
+		// batches run in an implicit transaction
+		batch.Queue("SET transaction_read_only TO on")
+		// TODO(espadolini): figure out if we want transaction_deferred enabled
+		// for GetRange
+
+		var items []backend.Item
+		batch.Queue(
+			"SELECT key, value, expires, revision FROM kv"+
+				" WHERE key BETWEEN $1 AND $2 AND (expires IS NULL OR expires > now())"+
+				" ORDER BY key LIMIT $3",
+			startKey, endKey, limit,
+		).Query(func(rows pgx.Rows) error {
+			var err error
+			items, err = pgx.CollectRows(rows, func(row pgx.CollectableRow) (backend.Item, error) {
+				var key, value []byte
+				var expires zeronull.Timestamptz
+				var revision pgtype.UUID
+				if err := row.Scan(&key, &value, &expires, &revision); err != nil {
+					return backend.Item{}, err
+				}
+				return backend.Item{
+					Key:     key,
+					Value:   value,
+					Expires: time.Time(expires).UTC(),
+					// revision isn't supported in backend.Item yet
+				}, nil
+			})
+			return trace.Wrap(err)
+		})
+
+		if err := b.pool.SendBatch(ctx, batch).Close(); err != nil {
+			return nil, trace.Wrap(err)
+		}
+
+		return items, nil
+	})
+	if err != nil {
+		return nil, trace.Wrap(err)
+	}
+
+	return &backend.GetResult{Items: items}, nil
+}
+
+// Delete implements [backend.Backend].
+func (b *Backend) Delete(ctx context.Context, key []byte) error {
+	deleted, err := pgcommon.Retry(ctx, b.log, func() (bool, error) {
+		tag, err := b.pool.Exec(ctx,
+			"DELETE FROM kv WHERE key = $1 AND (expires IS NULL OR expires > now())", key)
+		if err != nil {
+			return false, trace.Wrap(err)
+		}
+		return tag.RowsAffected() > 0, nil
+	})
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
+	if !deleted {
+		return trace.NotFound("key %q does not exist", key)
+	}
+	return nil
+}
+
+// DeleteRange implements [backend.Backend].
+func (b *Backend) DeleteRange(ctx context.Context, startKey []byte, endKey []byte) error {
+	// this is the only backend operation that might affect a disproportionate
+	// amount of rows at the same time; in actual operation, DeleteRange hardly
+	// ever deletes more than dozens of items at once, and logical decoding
+	// starts having performance issues when a transaction affects _thousands_
+	// of rows at once, so we're good here (but see [Backend.backgroundExpiry])
+	if _, err := pgcommon.Retry(ctx, b.log, func() (struct{}, error) {
+		_, err := b.pool.Exec(ctx,
+			"DELETE FROM kv WHERE key BETWEEN $1 AND $2",
+			startKey, endKey,
+		)
+		return struct{}{}, trace.Wrap(err)
+	}); err != nil {
+		return trace.Wrap(err)
+	}
+
+	return nil
+}
+
+// KeepAlive implements [backend.Backend].
+func (b *Backend) KeepAlive(ctx context.Context, lease backend.Lease, expires time.Time) error {
+	revision := newRevision()
+	updated, err := pgcommon.Retry(ctx, b.log, func() (bool, error) {
+		tag, err := b.pool.Exec(ctx,
+			"UPDATE kv SET expires = $1, revision = $2 WHERE key = $3 AND (expires IS NULL OR expires > now())",
+			zeronull.Timestamptz(expires.UTC()), revision, lease.Key)
+		if err != nil {
+			return false, trace.Wrap(err)
+		}
+		return tag.RowsAffected() > 0, nil
+	})
+	if err != nil {
+		return trace.Wrap(err)
+	}
+
+	if !updated {
+		return trace.NotFound("key %q does not exist", lease.Key)
+	}
+	return nil
+}
+
+// NewWatcher implements [backend.Backend].
+func (b *Backend) NewWatcher(ctx context.Context, watch backend.Watch) (backend.Watcher, error) {
+	return b.buf.NewWatcher(ctx, watch)
+}
+
+// CloseWatchers implements [backend.Backend].
+func (b *Backend) CloseWatchers() { b.buf.Clear() }
+
+// Clock implements [backend.Backend].
+func (b *Backend) Clock() clockwork.Clock {
+	// we don't support a custom clock, because deciding which items still exist
+	// in the backend depends on which items are still stored but expired, and
+	// it's much cleaner to just rely on the server transaction time (which is
+	// shared between all auth servers) for that
+	return clockwork.NewRealClock()
+}
@@ -0,0 +1,71 @@
+// Copyright 2023 Gravitational, Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pgbk
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"testing"
+
+	"github.com/gravitational/trace"
+	"github.com/jonboulle/clockwork"
+	"github.com/stretchr/testify/require"
+
+	"github.com/gravitational/teleport/lib/backend"
+	"github.com/gravitational/teleport/lib/backend/test"
+	"github.com/gravitational/teleport/lib/utils"
+)
+
+func TestMain(m *testing.M) {
+	utils.InitLoggerForTests()
+	os.Exit(m.Run())
+}
+
+func TestPostgresBackend(t *testing.T) {
+	// expiry_interval needs to be really short to pass some of the tests, and a
+	// faster poll interval helps a bit with runtime:
+	// {"conn_string":"...","expiry_interval":"500ms","change_feed_poll_interval":"500ms"}
+	paramString := os.Getenv("TELEPORT_PGBK_TEST_PARAMS_JSON")
+	if paramString == "" {
+		t.Skip("Postgres backend tests are disabled. Enable them by setting the TELEPORT_PGBK_TEST_PARAMS_JSON variable.")
+	}
+
+	newBackend := func(options ...test.ConstructionOption) (backend.Backend, clockwork.FakeClock, error) {
+		testCfg, err := test.ApplyOptions(options)
+		if err != nil {
+			return nil, nil, trace.Wrap(err)
+		}
+
+		if testCfg.MirrorMode {
+			return nil, nil, test.ErrMirrorNotSupported
+		}
+
+		if testCfg.ConcurrentBackend != nil {
+			return nil, nil, test.ErrConcurrentAccessNotSupported
+		}
+
+		var params backend.Params
+		require.NoError(t, json.Unmarshal([]byte(paramString), &params))
+
+		uut, err := NewFromParams(context.Background(), params)
+		if err != nil {
+			return nil, nil, trace.Wrap(err)
+		}
+		return uut, test.BlockingFakeClock{Clock: clockwork.NewRealClock()}, nil
+	}
+
+	test.RunBackendComplianceSuite(t, newBackend)
+}
@@ -0,0 +1,41 @@
+// Copyright 2023 Gravitational, Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pgbk
+
+import (
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5/pgtype"
+
+	"github.com/gravitational/teleport/lib/backend"
+)
+
+// newLease returns a non-nil [*backend.Lease] that's filled in with the details
+// of the item (i.e. its key) if the item has an expiration time.
+func newLease(i backend.Item) *backend.Lease {
+	var lease backend.Lease
+	if !i.Expires.IsZero() {
+		lease.Key = i.Key
+	}
+	return &lease
+}
+
+// newRevision returns a random, non-null [pgtype.UUID] to use as a row
+// revision.
+func newRevision() pgtype.UUID {
+	return pgtype.UUID{
+		Bytes: uuid.New(),
+		Valid: true,
+	}
+}
@@ -41,9 +41,6 @@ import (
 	"github.com/gravitational/teleport/lib/session"
 )

-// clientIDParam is the parameter in the fragment that specifies the optional client ID.
-const clientIDParam = "azure_client_id"
-
 // sessionContainerParam and inprogressContainerParam are the parameters in the
 // fragment that specify the containers to use for finalized session recordings
 // and in-progress data.
@@ -109,10 +106,6 @@ type Config struct {
 	// [defaultInprogressContainerName].
 	InprogressContainerName string

-	// ClientID, when set, defines the managed identity's client ID to use for
-	// authentication.
-	ClientID string
-
 	// Log is the logger to use. If unset, it will default to the global logger
 	// with a component of "azblob".
 	Log logrus.FieldLogger
@@ -146,7 +139,6 @@ func (c *Config) SetFromURL(u *url.URL) error {
 	c.ServiceURL.Fragment = ""
 	c.ServiceURL.RawFragment = ""

-	c.ClientID = params.Get(clientIDParam)
 	c.SessionContainerName = params.Get(sessionContainerParam)
 	c.InprogressContainerName = params.Get(inprogressContainerParam)

@@ -174,21 +166,9 @@ func NewHandler(ctx context.Context, cfg Config) (*Handler, error) {
 		return nil, trace.Wrap(err)
 	}

-	var cred azcore.TokenCredential
-	if cfg.ClientID != "" {
-		c, err := azidentity.NewManagedIdentityCredential(&azidentity.ManagedIdentityCredentialOptions{
-			ID: azidentity.ClientID(cfg.ClientID),
-		})
-		if err != nil {
-			return nil, trace.Wrap(err, "creating Azure managed identity credentials")
-		}
-		cred = c
-	} else {
-		c, err := azidentity.NewDefaultAzureCredential(nil)
-		if err != nil {
-			return nil, trace.Wrap(err, "creating default Azure credentials")
-		}
-		cred = c
+	cred, err := azidentity.NewDefaultAzureCredential(nil)
+	if err != nil {
+		return nil, trace.Wrap(err, "creating Azure credentials")
 	}

 	ensureContainer := func(name string) (*container.Client, error) {
@@ -52,8 +52,7 @@ const (

 // URL parameters for configuration.
 const (
-	authModeParam      = "auth_mode"
-	azureClientIDParam = "azure_client_id"
+	authModeParam = "auth_mode"

 	disableCleanupParam  = "disable_cleanup"
 	cleanupIntervalParam = "cleanup_interval"
@@ -88,8 +87,7 @@ type Config struct {
 	Log        logrus.FieldLogger
 	PoolConfig *pgxpool.Config

-	AuthMode      AuthMode
-	AzureClientID string
+	AuthMode AuthMode

 	DisableCleanup  bool
 	RetentionPeriod time.Duration
@@ -119,7 +117,6 @@ func (c *Config) SetFromURL(u *url.URL) error {
 	c.PoolConfig = poolConfig

 	c.AuthMode = AuthMode(params.Get(authModeParam))
-	c.AzureClientID = params.Get(azureClientIDParam)

 	if s := params.Get(disableCleanupParam); s != "" {
 		b, err := strconv.ParseBool(s)
@@ -159,10 +156,6 @@ func (c *Config) CheckAndSetDefaults() error {
 		return trace.Wrap(err)
 	}

-	if c.AzureClientID != "" && c.AuthMode != AzureADAuth {
-		return trace.BadParameter("azure client ID requires azure auth mode")
-	}
-
 	if c.RetentionPeriod < 0 {
 		return trace.BadParameter("retention period must be non-negative")
 	}
@@ -192,7 +185,7 @@ func New(ctx context.Context, cfg Config) (*Log, error) {
 	}

 	if cfg.AuthMode == AzureADAuth {
-		bc, err := pgcommon.AzureBeforeConnect(cfg.AzureClientID, cfg.Log)
+		bc, err := pgcommon.AzureBeforeConnect(cfg.Log)
 		if err != nil {
 			return nil, trace.Wrap(err)
 		}
@@ -100,8 +100,7 @@ func TestConfig(t *testing.T) {
 			CleanupInterval: defaultCleanupInterval,
 		},

-		"postgres://foo#azure_client_id=not-in-azure-auth-mode": nil,
-		"postgres://foo#auth_mode=invalid-auth-mode":            nil,
+		"postgres://foo#auth_mode=invalid-auth-mode": nil,
 	}

 	for u, expectedConfig := range configs {
@@ -83,6 +83,7 @@ import (
 	"github.com/gravitational/teleport/lib/backend/kubernetes"
 	"github.com/gravitational/teleport/lib/backend/lite"
 	"github.com/gravitational/teleport/lib/backend/memory"
+	"github.com/gravitational/teleport/lib/backend/pgbk"
 	"github.com/gravitational/teleport/lib/bpf"
 	"github.com/gravitational/teleport/lib/cache"
 	"github.com/gravitational/teleport/lib/cloud"
@@ -5172,6 +5173,9 @@ func (process *TeleportProcess) initAuthStorage() (bk backend.Backend, err error
 	// etcd backend.
 	case etcdbk.GetName():
 		bk, err = etcdbk.New(ctx, bc.Params)
+	// PostgreSQL backend
+	case pgbk.Name, pgbk.AltName:
+		bk, err = pgbk.NewFromParams(ctx, bc.Params)
 	default:
 		err = trace.BadParameter("unsupported secrets storage type: %q", bc.Type)
 	}
@@ -105,7 +105,6 @@ teleport:
    type: postgresql
    conn_string: host=testdb.postgres.database.azure.com port=5432 sslmode=verify-full user=teleport_auth dbname=teleport_backend
    auth_mode: azure
-    azure_client_id: ""

    change_feed_poll_interval: "1s"
    change_feed_batch_size: 10000
@@ -120,7 +119,7 @@ teleport:

 The `conn_string` used by the backend is a [`libpq`-compatible Postgres connection string](https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING), and can be specified in both keyword-value format and as a URI. The audit log URI can only be specified as a URI. In addition to the standard Postgres options, the options supported by [`pgxpool`](https://pkg.go.dev/github.com/jackc/pgx/v5/pgxpool#ParseConfig) are also supported - `pool_max_conns` being the most useful. Non-Postgres parameters for the audit log are passed as query parameters in the fragment of the URI.

-The `auth_mode: azure` field for the backend and the `auth_mode=azure` query parameter enable authentication to the database via Azure AD. The optional `azure_client_id` field (for both the backend and the audit log) override the default azure credential selection and the selection in the `AZURE_CLIENT_ID` envvar, so it's possible to select two different managed identities for backend and audit log. A blank (or unset) `auth_mode` will use the default Postgres authentication methods.
+The `auth_mode: azure` field for the backend and the `auth_mode=azure` query parameter enable authentication to the database via Azure AD.

 The `change_feed_*` options regulate the internal behavior of the change feed polling, the `expiry_*` and `disable_expiry` regulate the behavior of the TTL deletions. We're not expecting users to have to tweak these options (except maybe `disable_expiry` if TTL deletion is going to be run externally with `pg_cron` or similar methods), they're just exposed for performance tuning.

@@ -32,9 +32,9 @@ teleport:
    audit_sessions_uri: "azblob://<storage account name>.blob.core.windows.net"
 ```

-It's possible to specify a client ID for a managed identity (to allow for different identities to be used for backend, events and sessions, as opposed to setting the `AZURE_CLIENT_ID` envvar and using the default credentials) by specifying a URL fragment of `#azure_client_id=11111111-2222-3333-4444-555555555555`. The `azblob` schema is only used by Teleport to identify the storage backend, and will actually result in a `https` URL. For testing against simulators (or other services with the same API surface) over `http`, the `azblob-http` schema is also supported (but we will leave it undocumented for end users).
+The `azblob` schema is only used by Teleport to identify the storage backend, and will actually result in a `https` URL. For testing against simulators (or other services with the same API surface) over `http`, the `azblob-http` schema is also supported (but we will leave it undocumented for end users).

-Teleport will require access to two containers in the storage account, named `session`, which will store the completed session files, and `inprogress`, which will be used as temporary space to hold parts. The container names can be tweaked with query parameters in the fragment: `session_container=foo&inprogress_container=bar`.
+Teleport will require access to two containers in the storage account, named `session`, which will store the completed session files, and `inprogress`, which will be used as temporary space to hold parts. The container names can be tweaked with query parameters in the fragment of the URL: `#session_container=foo&inprogress_container=bar`.

 To signify the creation of an upload, a random UUID is generated, and an empty blob is created in `inprogress` with name `upload/<session id>/<upload id>`; this will be used when listing unfinished uploads, by listing blobs whose name begins with `upload/`. Parts are uploaded in `inprogress` with name `part/<session id>/<upload id>/<part number>`, which will allow listing the parts for a specific upload. To complete an upload, parts will be used as blocks to compose the final session recording in the `session` container (in a blob that has the session ID as its name); this can be done efficiently with the "Put Block By URL" call.