| package bt |
| |
| import ( |
| "context" |
| |
| "cloud.google.com/go/bigtable" |
| "go.skia.org/infra/go/sklog" |
| "go.skia.org/infra/go/util" |
| "google.golang.org/grpc/codes" |
| "google.golang.org/grpc/status" |
| ) |
| |
| // TableConfig maps a table name to a list of column families, describing which |
| // tables and column InitBigtable should create. |
| type TableConfig map[string][]string |
| |
| // InitBigtable takes a list of TableConfigs and creates the given tables and |
| // column families if they don't exist already. |
| func InitBigtable(projectID, instanceID string, tableConfigs ...TableConfig) error { |
| ctx := context.TODO() |
| |
| // Set up admin client, tables, and column families. |
| adminClient, err := bigtable.NewAdminClient(ctx, projectID, instanceID) |
| if err != nil { |
| return sklog.FmtErrorf("Unable to create admin client: %s", err) |
| } |
| |
| for _, tConfig := range tableConfigs { |
| for tableName, colFamilies := range tConfig { |
| // Create the table. Ignore error if it already existed. |
| err, code := ErrToCode(adminClient.CreateTable(ctx, tableName)) |
| if err != nil && code != codes.AlreadyExists { |
| return sklog.FmtErrorf("Error creating table %s: %s", err) |
| } else { |
| sklog.Infof("Created table: %s", tableName) |
| } |
| |
| // Create the column families. Ignore errors if they already existed. |
| for _, colFamName := range colFamilies { |
| err, code = ErrToCode(adminClient.CreateColumnFamily(ctx, tableName, colFamName)) |
| if err != nil && code != codes.AlreadyExists { |
| return sklog.FmtErrorf("Error creating column family %s in table %s: %s", colFamName, tableName, err) |
| } |
| } |
| } |
| } |
| return nil |
| } |
| |
| // DeleteTables deletes the tables given in the TableConfig. |
| func DeleteTables(projectID, instanceID string, tableConfigs ...TableConfig) (err error) { |
| ctx := context.TODO() |
| |
| // Set up admin client, tables, and column families. |
| adminClient, err := bigtable.NewAdminClient(ctx, projectID, instanceID) |
| if err != nil { |
| return sklog.FmtErrorf("Unable to create admin client: %s", err) |
| } |
| defer func() { |
| if err != nil { |
| util.Close(adminClient) |
| } else { |
| err = adminClient.Close() |
| } |
| }() |
| |
| // Delete all tables if they exist. |
| for _, tConfig := range tableConfigs { |
| for tableName := range tConfig { |
| // Ignore NotFound errors. |
| err, code := ErrToCode(adminClient.DeleteTable(ctx, tableName)) |
| if err != nil && code != codes.NotFound { |
| return err |
| } |
| } |
| } |
| return nil |
| } |
| |
| // ErrToCode returns the error that is passed and a gRPC code extracted from the error. |
| // If the error did not originate in gRPC the returned code is codes.Unknown. |
| // See https://godoc.org/google.golang.org/grpc/codes for a list of codes. |
| func ErrToCode(err error) (error, codes.Code) { |
| st, _ := status.FromError(err) |
| return err, st.Code() |
| } |