internal/fstore: support for firestore

Add a package that provides common support functions for
the GCP Firestore service.

Use it in the internal/jobs package.

Change-Id: Iac3bdcc22bb5abe6a3609d9545b66958c37e0cb9
Reviewed-on: https://go-review.googlesource.com/c/pkgsite-metrics/+/551755
TryBot-Result: Gopher Robot <[email protected]>
Reviewed-by: Zvonimir Pavlinovic <[email protected]>
Run-TryBot: Jonathan Amsterdam <[email protected]>

internal/fstore/fstore.go

@@ -0,0 +1,77 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package fstore provides general support for Firestore.
+// Its main feature is separate namespaces, to mimic separate
+// databases for different purposes (prod, dev, test, etc.).
+package fstore
+
+import (
+	"context"
+	"errors"
+
+	"cloud.google.com/go/firestore"
+	"golang.org/x/pkgsite-metrics/internal/derrors"
+)
+
+const namespaceCollection = "Namespaces"
+
+// A Namespace is a top-level collection for partitioning a Firestore
+// database into separate segments.
+type Namespace struct {
+	client *firestore.Client
+	name   string
+	doc    *firestore.DocumentRef
+}
+
+// OpenNamespace creates a new Firestore client whose collections will be located in the given namespace.
+func OpenNamespace(ctx context.Context, projectID, name string) (_ *Namespace, err error) {
+	defer derrors.Wrap(&err, "OpenNamespace(%q, %q)", projectID, name)
+
+	if name == "" {
+		return nil, errors.New("empty namespace")
+	}
+	client, err := firestore.NewClient(ctx, projectID)
+	if err != nil {
+		return nil, err
+	}
+	return &Namespace{
+		client: client,
+		name:   name,
+		doc:    client.Collection(namespaceCollection).Doc(name),
+	}, nil
+}
+
+// Name returns the Namespace's name.
+func (ns *Namespace) Name() string { return ns.name }
+
+// Client returns the underlying Firestore client.
+func (ns *Namespace) Client() *firestore.Client { return ns.client }
+
+// Close closes the underlying client.
+func (ns *Namespace) Close() error { return ns.client.Close() }
+
+// Collection returns a reference to the named collection in the namespace.
+func (ns *Namespace) Collection(name string) *firestore.CollectionRef {
+	return ns.doc.Collection(name)
+}
+
+// Get gets the DocumentRef and decodes the result to a value of type T.
+func Get[T any](ctx context.Context, dr *firestore.DocumentRef) (_ *T, err error) {
+	defer derrors.Wrap(&err, "fstore.Get(%q)", dr.Path)
+	docsnap, err := dr.Get(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return Decode[T](docsnap)
+}
+
+// Decode decodes a DocumentSnapshot into a value of type T.
+func Decode[T any](ds *firestore.DocumentSnapshot) (*T, error) {
+	var t T
+	if err := ds.DataTo(&t); err != nil {
+		return nil, err
+	}
+	return &t, nil
+}

internal/jobs/firestore.go

@@ -6,42 +6,27 @@ package jobs
 
 import (
 	"context"
-	"errors"
 	"time"
 
 	"cloud.google.com/go/firestore"
 	"golang.org/x/pkgsite-metrics/internal/derrors"
+	"golang.org/x/pkgsite-metrics/internal/fstore"
 	"google.golang.org/api/iterator"
 )
 
-// A DB is a client for a database that stores Jobs.
+const jobCollection = "Jobs"
+
 type DB struct {
-	namespace string
-	client    *firestore.Client
-	nsDoc     *firestore.DocumentRef // the namespace for this db
+	ns *fstore.Namespace
 }
 
-const (
-	namespaceCollection = "Namespaces"
-	jobCollection       = "Jobs"
-)
-
 // NewDB creates a new database client for jobs.
 func NewDB(ctx context.Context, projectID, namespace string) (_ *DB, err error) {
-	defer derrors.Wrap(&err, "job.NewDB(%q, %q)", projectID, namespace)
-
-	if namespace == "" {
-		return nil, errors.New("empty namespace")
-	}
-	client, err := firestore.NewClient(ctx, projectID)
+	ns, err := fstore.OpenNamespace(ctx, projectID, namespace)
 	if err != nil {
 		return nil, err
 	}
-	return &DB{
-		namespace: namespace,
-		client:    client,
-		nsDoc:     client.Collection(namespaceCollection).Doc(namespace),
-	}, nil
+	return &DB{ns}, nil
 }
 
 // CreateJob creates a new job. It returns an error if a job with the same ID already exists.
@@ -62,25 +47,21 @@ func (d *DB) DeleteJob(ctx context.Context, id string) (err error) {
 // GetJob retrieves the job with the given ID. It returns an error if the job does not exist.
 func (d *DB) GetJob(ctx context.Context, id string) (_ *Job, err error) {
 	defer derrors.Wrap(&err, "job.DB.GetJob(%s)", id)
-	docsnap, err := d.jobRef(id).Get(ctx)
-	if err != nil {
-		return nil, err
-	}
-	return docsnapToJob(docsnap)
+	return fstore.Get[Job](ctx, d.jobRef(id))
 }
 
 // UpdateJob gets the job with the given ID, which must exist, then calls f on
 // it, then writes it back to the database. These actions occur atomically.
 // If f returns an error, that error is returned and no update occurs.
 func (d *DB) UpdateJob(ctx context.Context, id string, f func(*Job) error) (err error) {
 	defer derrors.Wrap(&err, "job.DB.UpdateJob(%s)", id)
-	return d.client.RunTransaction(ctx, func(ctx context.Context, tx *firestore.Transaction) error {
+	return d.ns.Client().RunTransaction(ctx, func(ctx context.Context, tx *firestore.Transaction) error {
 		docref := d.jobRef(id)
 		docsnap, err := tx.Get(docref)
 		if err != nil {
 			return err
 		}
-		j, err := docsnapToJob(docsnap)
+		j, err := fstore.Decode[Job](docsnap)
 		if err != nil {
 			return err
 		}
@@ -108,7 +89,7 @@ func (d *DB) Increment(ctx context.Context, id, name string, n int) (err error)
 func (d *DB) ListJobs(ctx context.Context, f func(_ *Job, lastUpdate time.Time) error) (err error) {
 	defer derrors.Wrap(&err, "job.DB.ListJobs()")
 
-	q := d.nsDoc.Collection(jobCollection).OrderBy("StartedAt", firestore.Desc)
+	q := d.ns.Collection(jobCollection).OrderBy("StartedAt", firestore.Desc)
 	iter := q.Documents(ctx)
 	defer iter.Stop()
 	for {
@@ -119,7 +100,7 @@ func (d *DB) ListJobs(ctx context.Context, f func(_ *Job, lastUpdate time.Time)
 		if err != nil {
 			return err
 		}
-		job, err := docsnapToJob(docsnap)
+		job, err := fstore.Decode[Job](docsnap)
 		if err != nil {
 			return err
 		}
@@ -132,14 +113,5 @@ func (d *DB) ListJobs(ctx context.Context, f func(_ *Job, lastUpdate time.Time)
 
 // jobRef returns the DocumentRef for a job with the given ID.
 func (d *DB) jobRef(id string) *firestore.DocumentRef {
-	return d.nsDoc.Collection(jobCollection).Doc(id)
-}
-
-// docsnapToJob converts a DocumentSnapshot to a Job.
-func docsnapToJob(ds *firestore.DocumentSnapshot) (*Job, error) {
-	var j Job
-	if err := ds.DataTo(&j); err != nil {
-		return nil, err
-	}
-	return &j, nil
+	return d.ns.Collection(jobCollection).Doc(id)
 }