github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/datastore/postgres/gc.go (about)

     1  package postgres
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strings"
     7  	"time"
     8  
     9  	sq "github.com/Masterminds/squirrel"
    10  
    11  	"github.com/authzed/spicedb/internal/datastore/common"
    12  	"github.com/authzed/spicedb/pkg/datastore"
    13  )
    14  
    15  var (
    16  	_ common.GarbageCollector = (*pgDatastore)(nil)
    17  
    18  	// we are using "tableoid" to globally identify the row through the "ctid" in partitioned environments
    19  	// as it's not guaranteed 2 rows in different partitions have different "ctid" values
    20  	// See https://www.postgresql.org/docs/current/ddl-system-columns.html#DDL-SYSTEM-COLUMNS-TABLEOID
    21  	gcPKCols = []string{"tableoid", "ctid"}
    22  )
    23  
    24  func (pgd *pgDatastore) HasGCRun() bool {
    25  	return pgd.gcHasRun.Load()
    26  }
    27  
    28  func (pgd *pgDatastore) MarkGCCompleted() {
    29  	pgd.gcHasRun.Store(true)
    30  }
    31  
    32  func (pgd *pgDatastore) ResetGCCompleted() {
    33  	pgd.gcHasRun.Store(false)
    34  }
    35  
    36  func (pgd *pgDatastore) Now(ctx context.Context) (time.Time, error) {
    37  	// Retrieve the `now` time from the database.
    38  	nowSQL, nowArgs, err := getNow.ToSql()
    39  	if err != nil {
    40  		return time.Time{}, err
    41  	}
    42  
    43  	var now time.Time
    44  	err = pgd.readPool.QueryRow(ctx, nowSQL, nowArgs...).Scan(&now)
    45  	if err != nil {
    46  		return time.Time{}, err
    47  	}
    48  
    49  	// RelationTupleTransaction is not timezone aware -- explicitly use UTC
    50  	// before using as a query arg.
    51  	return now.UTC(), nil
    52  }
    53  
    54  func (pgd *pgDatastore) TxIDBefore(ctx context.Context, before time.Time) (datastore.Revision, error) {
    55  	// Find the highest transaction ID before the GC window.
    56  	sql, args, err := getRevision.Where(sq.Lt{colTimestamp: before}).ToSql()
    57  	if err != nil {
    58  		return datastore.NoRevision, err
    59  	}
    60  
    61  	var value xid8
    62  	var snapshot pgSnapshot
    63  	err = pgd.readPool.QueryRow(ctx, sql, args...).Scan(&value, &snapshot)
    64  	if err != nil {
    65  		return datastore.NoRevision, err
    66  	}
    67  
    68  	return postgresRevision{snapshot}, nil
    69  }
    70  
    71  func (pgd *pgDatastore) DeleteBeforeTx(ctx context.Context, txID datastore.Revision) (common.DeletionCounts, error) {
    72  	revision := txID.(postgresRevision)
    73  
    74  	minTxAlive := newXid8(revision.snapshot.xmin)
    75  	removed := common.DeletionCounts{}
    76  	var err error
    77  	// Delete any relationship rows that were already dead when this transaction started
    78  	removed.Relationships, err = pgd.batchDelete(
    79  		ctx,
    80  		tableTuple,
    81  		gcPKCols,
    82  		sq.Lt{colDeletedXid: minTxAlive},
    83  	)
    84  	if err != nil {
    85  		return removed, fmt.Errorf("failed to GC relationships table: %w", err)
    86  	}
    87  
    88  	// Delete all transaction rows with ID < the transaction ID.
    89  	//
    90  	// We don't delete the transaction itself to ensure there is always at least
    91  	// one transaction present.
    92  	removed.Transactions, err = pgd.batchDelete(
    93  		ctx,
    94  		tableTransaction,
    95  		gcPKCols,
    96  		sq.Lt{colXID: minTxAlive},
    97  	)
    98  	if err != nil {
    99  		return removed, fmt.Errorf("failed to GC transactions table: %w", err)
   100  	}
   101  
   102  	// Delete any namespace rows with deleted_transaction <= the transaction ID.
   103  	removed.Namespaces, err = pgd.batchDelete(
   104  		ctx,
   105  		tableNamespace,
   106  		gcPKCols,
   107  		sq.Lt{colDeletedXid: minTxAlive},
   108  	)
   109  	if err != nil {
   110  		return removed, fmt.Errorf("failed to GC namespaces table: %w", err)
   111  	}
   112  
   113  	return removed, err
   114  }
   115  
   116  func (pgd *pgDatastore) batchDelete(
   117  	ctx context.Context,
   118  	tableName string,
   119  	pkCols []string,
   120  	filter sqlFilter,
   121  ) (int64, error) {
   122  	sql, args, err := psql.Select(pkCols...).From(tableName).Where(filter).Limit(gcBatchDeleteSize).ToSql()
   123  	if err != nil {
   124  		return -1, err
   125  	}
   126  
   127  	pkColsExpression := strings.Join(pkCols, ", ")
   128  	query := fmt.Sprintf(`WITH rows AS (%[1]s)
   129  		  DELETE FROM %[2]s
   130  		  WHERE (%[3]s) IN (SELECT %[3]s FROM rows);
   131  	`, sql, tableName, pkColsExpression)
   132  
   133  	var deletedCount int64
   134  	for {
   135  		cr, err := pgd.writePool.Exec(ctx, query, args...)
   136  		if err != nil {
   137  			return deletedCount, err
   138  		}
   139  
   140  		rowsDeleted := cr.RowsAffected()
   141  		deletedCount += rowsDeleted
   142  		if rowsDeleted < gcBatchDeleteSize {
   143  			break
   144  		}
   145  	}
   146  
   147  	return deletedCount, nil
   148  }