go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/tq/txn/spanner/spanner.go (about)

     1  // Copyright 2020 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package spanner
    16  
    17  import (
    18  	"context"
    19  
    20  	"cloud.google.com/go/spanner"
    21  
    22  	"go.chromium.org/luci/common/errors"
    23  	"go.chromium.org/luci/common/retry/transient"
    24  
    25  	"go.chromium.org/luci/server/span"
    26  	"go.chromium.org/luci/server/tq/internal/reminder"
    27  )
    28  
    29  // tableName is the name of the table that user must create in their Spanner
    30  // database prior to using this package.
    31  //
    32  // CREATE TABLE TQReminders (
    33  //
    34  //	ID STRING(MAX) NOT NULL,
    35  //	FreshUntil TIMESTAMP NOT NULL,
    36  //	Payload BYTES(102400) NOT NULL,
    37  //
    38  // ) PRIMARY KEY (ID ASC);
    39  //
    40  // If you ever need to change this, change also user-visible server/tq doc.
    41  const tableName = "TQReminders"
    42  
    43  type spanDB struct{}
    44  
    45  func (spanDB) Kind() string {
    46  	return "spanner"
    47  }
    48  
    49  func (spanDB) Defer(ctx context.Context, cb func(context.Context)) {
    50  	span.Defer(ctx, cb)
    51  }
    52  
    53  func (spanDB) SaveReminder(ctx context.Context, r *reminder.Reminder) error {
    54  	span.BufferWrite(ctx, spanner.Insert(
    55  		tableName,
    56  		[]string{"ID", "FreshUntil", "Payload"},
    57  		[]any{r.ID, r.FreshUntil, r.RawPayload}))
    58  	return nil
    59  }
    60  
    61  func (spanDB) DeleteReminder(ctx context.Context, r *reminder.Reminder) error {
    62  	_, err := span.Apply(ctx, []*spanner.Mutation{
    63  		spanner.Delete(tableName, spanner.Key{r.ID}),
    64  	}, spanner.ApplyAtLeastOnce())
    65  	if err != nil {
    66  		return errors.Annotate(err, "failed to delete the Reminder %s", r.ID).Tag(transient.Tag).Err()
    67  	}
    68  	return nil
    69  }
    70  
    71  func (spanDB) FetchRemindersMeta(ctx context.Context, low string, high string, limit int) (res []*reminder.Reminder, err error) {
    72  	iter := span.ReadWithOptions(span.Single(ctx),
    73  		tableName,
    74  		spanner.KeyRange{Start: spanner.Key{low}, End: spanner.Key{high}},
    75  		[]string{"ID", "FreshUntil"},
    76  		&spanner.ReadOptions{Limit: limit})
    77  	err = iter.Do(func(row *spanner.Row) error {
    78  		r := &reminder.Reminder{}
    79  		if err := row.Columns(&r.ID, &r.FreshUntil); err != nil {
    80  			return err
    81  		}
    82  		res = append(res, r)
    83  		return nil
    84  	})
    85  	if err != nil && err != context.DeadlineExceeded {
    86  		err = errors.Annotate(err, "failed to fetch Reminder keys").Tag(transient.Tag).Err()
    87  	}
    88  	return
    89  }
    90  
    91  func (spanDB) FetchReminderRawPayloads(ctx context.Context, batch []*reminder.Reminder) ([]*reminder.Reminder, error) {
    92  	ks := make([]spanner.KeySet, len(batch))
    93  	for i, r := range batch {
    94  		ks[i] = spanner.Key{r.ID}
    95  	}
    96  
    97  	iter := span.Read(span.Single(ctx), tableName, spanner.KeySets(ks...),
    98  		[]string{"ID", "FreshUntil", "Payload"})
    99  	i := 0
   100  	err := iter.Do(func(row *spanner.Row) error {
   101  		if err := row.Columns(&batch[i].ID, &batch[i].FreshUntil, &batch[i].RawPayload); err != nil {
   102  			return err
   103  		}
   104  		i++
   105  		return nil
   106  	})
   107  	if err != nil {
   108  		return batch[:i], errors.Annotate(err, "failed to fetch Reminders").Tag(transient.Tag).Err()
   109  	}
   110  	return batch[:i], nil
   111  }