diff --git a/bigtable/admin.go b/bigtable/admin.go index f569c0b757be..e8068d3766b0 100644 --- a/bigtable/admin.go +++ b/bigtable/admin.go @@ -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" @@ -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 @@ -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 } @@ -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) } @@ -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 @@ -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 } @@ -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 diff --git a/bigtable/admin/apiv2/table_admin.go b/bigtable/admin/apiv2/table_admin.go new file mode 100644 index 000000000000..bae54af58610 --- /dev/null +++ b/bigtable/admin/apiv2/table_admin.go @@ -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) + 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: + } + } +} diff --git a/bigtable/admin/apiv2/table_admin_integration_test.go b/bigtable/admin/apiv2/table_admin_integration_test.go new file mode 100644 index 000000000000..97a858ea99b6 --- /dev/null +++ b/bigtable/admin/apiv2/table_admin_integration_test.go @@ -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) + } +} diff --git a/bigtable/admin/apiv2/table_admin_test.go b/bigtable/admin/apiv2/table_admin_test.go new file mode 100644 index 000000000000..97df502ead17 --- /dev/null +++ b/bigtable/admin/apiv2/table_admin_test.go @@ -0,0 +1,278 @@ +// 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" + "net" + "testing" + + adminpb "cloud.google.com/go/bigtable/admin/apiv2/adminpb" + longrunningpb "cloud.google.com/go/longrunning/autogen/longrunningpb" + "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/grpc/test/bufconn" + "google.golang.org/protobuf/types/known/anypb" +) + +type mockServer struct { + adminpb.UnimplementedBigtableTableAdminServer + longrunningpb.UnimplementedOperationsServer + + restoreTableFunc func(context.Context, *adminpb.RestoreTableRequest) (*longrunningpb.Operation, error) + getOperationFunc func(context.Context, *longrunningpb.GetOperationRequest) (*longrunningpb.Operation, error) + generateConsistencyTokenFunc func(context.Context, *adminpb.GenerateConsistencyTokenRequest) (*adminpb.GenerateConsistencyTokenResponse, error) + checkConsistencyFunc func(context.Context, *adminpb.CheckConsistencyRequest) (*adminpb.CheckConsistencyResponse, error) +} + +func (m *mockServer) RestoreTable(ctx context.Context, req *adminpb.RestoreTableRequest) (*longrunningpb.Operation, error) { + if m.restoreTableFunc != nil { + return m.restoreTableFunc(ctx, req) + } + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockServer) GetOperation(ctx context.Context, req *longrunningpb.GetOperationRequest) (*longrunningpb.Operation, error) { + if m.getOperationFunc != nil { + return m.getOperationFunc(ctx, req) + } + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockServer) GenerateConsistencyToken(ctx context.Context, req *adminpb.GenerateConsistencyTokenRequest) (*adminpb.GenerateConsistencyTokenResponse, error) { + if m.generateConsistencyTokenFunc != nil { + return m.generateConsistencyTokenFunc(ctx, req) + } + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func (m *mockServer) CheckConsistency(ctx context.Context, req *adminpb.CheckConsistencyRequest) (*adminpb.CheckConsistencyResponse, error) { + if m.checkConsistencyFunc != nil { + return m.checkConsistencyFunc(ctx, req) + } + return nil, status.Error(codes.Unimplemented, "not implemented") +} + +func setupMockClient(t *testing.T, mock *mockServer) (*BigtableTableAdminClient, func()) { + const bufSize = 1024 * 1024 + lis := bufconn.Listen(bufSize) + + s := grpc.NewServer() + adminpb.RegisterBigtableTableAdminServer(s, mock) + longrunningpb.RegisterOperationsServer(s, mock) + + go func() { + if err := s.Serve(lis); err != nil && err != grpc.ErrServerStopped { + panic(fmt.Sprintf("Server exited with error: %v", err)) + } + }() + + dialer := func(context.Context, string) (net.Conn, error) { + return lis.Dial() + } + + ctx := context.Background() + conn, err := grpc.DialContext(ctx, "bufnet", + grpc.WithContextDialer(dialer), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + t.Fatalf("Failed to dial bufnet: %v", err) + } + + client, err := NewBigtableTableAdminClient(ctx, option.WithGRPCConn(conn)) + if err != nil { + t.Fatalf("Failed to create admin client: %v", err) + } + + cleanup := func() { + client.Close() + s.Stop() + lis.Close() + } + + return client, cleanup +} + +func TestRestoreTable_Success(t *testing.T) { + mock := &mockServer{} + const opName = "operations/restore-test-op" + + // 1. RestoreTable returns a pending LRO + mock.restoreTableFunc = func(ctx context.Context, req *adminpb.RestoreTableRequest) (*longrunningpb.Operation, error) { + return &longrunningpb.Operation{ + Name: opName, + Done: false, + }, nil + } + + // 2. GetOperation returns Done: true on second call + getCalls := 0 + mock.getOperationFunc = func(ctx context.Context, req *longrunningpb.GetOperationRequest) (*longrunningpb.Operation, error) { + if req.GetName() != opName { + return nil, status.Errorf(codes.NotFound, "operation not found: %s", req.GetName()) + } + getCalls++ + if getCalls < 2 { + return &longrunningpb.Operation{ + Name: opName, + Done: false, + }, nil + } + + table := &adminpb.Table{ + Name: "projects/p/instances/i/tables/restored-table", + } + anyTable, err := anypb.New(table) + if err != nil { + return nil, err + } + + return &longrunningpb.Operation{ + Name: opName, + Done: true, + Result: &longrunningpb.Operation_Response{ + Response: anyTable, + }, + }, nil + } + + client, cleanup := setupMockClient(t, mock) + defer cleanup() + + req := &adminpb.RestoreTableRequest{ + Parent: "projects/p/instances/i", + TableId: "restored-table", + } + + err := client.RestoreTable(context.Background(), req) + if err != nil { + t.Fatalf("RestoreTable failed: %v", err) + } + + if getCalls != 2 { + t.Errorf("Expected 2 GetOperation calls, got %d", getCalls) + } +} + +func TestRestoreTable_Error(t *testing.T) { + mock := &mockServer{} + const opName = "operations/restore-test-op" + + mock.restoreTableFunc = func(ctx context.Context, req *adminpb.RestoreTableRequest) (*longrunningpb.Operation, error) { + return &longrunningpb.Operation{ + Name: opName, + Done: false, + }, nil + } + + mock.getOperationFunc = func(ctx context.Context, req *longrunningpb.GetOperationRequest) (*longrunningpb.Operation, error) { + return &longrunningpb.Operation{ + Name: opName, + Done: true, + Result: &longrunningpb.Operation_Error{ + Error: status.New(codes.Aborted, "restore aborted by server").Proto(), + }, + }, nil + } + + client, cleanup := setupMockClient(t, mock) + defer cleanup() + + req := &adminpb.RestoreTableRequest{ + Parent: "projects/p/instances/i", + TableId: "restored-table", + } + + err := client.RestoreTable(context.Background(), req) + if err == nil { + t.Fatal("RestoreTable succeeded, wanted LRO failure error") + } + + st, ok := status.FromError(err) + if !ok { + t.Fatalf("Expected gRPC status error, got %v", err) + } + if st.Code() != codes.Aborted { + t.Errorf("RestoreTable error code: %v, want %v", st.Code(), codes.Aborted) + } +} + +func TestWaitForConsistency_Success(t *testing.T) { + mock := &mockServer{} + const token = "test-consistency-token" + + mock.generateConsistencyTokenFunc = func(ctx context.Context, req *adminpb.GenerateConsistencyTokenRequest) (*adminpb.GenerateConsistencyTokenResponse, error) { + return &adminpb.GenerateConsistencyTokenResponse{ + ConsistencyToken: token, + }, nil + } + + checkCalls := 0 + mock.checkConsistencyFunc = func(ctx context.Context, req *adminpb.CheckConsistencyRequest) (*adminpb.CheckConsistencyResponse, error) { + if req.GetConsistencyToken() != token { + return nil, status.Error(codes.InvalidArgument, "invalid token") + } + checkCalls++ + return &adminpb.CheckConsistencyResponse{ + Consistent: true, + }, nil + } + + client, cleanup := setupMockClient(t, mock) + defer cleanup() + + err := client.WaitForConsistency(context.Background(), "projects/p/instances/i/tables/t") + if err != nil { + t.Fatalf("WaitForConsistency failed: %v", err) + } + + if checkCalls != 1 { + t.Errorf("Expected 1 CheckConsistency call, got %d", checkCalls) + } +} + +func TestWaitForConsistency_ContextCancelled(t *testing.T) { + mock := &mockServer{} + const token = "test-consistency-token" + + ctx, cancel := context.WithCancel(context.Background()) + + mock.generateConsistencyTokenFunc = func(ctx context.Context, req *adminpb.GenerateConsistencyTokenRequest) (*adminpb.GenerateConsistencyTokenResponse, error) { + return &adminpb.GenerateConsistencyTokenResponse{ + ConsistencyToken: token, + }, nil + } + + // CheckConsistency returns Consistent: false, prompting a sleep/tick + mock.checkConsistencyFunc = func(ctx context.Context, req *adminpb.CheckConsistencyRequest) (*adminpb.CheckConsistencyResponse, error) { + cancel() // Cancel the context to trigger the select case + return &adminpb.CheckConsistencyResponse{ + Consistent: false, + }, nil + } + + client, cleanup := setupMockClient(t, mock) + defer cleanup() + + err := client.WaitForConsistency(ctx, "projects/p/instances/i/tables/t") + if err != context.Canceled && status.Code(err) != codes.Canceled { + t.Fatalf("WaitForConsistency error: %v, want %v or gRPC Canceled status", err, context.Canceled) + } +} diff --git a/bigtable/admin_test.go b/bigtable/admin_test.go index 625cd7a466ee..461d07f526df 100644 --- a/bigtable/admin_test.go +++ b/bigtable/admin_test.go @@ -28,6 +28,7 @@ import ( "cloud.google.com/go/internal/testutil" longrunning "cloud.google.com/go/longrunning/autogen/longrunningpb" "github.com/google/go-cmp/cmp" + "google.golang.org/api/option" "google.golang.org/grpc" "google.golang.org/protobuf/types/known/anypb" "google.golang.org/protobuf/types/known/durationpb" @@ -1805,3 +1806,59 @@ func TestInstanceAdmin_UpdateAppProfile(t *testing.T) { }) } } + +func TestTableAdminClientV2(t *testing.T) { + ctx := context.Background() + c, err := NewAdminClient(ctx, "my-cool-project", "my-cool-instance", option.WithoutAuthentication()) + if err != nil { + t.Fatalf("NewAdminClient failed: %v", err) + } + + gapicClient := c.TableAdminClientV2() + if gapicClient == nil { + t.Fatal("Expected non-nil BigtableTableAdminClient") + } + + // Close the gapic client. It should close the underlying connPool. + if err := gapicClient.Close(); err != nil { + t.Errorf("gapicClient.Close() failed: %v", err) + } + + // Now trying to use AdminClient should fail because the connection pool is closed. + err = c.CreateTable(ctx, "some-table") + if err == nil { + t.Error("Expected error when calling method on closed client") + } else { + if !strings.Contains(err.Error(), "closing") && !strings.Contains(err.Error(), "closed") && !strings.Contains(err.Error(), "Shutdown") { + t.Errorf("Unexpected error: %v", err) + } + } +} + +func TestInstanceAdminClientV2(t *testing.T) { + ctx := context.Background() + c, err := NewInstanceAdminClient(ctx, "my-cool-project", option.WithoutAuthentication()) + if err != nil { + t.Fatalf("NewInstanceAdminClient failed: %v", err) + } + + gapicClient := c.InstanceAdminClientV2() + if gapicClient == nil { + t.Fatal("Expected non-nil BigtableInstanceAdminClient") + } + + // Close the gapic client. It should close the underlying connPool. + if err := gapicClient.Close(); err != nil { + t.Errorf("gapicClient.Close() failed: %v", err) + } + + // Now trying to use InstanceAdminClient should fail because the connection pool is closed. + _, err = c.Instances(ctx) + if err == nil { + t.Error("Expected error when calling method on closed client") + } else { + if !strings.Contains(err.Error(), "closing") && !strings.Contains(err.Error(), "closed") && !strings.Contains(err.Error(), "Shutdown") { + t.Errorf("Unexpected error: %v", err) + } + } +} diff --git a/bigtable/integration_test.go b/bigtable/integration_test.go index 685ee1f0a719..8b92f000759e 100644 --- a/bigtable/integration_test.go +++ b/bigtable/integration_test.go @@ -4184,6 +4184,26 @@ func TestIntegration_AdminBackup(t *testing.T) { t.Fatalf("Restored TableInfo: %v", err) } + // Test V2 client RestoreTable + restoredTableV2 := tblConf.TableID + "-restored-v2" + t.Cleanup(func() { deleteTable(context.Background(), t, adminClient, restoredTableV2) }) + + v2Client := adminClient.TableAdminClientV2() + parentPath := fmt.Sprintf("projects/%s/instances/%s", testEnv.Config().Project, testEnv.Config().Instance) + backupPath := fmt.Sprintf("%s/clusters/%s/backups/%s", parentPath, sourceCluster, stdBkpName) + + req := &btapb.RestoreTableRequest{ + Parent: parentPath, + TableId: restoredTableV2, + Source: &btapb.RestoreTableRequest_Backup{Backup: backupPath}, + } + if err = v2Client.RestoreTable(ctx, req); err != nil { + t.Fatalf("V2 RestoreTable: %v", err) + } + if _, err := adminClient.TableInfo(ctx, restoredTableV2); err != nil { + t.Fatalf("V2 Restored TableInfo: %v", err) + } + // If 'it.run-create-instance-tests' flag is set while running the tests, // instanceToCreate will be non-empty string. // Add more testcases if instanceToCreate is non-empty string