Skip to content

File tree

bigtable/admin.go

Lines changed: 90 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,26 @@ type TableAutomatedBackupPolicy struct {
316316

317317
func (*TableAutomatedBackupPolicy) isTableAutomatedBackupConfig() {}
318318

319+
// TieredStorageConfig defines a tiered storage configuration for a table.
320+
type TieredStorageConfig struct {
321+
// Rule to specify what data is stored in the infrequent access(IA) tier.
322+
// The IA tier allows storing more data per node with reduced performance.
323+
InfrequentAccess TieredStorageRule
324+
}
325+
326+
// TieredStorageRule defines a tiered storage rule for a table.
327+
type TieredStorageRule interface {
328+
isTieredStorageRule()
329+
}
330+
331+
// TieredStorageIncludeIfOlderThan defines a tiered storage rule that includes
332+
// cells older than the given age.
333+
type TieredStorageIncludeIfOlderThan struct {
334+
Duration optional.Duration
335+
}
336+
337+
func (*TieredStorageIncludeIfOlderThan) isTieredStorageRule() {}
338+
319339
func toAutomatedBackupConfigProto(automatedBackupConfig TableAutomatedBackupConfig) (*btapb.Table_AutomatedBackupPolicy_, error) {
320340
if automatedBackupConfig == nil {
321341
return nil, nil
@@ -347,6 +367,26 @@ func (abp *TableAutomatedBackupPolicy) toProto() (*btapb.Table_AutomatedBackupPo
347367
}, nil
348368
}
349369

370+
func (tsc *TieredStorageConfig) toProto() (*btapb.TieredStorageConfig, error) {
371+
if tsc == nil {
372+
return nil, nil
373+
}
374+
pb := &btapb.TieredStorageConfig{}
375+
if tsc.InfrequentAccess != nil {
376+
switch rule := tsc.InfrequentAccess.(type) {
377+
case *TieredStorageIncludeIfOlderThan:
378+
pb.InfrequentAccess = &btapb.TieredStorageRule{
379+
Rule: &btapb.TieredStorageRule_IncludeIfOlderThan{
380+
IncludeIfOlderThan: durationpb.New(optional.ToDuration(rule.Duration)),
381+
},
382+
}
383+
default:
384+
return nil, fmt.Errorf("bigtable: unknown tiered storage rule type: %T", rule)
385+
}
386+
}
387+
return pb, nil
388+
}
389+
350390
// Family represents a column family with its optional GC policy and value type.
351391
type Family struct {
352392
GCPolicy GCPolicy
@@ -371,10 +411,12 @@ type TableConf struct {
371411
// set to protected to make the table protected against data loss
372412
DeletionProtection DeletionProtection
373413
ChangeStreamRetention ChangeStreamRetention
374-
// Configure an automated backup policy for the table
414+
// AutomatedBackupConfig defines the automated backup policy for the table.
375415
AutomatedBackupConfig TableAutomatedBackupConfig
376-
// Configure a row key schema for the table
416+
// RowKeySchema describes the structure of the row keys in the table.
377417
RowKeySchema *StructType
418+
// TieredStorageConfig represents rules for the table to specify what data is stored in each storage tier.
419+
TieredStorageConfig *TieredStorageConfig
378420
}
379421

380422
// CreateTable creates a new table in the instance.
@@ -427,6 +469,14 @@ func (ac *AdminClient) CreateTableFromConf(ctx context.Context, conf *TableConf)
427469
tbl.RowKeySchema = conf.RowKeySchema.proto().GetStructType()
428470
}
429471

472+
if conf.TieredStorageConfig != nil {
473+
proto, err := conf.TieredStorageConfig.toProto()
474+
if err != nil {
475+
return err
476+
}
477+
tbl.TieredStorageConfig = proto
478+
}
479+
430480
if conf.Families != nil && conf.ColumnFamilies != nil {
431481
return errors.New("only one of Families or ColumnFamilies may be set, not both")
432482
}
@@ -512,6 +562,7 @@ const (
512562
retentionPeriodFieldMaskPath = "retention_period"
513563
frequencyFieldMaskPath = "frequency"
514564
rowKeySchemaMaskPath = "row_key_schema"
565+
tieredStorageConfigFieldMask = "tiered_storage_config"
515566
)
516567

517568
func (ac *AdminClient) newUpdateTableRequestProto(tableID string) (*btapb.UpdateTableRequest, error) {
@@ -641,6 +692,31 @@ func (ac *AdminClient) UpdateTableRemoveRowKeySchema(ctx context.Context, tableI
641692
return ac.updateTableAndWait(ctx, req)
642693
}
643694

695+
// UpdateTableWithTieredStorageConfig updates a table with TieredStorageConfig.
696+
func (ac *AdminClient) UpdateTableWithTieredStorageConfig(ctx context.Context, tableID string, tieredStorageConfig *TieredStorageConfig) error {
697+
req, err := ac.newUpdateTableRequestProto(tableID)
698+
if err != nil {
699+
return err
700+
}
701+
proto, err := tieredStorageConfig.toProto()
702+
if err != nil {
703+
return err
704+
}
705+
req.UpdateMask.Paths = append(req.UpdateMask.Paths, tieredStorageConfigFieldMask)
706+
req.Table.TieredStorageConfig = proto
707+
return ac.updateTableAndWait(ctx, req)
708+
}
709+
710+
// UpdateTableRemoveTieredStorageConfig removes a TieredStorageConfig from a table.
711+
func (ac *AdminClient) UpdateTableRemoveTieredStorageConfig(ctx context.Context, tableID string) error {
712+
req, err := ac.newUpdateTableRequestProto(tableID)
713+
if err != nil {
714+
return err
715+
}
716+
req.UpdateMask.Paths = append(req.UpdateMask.Paths, tieredStorageConfigFieldMask)
717+
return ac.updateTableAndWait(ctx, req)
718+
}
719+
644720
// DeleteTable deletes a table and all of its data.
645721
func (ac *AdminClient) DeleteTable(ctx context.Context, table string) error {
646722
ctx = mergeOutgoingMetadata(ctx, ac.md)
@@ -679,6 +755,7 @@ type TableInfo struct {
679755
ChangeStreamRetention ChangeStreamRetention
680756
AutomatedBackupConfig TableAutomatedBackupConfig
681757
RowKeySchema *StructType
758+
TieredStorageConfig *TieredStorageConfig
682759
}
683760

684761
// FamilyInfo represents information about a column family.
@@ -754,6 +831,17 @@ func (ac *AdminClient) TableInfo(ctx context.Context, table string) (*TableInfo,
754831
structType := structProtoToType(res.RowKeySchema).(StructType)
755832
ti.RowKeySchema = &structType
756833
}
834+
if res.TieredStorageConfig != nil {
835+
ti.TieredStorageConfig = &TieredStorageConfig{}
836+
if res.TieredStorageConfig.InfrequentAccess != nil {
837+
switch rule := res.TieredStorageConfig.InfrequentAccess.Rule.(type) {
838+
case *btapb.TieredStorageRule_IncludeIfOlderThan:
839+
ti.TieredStorageConfig.InfrequentAccess = &TieredStorageIncludeIfOlderThan{
840+
Duration: rule.IncludeIfOlderThan.AsDuration(),
841+
}
842+
}
843+
}
844+
}
757845

758846
return ti, nil
759847
}

bigtable/admin_test.go

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import (
3030
"github.com/google/go-cmp/cmp"
3131
"google.golang.org/grpc"
3232
"google.golang.org/protobuf/types/known/anypb"
33+
"google.golang.org/protobuf/types/known/durationpb"
3334
"google.golang.org/protobuf/types/known/fieldmaskpb"
3435
"google.golang.org/protobuf/types/known/timestamppb"
3536
)
@@ -51,6 +52,16 @@ type mockTableAdminClock struct {
5152
createAuthorizedViewError error
5253
updateAuthorizedViewReq *btapb.UpdateAuthorizedViewRequest
5354
updateAuthorizedViewError error
55+
56+
getTableReq *btapb.GetTableRequest
57+
getTableResp *btapb.Table
58+
}
59+
60+
func (c *mockTableAdminClock) GetTable(
61+
ctx context.Context, in *btapb.GetTableRequest, opts ...grpc.CallOption,
62+
) (*btapb.Table, error) {
63+
c.getTableReq = in
64+
return c.getTableResp, nil
5465
}
5566

5667
func (c *mockTableAdminClock) CreateTable(
@@ -207,6 +218,37 @@ func TestTableAdmin_CreateTableFromConf_AutomatedBackupPolicy_Valid(t *testing.T
207218
}
208219
}
209220

221+
func TestTableAdmin_CreateTableFromConf_TieredStorageConfig_Valid(t *testing.T) {
222+
mock := &mockTableAdminClock{}
223+
c := setupTableClient(t, mock)
224+
225+
tieredStorageConfig := TieredStorageConfig{
226+
InfrequentAccess: &TieredStorageIncludeIfOlderThan{Duration: 30 * 24 * time.Hour},
227+
}
228+
229+
err := c.CreateTableFromConf(context.Background(), &TableConf{
230+
TableID: "My-table",
231+
TieredStorageConfig: &tieredStorageConfig,
232+
})
233+
if err != nil {
234+
t.Fatalf("CreateTableFromConf failed: %v", err)
235+
}
236+
createTableReq := mock.createTableReq
237+
if !cmp.Equal(createTableReq.TableId, "My-table") {
238+
t.Errorf("Unexpected table ID: %v, expected %v", createTableReq.TableId, "My-table")
239+
}
240+
got := createTableReq.Table.TieredStorageConfig
241+
if got == nil {
242+
t.Fatal("TieredStorageConfig is nil")
243+
}
244+
if got.InfrequentAccess == nil {
245+
t.Fatal("InfrequentAccess is nil")
246+
}
247+
if got.InfrequentAccess.GetIncludeIfOlderThan().AsDuration() != 30*24*time.Hour {
248+
t.Errorf("Unexpected IncludeIfOlderThan: %v, expected %v", got.InfrequentAccess.GetIncludeIfOlderThan().AsDuration(), 30*24*time.Hour)
249+
}
250+
}
251+
210252
func TestTableAdmin_CreateTableFromConf_WithRowKeySchema(t *testing.T) {
211253
mock := &mockTableAdminClock{}
212254
c := setupTableClient(t, mock)
@@ -758,6 +800,88 @@ func TestTableAdmin_UpdateTableWithAutomatedBackupPolicy_ZeroFields_Invalid(t *t
758800
}
759801
}
760802

803+
func TestTableAdmin_UpdateTableWithTieredStorageConfig(t *testing.T) {
804+
mock := &mockTableAdminClock{}
805+
c := setupTableClient(t, mock)
806+
807+
tieredStorageConfig := TieredStorageConfig{
808+
InfrequentAccess: &TieredStorageIncludeIfOlderThan{Duration: 30 * 24 * time.Hour},
809+
}
810+
811+
err := c.UpdateTableWithTieredStorageConfig(context.Background(), "My-table", &tieredStorageConfig)
812+
if err != nil {
813+
t.Fatalf("UpdateTableWithTieredStorageConfig failed: %v", err)
814+
}
815+
816+
updateTableReq := mock.updateTableReq
817+
if updateTableReq.UpdateMask.Paths[0] != tieredStorageConfigFieldMask {
818+
t.Errorf("Unexpected update mask: %v, expected %v", updateTableReq.UpdateMask.Paths[0], tieredStorageConfigFieldMask)
819+
}
820+
got := updateTableReq.Table.TieredStorageConfig
821+
if got == nil {
822+
t.Fatal("TieredStorageConfig is nil")
823+
}
824+
if got.InfrequentAccess == nil {
825+
t.Fatal("InfrequentAccess is nil")
826+
}
827+
if got.InfrequentAccess.GetIncludeIfOlderThan().AsDuration() != 30*24*time.Hour {
828+
t.Errorf("Unexpected IncludeIfOlderThan: %v, expected %v", got.InfrequentAccess.GetIncludeIfOlderThan().AsDuration(), 30*24*time.Hour)
829+
}
830+
}
831+
832+
func TestTableAdmin_TableInfo_TieredStorageConfig(t *testing.T) {
833+
mock := &mockTableAdminClock{}
834+
c := setupTableClient(t, mock)
835+
836+
mock.getTableResp = &btapb.Table{
837+
Name: "projects/my-cool-project/instances/my-cool-instance/tables/My-table",
838+
TieredStorageConfig: &btapb.TieredStorageConfig{
839+
InfrequentAccess: &btapb.TieredStorageRule{
840+
Rule: &btapb.TieredStorageRule_IncludeIfOlderThan{
841+
IncludeIfOlderThan: durationpb.New(30 * 24 * time.Hour),
842+
},
843+
},
844+
},
845+
}
846+
847+
ti, err := c.TableInfo(context.Background(), "My-table")
848+
if err != nil {
849+
t.Fatalf("TableInfo failed: %v", err)
850+
}
851+
852+
if ti.TieredStorageConfig == nil {
853+
t.Fatal("TieredStorageConfig is nil")
854+
}
855+
if ti.TieredStorageConfig.InfrequentAccess == nil {
856+
t.Fatal("InfrequentAccess is nil")
857+
}
858+
rule, ok := ti.TieredStorageConfig.InfrequentAccess.(*TieredStorageIncludeIfOlderThan)
859+
if !ok {
860+
t.Fatalf("Unexpected rule type: %T", ti.TieredStorageConfig.InfrequentAccess)
861+
}
862+
if optional.ToDuration(rule.Duration) != 30*24*time.Hour {
863+
t.Errorf("Unexpected IncludeIfOlderThan: %v, expected %v", optional.ToDuration(rule.Duration), 30*24*time.Hour)
864+
}
865+
}
866+
867+
func TestTableAdmin_UpdateTableRemoveTieredStorageConfig(t *testing.T) {
868+
mock := &mockTableAdminClock{}
869+
c := setupTableClient(t, mock)
870+
871+
err := c.UpdateTableRemoveTieredStorageConfig(context.Background(), "My-table")
872+
if err != nil {
873+
t.Fatalf("UpdateTableRemoveTieredStorageConfig failed: %v", err)
874+
}
875+
876+
updateTableReq := mock.updateTableReq
877+
if updateTableReq.UpdateMask.Paths[0] != tieredStorageConfigFieldMask {
878+
t.Errorf("Unexpected update mask: %v, expected %v", updateTableReq.UpdateMask.Paths[0], tieredStorageConfigFieldMask)
879+
}
880+
if updateTableReq.Table.TieredStorageConfig != nil {
881+
t.Errorf("TieredStorageConfig should be nil")
882+
}
883+
}
884+
761885
type mockAdminClock struct {
762886
btapb.BigtableInstanceAdminClient
763887

bigtable/integration_test.go

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,86 @@ func cleanup(c IntegrationTestConfig) error {
188188
return nil
189189
}
190190

191+
func TestIntegration_TieredStorage(t *testing.T) {
192+
ctx := context.Background()
193+
testEnv, _, adminClient, _, _, cleanup, err := setupIntegration(ctx, t)
194+
if err != nil {
195+
t.Fatal(err)
196+
}
197+
defer cleanup()
198+
199+
if !testEnv.Config().UseProd {
200+
t.Skip("emulator doesn't support TieredStorage")
201+
}
202+
203+
tableName := tableNameSpace.New()
204+
conf := &TableConf{
205+
TableID: tableName,
206+
TieredStorageConfig: &TieredStorageConfig{
207+
InfrequentAccess: &TieredStorageIncludeIfOlderThan{
208+
Duration: 30 * 24 * time.Hour,
209+
},
210+
},
211+
}
212+
213+
if err := adminClient.CreateTableFromConf(ctx, conf); err != nil {
214+
t.Fatalf("CreateTableFromConf failed: %v", err)
215+
}
216+
defer adminClient.DeleteTable(ctx, tableName)
217+
218+
ti, err := adminClient.TableInfo(ctx, tableName)
219+
if err != nil {
220+
t.Fatalf("TableInfo failed: %v", err)
221+
}
222+
223+
if ti.TieredStorageConfig == nil {
224+
t.Fatal("TieredStorageConfig is nil")
225+
}
226+
rule, ok := ti.TieredStorageConfig.InfrequentAccess.(*TieredStorageIncludeIfOlderThan)
227+
if !ok {
228+
t.Fatalf("Unexpected rule type: %T", ti.TieredStorageConfig.InfrequentAccess)
229+
}
230+
if optional.ToDuration(rule.Duration) != 30*24*time.Hour {
231+
t.Errorf("Unexpected IncludeIfOlderThan: %v, expected %v", optional.ToDuration(rule.Duration), 30*24*time.Hour)
232+
}
233+
234+
// Update tiered storage config
235+
newDuration := 45 * 24 * time.Hour
236+
newConfig := TieredStorageConfig{
237+
InfrequentAccess: &TieredStorageIncludeIfOlderThan{
238+
Duration: newDuration,
239+
},
240+
}
241+
if err := adminClient.UpdateTableWithTieredStorageConfig(ctx, tableName, &newConfig); err != nil {
242+
t.Fatalf("UpdateTableWithTieredStorageConfig failed: %v", err)
243+
}
244+
245+
ti, err = adminClient.TableInfo(ctx, tableName)
246+
if err != nil {
247+
t.Fatalf("TableInfo failed after update: %v", err)
248+
}
249+
rule, ok = ti.TieredStorageConfig.InfrequentAccess.(*TieredStorageIncludeIfOlderThan)
250+
if !ok {
251+
t.Fatalf("Unexpected rule type after update: %T", ti.TieredStorageConfig.InfrequentAccess)
252+
}
253+
if optional.ToDuration(rule.Duration) != newDuration {
254+
t.Errorf("Unexpected IncludeIfOlderThan after update: %v, expected %v", optional.ToDuration(rule.Duration), newDuration)
255+
}
256+
257+
// Remove tiered storage config
258+
if err := adminClient.UpdateTableRemoveTieredStorageConfig(ctx, tableName); err != nil {
259+
t.Fatalf("UpdateTableRemoveTieredStorageConfig failed: %v", err)
260+
}
261+
262+
ti, err = adminClient.TableInfo(ctx, tableName)
263+
if err != nil {
264+
t.Fatalf("TableInfo failed after removal: %v", err)
265+
}
266+
if ti.TieredStorageConfig != nil {
267+
t.Errorf("TieredStorageConfig should be nil after removal, got %+v", ti.TieredStorageConfig)
268+
}
269+
}
270+
191271
func TestIntegration_Pinger(t *testing.T) {
192272
ctx := context.Background()
193273
testEnv, client, _, _, _, cleanup, err := setupIntegration(ctx, t)

0 commit comments

Comments
 (0)