Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
70 changes: 52 additions & 18 deletions bigtable/admin.go
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ import (
"strings"
"time"

admin "cloud.google.com/go/bigtable/admin/apiv2"
btapb "cloud.google.com/go/bigtable/admin/apiv2/adminpb"
btopt "cloud.google.com/go/bigtable/internal/option"
"cloud.google.com/go/iam"
Expand Down Expand Up @@ -101,9 +102,10 @@ func (e ErrPartiallyUnavailable) Error() string {

// AdminClient is a client type for performing admin operations within a specific instance.
type AdminClient struct {
connPool gtransport.ConnPool
tClient btapb.BigtableTableAdminClient
lroClient *lroauto.OperationsClient
connPool gtransport.ConnPool
tClient btapb.BigtableTableAdminClient
lroClient *lroauto.OperationsClient
tableAdminClient *admin.BigtableTableAdminClient

project, instance string

Expand Down Expand Up @@ -140,13 +142,19 @@ func NewAdminClient(ctx context.Context, project, instance string, opts ...optio
return nil, err
}

tableAdminClient, err := admin.NewBigtableTableAdminClient(ctx, gtransport.WithConnPool(connPool))
if err != nil {
return nil, err
}

return &AdminClient{
connPool: connPool,
tClient: btapb.NewBigtableTableAdminClient(connPool),
lroClient: lroClient,
project: project,
instance: instance,
md: metadata.Pairs(resourcePrefixHeader, fmt.Sprintf("projects/%s/instances/%s", project, instance)),
connPool: connPool,
tClient: btapb.NewBigtableTableAdminClient(connPool),
lroClient: lroClient,
tableAdminClient: tableAdminClient,
project: project,
instance: instance,
md: metadata.Pairs(resourcePrefixHeader, fmt.Sprintf("projects/%s/instances/%s", project, instance)),
}, nil
}

Expand All @@ -155,6 +163,16 @@ func (ac *AdminClient) Close() error {
return ac.connPool.Close()
}

// TableAdminClientV2 returns the GAPIC generated BigtableTableAdminClient.
//
// The returned client shares the underlying connection pool with AdminClient.
// Since the connection pool is shared, calling Close on either the returned client
// or the parent AdminClient will close the connection pool, making both clients
// unusable.
func (ac *AdminClient) TableAdminClientV2() *admin.BigtableTableAdminClient {
return ac.tableAdminClient
}

func (ac *AdminClient) instancePrefix() string {
return instancePrefix(ac.project, ac.instance)
}
Expand Down Expand Up @@ -1278,9 +1296,10 @@ const mtlsInstanceAdminAddr = "bigtableadmin.mtls.googleapis.com:443"
// InstanceAdminClient is a client type for performing admin operations on instances.
// These operations can be substantially more dangerous than those provided by AdminClient.
type InstanceAdminClient struct {
connPool gtransport.ConnPool
iClient btapb.BigtableInstanceAdminClient
lroClient *lroauto.OperationsClient
connPool gtransport.ConnPool
iClient btapb.BigtableInstanceAdminClient
lroClient *lroauto.OperationsClient
instanceAdminClient *admin.BigtableInstanceAdminClient

project string

Expand Down Expand Up @@ -1313,13 +1332,18 @@ func NewInstanceAdminClient(ctx context.Context, project string, opts ...option.
return nil, err
}

return &InstanceAdminClient{
connPool: connPool,
iClient: btapb.NewBigtableInstanceAdminClient(connPool),
lroClient: lroClient,
instanceAdminClient, err := admin.NewBigtableInstanceAdminClient(ctx, gtransport.WithConnPool(connPool))
if err != nil {
return nil, err
}

project: project,
md: metadata.Pairs(resourcePrefixHeader, "projects/"+project),
return &InstanceAdminClient{
connPool: connPool,
iClient: btapb.NewBigtableInstanceAdminClient(connPool),
lroClient: lroClient,
instanceAdminClient: instanceAdminClient,
project: project,
md: metadata.Pairs(resourcePrefixHeader, "projects/"+project),
}, nil
}

Expand All @@ -1328,6 +1352,16 @@ func (iac *InstanceAdminClient) Close() error {
return iac.connPool.Close()
}

// InstanceAdminClientV2 returns the GAPIC generated BigtableInstanceAdminClient.
//
// The returned client shares the underlying connection pool with InstanceAdminClient.
// Since the connection pool is shared, calling Close on either the returned client
// or the parent InstanceAdminClient will close the connection pool, making both clients
// unusable.
func (iac *InstanceAdminClient) InstanceAdminClientV2() *admin.BigtableInstanceAdminClient {
return iac.instanceAdminClient
}

// StorageType is the type of storage used for all tables in an instance
type StorageType int

Expand Down
72 changes: 72 additions & 0 deletions bigtable/admin/apiv2/table_admin.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package admin

import (
"context"
"time"

adminpb "cloud.google.com/go/bigtable/admin/apiv2/adminpb"
gax "github.com/googleapis/gax-go/v2"
)

// RestoreTable creates a new table by restoring from a backup.
func (c *BigtableTableAdminClient) RestoreTable(ctx context.Context, req *adminpb.RestoreTableRequest, opts ...gax.CallOption) error {
op, err := c.restoreTable(ctx, req, opts...)
if err != nil {
return err
}

// Poll the LRO until the table is usable
if _, err := op.Wait(ctx, opts...); err != nil {
return err
}

return nil
}

// WaitForConsistency waits until all the writes committed before the call started have been propagated to all the clusters in the instance via replication.
func (c *BigtableTableAdminClient) WaitForConsistency(ctx context.Context, tableName string, opts ...gax.CallOption) error {
// Get the token.
tokenResp, err := c.GenerateConsistencyToken(ctx, &adminpb.GenerateConsistencyTokenRequest{
Name: tableName,
}, opts...)
if err != nil {
return err
}
token := tokenResp.GetConsistencyToken()

// Periodically check if the token is consistent.
timer := time.NewTicker(time.Second * 10)
Comment thread
bhshkh marked this conversation as resolved.
defer timer.Stop()
for {
consistentResp, err := c.CheckConsistency(ctx, &adminpb.CheckConsistencyRequest{
Name: tableName,
ConsistencyToken: token,
}, opts...)
if err != nil {
return err
}
if consistentResp.GetConsistent() {
return nil
}
// Sleep for a bit or until the ctx is cancelled.
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
}
}
}
204 changes: 204 additions & 0 deletions bigtable/admin/apiv2/table_admin_integration_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,204 @@
// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package admin

import (
"context"
"fmt"
"os"
"testing"
"time"

adminpb "cloud.google.com/go/bigtable/admin/apiv2/adminpb"
"google.golang.org/api/option"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/status"
"google.golang.org/protobuf/types/known/timestamppb"
)

type testEnv struct {
client *BigtableTableAdminClient
project string
instance string
cluster string
}

func setupIntegration(t *testing.T) *testEnv {
t.Helper()
if testing.Short() {
t.Skip("Skipping integration test in short mode")
}

ctx := context.Background()
emulatorHost := os.Getenv("BIGTABLE_EMULATOR_HOST")
var opts []option.ClientOption
project := os.Getenv("GCLOUD_TESTS_GOLANG_PROJECT_ID")
instance := os.Getenv("GCLOUD_TESTS_BIGTABLE_INSTANCE")
cluster := os.Getenv("GCLOUD_TESTS_BIGTABLE_CLUSTER")

if emulatorHost != "" {
t.Logf("Using emulator at %s", emulatorHost)
opts = append(opts,
option.WithEndpoint(emulatorHost),
option.WithGRPCDialOption(grpc.WithTransportCredentials(insecure.NewCredentials())),
)
if project == "" {
project = "test-project"
}
if instance == "" {
instance = "test-instance"
}
if cluster == "" {
cluster = "test-cluster"
}
} else if project == "" || instance == "" {
t.Skip("Missing GCLOUD_TESTS_GOLANG_PROJECT_ID or GCLOUD_TESTS_BIGTABLE_INSTANCE for non-emulator run")
}

client, err := NewBigtableTableAdminClient(ctx, opts...)
if err != nil {
t.Fatalf("Failed to create client: %v", err)
}
t.Cleanup(func() {
client.Close()
})

return &testEnv{
client: client,
project: project,
instance: instance,
cluster: cluster,
}
}

func TestIntegration_RestoreTable(t *testing.T) {
env := setupIntegration(t)
if os.Getenv("BIGTABLE_EMULATOR_HOST") == "" && env.cluster == "" {
t.Skip("Missing GCLOUD_TESTS_BIGTABLE_CLUSTER for non-emulator run")
}

ctx := context.Background()
client := env.client

suffix := time.Now().Format("20060102-150405")
sourceTableID := fmt.Sprintf("src-table-%s", suffix)
backupID := fmt.Sprintf("backup-%s", suffix)
restoredTableID := fmt.Sprintf("restored-table-%s", suffix)

instancePath := fmt.Sprintf("projects/%s/instances/%s", env.project, env.instance)
sourceTablePath := fmt.Sprintf("%s/tables/%s", instancePath, sourceTableID)
clusterPath := fmt.Sprintf("%s/clusters/%s", instancePath, env.cluster)
backupPath := fmt.Sprintf("%s/backups/%s", clusterPath, backupID)
restoredTablePath := fmt.Sprintf("%s/tables/%s", instancePath, restoredTableID)

// 1. Create source table
_, err := client.CreateTable(ctx, &adminpb.CreateTableRequest{
Parent: instancePath,
TableId: sourceTableID,
Table: &adminpb.Table{},
})
if err != nil {
t.Fatalf("Failed to create source table: %v", err)
}
t.Cleanup(func() {
client.DeleteTable(ctx, &adminpb.DeleteTableRequest{Name: sourceTablePath})
})

// 2. Create backup
expireTime := time.Now().Add(7 * time.Hour)
opCreateBackup, err := client.CreateBackup(ctx, &adminpb.CreateBackupRequest{
Parent: clusterPath,
BackupId: backupID,
Backup: &adminpb.Backup{
SourceTable: sourceTablePath,
ExpireTime: timestamppb.New(expireTime),
},
})
if err != nil {
st, ok := status.FromError(err)
if ok && st.Code() == codes.Unimplemented {
t.Skip("Emulator does not support CreateBackup")
}
t.Fatalf("Failed to initiate backup: %v", err)
}
t.Cleanup(func() {
client.DeleteBackup(ctx, &adminpb.DeleteBackupRequest{Name: backupPath})
})

_, err = opCreateBackup.Wait(ctx)
if err != nil {
t.Fatalf("Backup LRO failed: %v", err)
}

// 3. Restore table
err = client.RestoreTable(ctx, &adminpb.RestoreTableRequest{
Parent: instancePath,
TableId: restoredTableID,
Source: &adminpb.RestoreTableRequest_Backup{
Backup: backupPath,
},
})
if err != nil {
t.Fatalf("RestoreTable failed: %v", err)
}
t.Cleanup(func() {
client.DeleteTable(ctx, &adminpb.DeleteTableRequest{Name: restoredTablePath})
})

// 4. Verify restored table exists
restoredTable, err := client.GetTable(ctx, &adminpb.GetTableRequest{
Name: restoredTablePath,
})
if err != nil {
t.Fatalf("Failed to get restored table: %v", err)
}

if restoredTable.Name != restoredTablePath {
t.Errorf("Expected restored table name %q, got %q", restoredTablePath, restoredTable.Name)
}
}

func TestIntegration_WaitForConsistency(t *testing.T) {
env := setupIntegration(t)

ctx := context.Background()
client := env.client

suffix := time.Now().Format("20060102-150405")
tableID := fmt.Sprintf("repl-test-table-%s", suffix)
instancePath := fmt.Sprintf("projects/%s/instances/%s", env.project, env.instance)
tablePath := fmt.Sprintf("%s/tables/%s", instancePath, tableID)

// Create table
_, err := client.CreateTable(ctx, &adminpb.CreateTableRequest{
Parent: instancePath,
TableId: tableID,
Table: &adminpb.Table{},
})
if err != nil {
t.Fatalf("Failed to create table: %v", err)
}
t.Cleanup(func() {
client.DeleteTable(ctx, &adminpb.DeleteTableRequest{Name: tablePath})
})

// Wait for replication
err = client.WaitForConsistency(ctx, tablePath)
if err != nil {
t.Fatalf("WaitForConsistency failed: %v", err)
}
}
Loading
Loading