Demo - Cloud Bigtable from Go - Basic CRUD/Filter operations!

bigtable.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"cloud.google.com/go/bigtable"
)

var (
	operation   string
	allVersions bool

	customerName   string
	customerEmail  string
	customerDomain string

	ProjectID = ""
	Instance  = "bt-instance-us"
)

const (
	Table                = "Customer"
	CustomerColumnFamily = "personalData"
)

func init() {
	flag.StringVar(&ProjectID, "project", os.Getenv("GOOGLE_CLOUD_PROJECT"), "The Google Cloud Project ID where the instance is located.")
	flag.StringVar(&Instance, "instance", "bt-instance-us", "The Bigtable instance id.")
	flag.StringVar(&operation, "operation", "list", "The opration to be executed: list, insert, delete")
	flag.BoolVar(&allVersions, "all-versions", false, "If the List operation shows all versions")

	flag.StringVar(&customerEmail, "email", "", "The email address of the customer")
	flag.StringVar(&customerDomain, "domain", "", "The domain of the customer")
	flag.StringVar(&customerName, "name", "", "The name of the customer")
}

func main() {
	flag.Parse()

	app := Connect()
	fmt.Println("Successfully created a connection to a Cloud Bigtable instance")

	switch operation {
	case "list":
		app.ListCustomers(customerDomain)
	case "insert":
		if customerEmail == "" {
			log.Println("Setting up sample customers")
			checkError(app.InsertCustomer("Ronoaldo Pereira", "[email protected]"))
			checkError(app.InsertCustomer("Albert Eistain", "[email protected]"))
			checkError(app.InsertCustomer("Isaac Newton", "[email protected]"))
			checkError(app.InsertCustomer("Marie Curie", "[email protected]"))
			checkError(app.InsertCustomer("Nelson Mandela", "[email protected]"))
			return
		}
		checkError(app.InsertCustomer(customerName, customerEmail))
	case "delete":
		if customerEmail == "all" {
			app.RemoveAll()
			return
		}
		checkError(app.RemoveCustomer(customerEmail))
	}
}

type App struct {
	ctx      context.Context
	bigtable *bigtable.Client
}

func Connect() *App {
	ctx := context.Background()
	bt, err := bigtable.NewClient(ctx, ProjectID, Instance)
	if err != nil {
		panic(fmt.Errorf("bigtable.NewAdminClient: %v", err))
	}
	return &App{
		ctx:      ctx,
		bigtable: bt,
	}
}

func NewCustomerKey(email string) string {
	if strings.Count(email, "@") != 1 {
		panic(fmt.Errorf("invalid email: %s", email))
	}
	parts := strings.Split(email, "@")
	_, domain := parts[0], parts[1]
	switch domain {
	case "gmail.com", "yahoo.com", "outlook.com":
		domain = "public"
	}
	return fmt.Sprintf("customer#%s#%s", domain, email)
}

// InsertCustomer adds a new customer to the database
func (app *App) InsertCustomer(name string, email string) error {
	tbl := app.bigtable.Open(Table)

	timestamp := bigtable.Now()
	mut := bigtable.NewMutation()
	mut.Set(CustomerColumnFamily, "name", timestamp, []byte(name))
	mut.Set(CustomerColumnFamily, "email", timestamp, []byte(email))

	rowKey := NewCustomerKey(email)
	log.Printf("Writing row: %s\n", rowKey)
	if err := tbl.Apply(app.ctx, rowKey, mut); err != nil {
		log.Printf("Apply: %v", err)
		return err
	}
	log.Printf("Successfully wrote row: %s\n", rowKey)
	return nil
}

func (app *App) RemoveCustomer(email string) error {
	tbl := app.bigtable.Open(Table)
	rowKey := NewCustomerKey(email)
	mut := bigtable.NewMutation()
	mut.DeleteRow()
	if err := tbl.Apply(app.ctx, rowKey, mut); err != nil {
		log.Printf("DeleteRow: %v", err)
		return err
	}
	log.Printf("Successfully deleted row: %s\n", rowKey)
	return nil
}

func (app *App) RemoveAll() {
	keys := app.CustomerKeys()
	muts := []*bigtable.Mutation{}
	tbl := app.bigtable.Open(Table)

	for i := 0; i < len(keys); i++ {
		mut := bigtable.NewMutation()
		mut.DeleteRow()
		muts = append(muts, mut)
	}

	if _, err := tbl.ApplyBulk(app.ctx, keys, muts); err != nil {
		log.Printf("ApplyBulk: %v", err)
	}

	log.Printf("Successfully deleted all rows\n")
}

func (app *App) ListCustomers(domain string) {
	count := 0
	tbl := app.bigtable.Open(Table)
	var filters []bigtable.ReadOption
	if !allVersions {
		filters = append(filters, bigtable.RowFilter(bigtable.LatestNFilter(1)))
	}
	prefix := bigtable.PrefixRange(fmt.Sprintf("customer#%s", domain))
	callback := func(row bigtable.Row) bool {
		count++
		return printRow(row)
	}
	err := tbl.ReadRows(app.ctx, prefix, callback, filters...)
	if err != nil {
		log.Fatalf("tbl.ReadRows: %v", err)
	}
	fmt.Printf("Found %d customers\n", count)
}

func (app *App) CustomerKeys() (keys []string) {
	tbl := app.bigtable.Open(Table)
	tbl.ReadRows(app.ctx, bigtable.PrefixRange("customer#"), func(row bigtable.Row) bool {
		keys = append(keys, row.Key())
		return true
	}, bigtable.RowFilter(bigtable.StripValueFilter()))
	return keys
}

func printRow(row bigtable.Row) bool {
	fmt.Printf("%s:\n", row.Key())
	for _, cols := range row {
		for _, col := range cols {
			fmt.Printf("  %s=%s@%v\n", col.Column, string(col.Value), col.Timestamp.Time())
		}
	}
	return true
}

func checkError(err error) {
	if err != nil {
		panic(err)
	}
}