Go+Google Cloud Firestore

Go の勉強がてら Google Cloud Datastore へ読み書きする簡単なプログラムを作成したため共有する。
Go は Java と比べてリフレクション周りが癖があって扱いにくかった。
あと、オーバーロードがないのが辛い。

fb.go

package fb

import (
	"context"
	"fmt"
	"reflect"

	"cloud.google.com/go/firestore"
	firebase "firebase.google.com/go"
	"google.golang.org/api/iterator"
)

type Document interface {
	ID() string
	SetID(id string)
	Collection() string
}

func NewApp(ctx context.Context) (*firebase.App, error) {
	return NewAppWithProjectID(ctx, "your project id here.") // FIXME
}

func NewAppWithProjectID(ctx context.Context, projectID string) (*firebase.App, error) {
	conf := &firebase.Config{ProjectID: projectID}
	app, err := firebase.NewApp(ctx, conf)
	if err == nil {
		return app, err
	}

	return nil, err
}

func NewClient(ctx context.Context) (*firestore.Client, error) {
	return NewClientWithProjectID(ctx, "your project id here.") // FIXME
}

func NewClientWithProjectID(ctx context.Context, projectID string) (*firestore.Client, error) {
	app, err := NewAppWithProjectID(ctx, projectID)
	if err != nil {
		return nil, err
	}

	client, err := app.Firestore(ctx)
	if err == nil {
		return client, err
	}

	return nil, err
}

func GetDoc(ctx context.Context, client *firestore.Client, doc Document) error {
	var dsnap *firestore.DocumentSnapshot
	var err error

	if client == nil {
		dsnap, err = Client.Collection(doc.Collection()).Doc(doc.ID()).Get(ctx)
	} else {
		dsnap, err = client.Collection(doc.Collection()).Doc(doc.ID()).Get(ctx)
	}

	if err != nil {
		return err
	}

	err = dsnap.DataTo(doc)
	if err != nil {
		return err
	}

	return nil
}

func GetDocs(ctx context.Context, q firestore.Query, dst interface{}) error {
	var elemIsPtr bool
	slice := reflect.ValueOf(dst).Elem() // dv はスライスになる
	elemType := slice.Type().Elem()      // elemType はスライス slice の要素の型になる

	switch elemType.Kind() {
	case reflect.Struct:
		elemIsPtr = false
	case reflect.Ptr:
		elemIsPtr = true
		elemType = elemType.Elem()
	default:
		return fmt.Errorf("unsupported slice element type: %v", elemType)
	}

	iter := q.Documents(ctx)

	for {
		doc, err := iter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return err
		}

		entity := reflect.New(elemType)
		entity.Interface().(Document).SetID(doc.Ref.ID)
		err = doc.DataTo(entity.Interface())
		if err != nil {
			return err
		}

		if elemIsPtr {
			slice.Set(reflect.Append(slice, entity))
		} else {
			slice.Set(reflect.Append(slice, entity.Elem()))
		}
	}

	return nil
}

func PutDoc(ctx context.Context, client *firestore.Client, doc Document) error {
	var ref *firestore.DocumentRef
	var err error

	if doc.ID() == "" {
		if client == nil {
			ref, _, err = Client.Collection(doc.Collection()).Add(ctx, doc)
		} else {
			ref, _, err = client.Collection(doc.Collection()).Add(ctx, doc)
		}
	} else {
		if client == nil {
			_, err = Client.Collection(doc.Collection()).Doc(doc.ID()).Set(ctx, doc)
		} else {
			_, err = client.Collection(doc.Collection()).Doc(doc.ID()).Set(ctx, doc)
		}
	}

	if err != nil {
		return err
	}

	if ref != nil {
		doc.SetID(ref.ID)
	}

	return nil
}

func PutDocs(ctx context.Context, client *firestore.Client, entities interface{}) error {
	var elemIsPtr bool
	slice := reflect.ValueOf(entities)
	elemType := slice.Type().Elem() // elemType はスライス slice の要素の型になる

	switch elemType.Kind() {
	case reflect.Struct:
		elemIsPtr = false
	case reflect.Ptr:
		elemIsPtr = true
		elemType = elemType.Elem()
	default:
		return fmt.Errorf("unsupported slice element type: %v", elemType)
	}

	var ref *firestore.CollectionRef
	var batch *firestore.WriteBatch

	if client == nil {
		ref = Client.Collection(reflect.New(elemType).Interface().(Document).Collection())
		batch = Client.Batch()
	} else {
		ref = client.Collection(reflect.New(elemType).Interface().(Document).Collection())
		batch = client.Batch()
	}

	for i := 0; i < slice.Len(); i++ {
		var doc Document
		if elemIsPtr {
			doc = slice.Index(i).Interface().(Document)
		} else {
			doc = slice.Index(i).Addr().Interface().(Document)
		}

		if doc.ID() == "" {
			newDoc := ref.NewDoc()
			_ = batch.Set(newDoc, doc)
			doc.SetID(newDoc.ID)
		} else {
			_ = batch.Set(ref.Doc(doc.ID()), doc)
		}
	}

	_, err := batch.Commit(ctx)
	if err != nil {
		return err
	}

	return nil
}

func DeleteDoc(ctx context.Context, client *firestore.Client, doc Document) error {
	var err error
	if client == nil {
		_, err = Client.Collection(doc.Collection()).Doc(doc.ID()).Delete(ctx)
	} else {
		_, err = client.Collection(doc.Collection()).Doc(doc.ID()).Delete(ctx)
	}

	if err != nil {
		return err
	}

	return nil
}

func DeleteDocs(ctx context.Context, client *firestore.Client, entities interface{}) error {
	var elemIsPtr bool
	slice := reflect.ValueOf(entities)
	elemType := slice.Type().Elem() // elemType はスライス slice の要素の型になる

	switch elemType.Kind() {
	case reflect.Struct:
		elemIsPtr = false
	case reflect.Ptr:
		elemIsPtr = true
		elemType = elemType.Elem()
	default:
		return fmt.Errorf("unsupported slice element type: %v", elemType)
	}

	var ref *firestore.CollectionRef
	var batch *firestore.WriteBatch

	if client == nil {
		ref = Client.Collection(reflect.New(elemType).Interface().(Document).Collection())
		batch = Client.Batch()
	} else {
		ref = client.Collection(reflect.New(elemType).Interface().(Document).Collection())
		batch = client.Batch()
	}

	for i := 0; i < slice.Len(); i++ {
		var doc Document
		if elemIsPtr {
			doc = slice.Index(i).Interface().(Document)
		} else {
			doc = slice.Index(i).Addr().Interface().(Document)
		}
		_ = batch.Delete(ref.Doc(doc.ID()))
	}

	_, err := batch.Commit(ctx)
	if err != nil {
		return err
	}

	return nil
}