[feat] RR Create API 2: Add storage and spanner accessor

accessors/clients/dataflow/dataflow_client.go

@@ -0,0 +1,43 @@
+// Copyright 2024 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
+//
+//	http://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 dataflowclient
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	dataflow "cloud.google.com/go/dataflow/apiv1beta3"
+)
+
+var once sync.Once
+var dfClient *dataflow.FlexTemplatesClient
+
+// This function is declared as a global variable to make it testable. The unit
+// tests update this function, acting like a double.
+var newFlexTemplatesClient = dataflow.NewFlexTemplatesClient
+
+func GetOrCreateClient(ctx context.Context) (*dataflow.FlexTemplatesClient, error) {
+	var err error
+	if dfClient == nil {
+		once.Do(func() {
+			dfClient, err = newFlexTemplatesClient(ctx)
+		})
+		if err != nil {
+			return nil, fmt.Errorf("failed to create dataflow client: %v", err)
+		}
+		return dfClient, nil
+	}
+	return dfClient, nil
+}

accessors/clients/dataflow/dataflow_client_test.go

@@ -0,0 +1,116 @@
+// Copyright 2024 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
+//
+//	http://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 dataflowclient
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"sync"
+	"testing"
+
+	dataflow "cloud.google.com/go/dataflow/apiv1beta3"
+	"github.com/GoogleCloudPlatform/spanner-migration-tool/logger"
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap"
+	"google.golang.org/api/option"
+)
+
+func init() {
+	logger.Log = zap.NewNop()
+}
+
+func TestMain(m *testing.M) {
+	res := m.Run()
+	os.Exit(res)
+}
+
+func resetTest() {
+	dfClient = nil
+	once = sync.Once{}
+}
+
+func TestGetOrCreateClient_Basic(t *testing.T) {
+	resetTest()
+	ctx := context.Background()
+	oldFunc := newFlexTemplatesClient
+	defer func() { newFlexTemplatesClient = oldFunc }()
+	newFlexTemplatesClient = func(ctx context.Context, opts ...option.ClientOption) (*dataflow.FlexTemplatesClient, error) {
+		return &dataflow.FlexTemplatesClient{}, nil
+	}
+	c, err := GetOrCreateClient(ctx)
+	assert.NotNil(t, c)
+	assert.Nil(t, err)
+}
+
+func TestGetOrCreateClient_OnlyOnceViaSync(t *testing.T) {
+	resetTest()
+	ctx := context.Background()
+	oldFunc := newFlexTemplatesClient
+	defer func() { newFlexTemplatesClient = oldFunc }()
+
+	newFlexTemplatesClient = func(ctx context.Context, opts ...option.ClientOption) (*dataflow.FlexTemplatesClient, error) {
+		return &dataflow.FlexTemplatesClient{}, nil
+	}
+	c, err := GetOrCreateClient(ctx)
+	assert.NotNil(t, c)
+	assert.Nil(t, err)
+	// Explicitly set the client to nil. Running GetOrCreateClient should not create a
+	// new client since sync would already be executed.
+	dfClient = nil
+	newFlexTemplatesClient = func(ctx context.Context, opts ...option.ClientOption) (*dataflow.FlexTemplatesClient, error) {
+		return nil, fmt.Errorf("test error")
+	}
+	c, err = GetOrCreateClient(ctx)
+	assert.Nil(t, c)
+	assert.Nil(t, err)
+}
+
+func TestGetOrCreateClient_OnlyOnceViaIf(t *testing.T) {
+	resetTest()
+	ctx := context.Background()
+	oldFunc := newFlexTemplatesClient
+	defer func() { newFlexTemplatesClient = oldFunc }()
+
+	newFlexTemplatesClient = func(ctx context.Context, opts ...option.ClientOption) (*dataflow.FlexTemplatesClient, error) {
+		return &dataflow.FlexTemplatesClient{}, nil
+	}
+	oldC, err := GetOrCreateClient(ctx)
+	assert.NotNil(t, oldC)
+	assert.Nil(t, err)
+
+	// Explicitly reset once. Running GetOrCreateClient should not create a
+	// new client the if condition should prevent it.
+	once = sync.Once{}
+	newFlexTemplatesClient = func(ctx context.Context, opts ...option.ClientOption) (*dataflow.FlexTemplatesClient, error) {
+		return nil, fmt.Errorf("test error")
+	}
+	newC, err := GetOrCreateClient(ctx)
+	assert.Equal(t, oldC, newC)
+	assert.Nil(t, err)
+}
+
+func TestGetOrCreateClient_Error(t *testing.T) {
+	resetTest()
+	ctx := context.Background()
+	oldFunc := newFlexTemplatesClient
+	defer func() { newFlexTemplatesClient = oldFunc }()
+
+	newFlexTemplatesClient = func(ctx context.Context, opts ...option.ClientOption) (*dataflow.FlexTemplatesClient, error) {
+		return nil, fmt.Errorf("test error")
+	}
+	c, err := GetOrCreateClient(ctx)
+	assert.Nil(t, c)
+	assert.NotNil(t, err)
+}

accessors/clients/dataflow/interface.go

@@ -0,0 +1,44 @@
+// Copyright 2024 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
+//
+//	http://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 dataflowclient
+
+import (
+	"context"
+
+	dataflow "cloud.google.com/go/dataflow/apiv1beta3"
+	"cloud.google.com/go/dataflow/apiv1beta3/dataflowpb"
+	"github.com/googleapis/gax-go/v2"
+)
+
+// Use this interface instead of dataflow.FlexTemplatesClient to support mocking.
+type DataflowClient interface {
+	LaunchFlexTemplate(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error)
+}
+
+// This implements the DataflowClient interface. This is the primary implementation that should be used in all places other than tests.
+type DataflowClientImpl struct {
+	client *dataflow.FlexTemplatesClient
+}
+
+func NewDataflowClientImpl(ctx context.Context) (*DataflowClientImpl, error) {
+	c, err := GetOrCreateClient(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return &DataflowClientImpl{client: c}, nil
+}
+
+func (c *DataflowClientImpl) LaunchFlexTemplate(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) {
+	return c.client.LaunchFlexTemplate(ctx, req, opts...)
+}

accessors/clients/dataflow/mocks.go

@@ -0,0 +1,31 @@
+// Copyright 2024 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
+//
+//	http://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 dataflowclient
+
+import (
+	"context"
+
+	"cloud.google.com/go/dataflow/apiv1beta3/dataflowpb"
+	"github.com/googleapis/gax-go/v2"
+)
+
+// Mock that implements the DataflowClient interface.
+// Pass in unit tests where DataflowClient is an input parameter.
+type DataflowClientMock struct {
+	LaunchFlexTemplateMock func(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error)
+}
+
+func (dcm *DataflowClientMock) LaunchFlexTemplate(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) {
+	return dcm.LaunchFlexTemplateMock(ctx, req, opts...)
+}

accessors/clients/spanner/admin/admin_client.go

@@ -0,0 +1,43 @@
+// Copyright 2024 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
+//
+//	http://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 spanneradmin
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	database "cloud.google.com/go/spanner/admin/database/apiv1"
+)
+
+var once sync.Once
+var spannerAdminClient *database.DatabaseAdminClient
+
+// This function is declared as a global variable to make it testable. The unit
+// tests update this function, acting like a double.
+var newDatabaseAdminClient = database.NewDatabaseAdminClient
+
+func GetOrCreateClient(ctx context.Context) (*database.DatabaseAdminClient, error) {
+	var err error
+	if spannerAdminClient == nil {
+		once.Do(func() {
+			spannerAdminClient, err = newDatabaseAdminClient(ctx)
+		})
+		if err != nil {
+			return nil, fmt.Errorf("failed to create spanner admin client: %v", err)
+		}
+		return spannerAdminClient, nil
+	}
+	return spannerAdminClient, nil
+}

accessors/clients/spanner/admin/admin_client_test.go

@@ -0,0 +1,116 @@
+// Copyright 2024 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
+//
+//	http://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 spanneradmin
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"sync"
+	"testing"
+
+	database "cloud.google.com/go/spanner/admin/database/apiv1"
+	"github.com/GoogleCloudPlatform/spanner-migration-tool/logger"
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap"
+	"google.golang.org/api/option"
+)
+
+func init() {
+	logger.Log = zap.NewNop()
+}
+
+func TestMain(m *testing.M) {
+	res := m.Run()
+	os.Exit(res)
+}
+
+func resetTest() {
+	spannerAdminClient = nil
+	once = sync.Once{}
+}
+
+func TestGetOrCreateClient_Basic(t *testing.T) {
+	resetTest()
+	ctx := context.Background()
+	oldFunc := newDatabaseAdminClient
+	defer func() { newDatabaseAdminClient = oldFunc }()
+	newDatabaseAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*database.DatabaseAdminClient, error) {
+		return &database.DatabaseAdminClient{}, nil
+	}
+	c, err := GetOrCreateClient(ctx)
+	assert.NotNil(t, c)
+	assert.Nil(t, err)
+}
+
+func TestGetOrCreateClient_OnlyOnceViaSync(t *testing.T) {
+	resetTest()
+	ctx := context.Background()
+	oldFunc := newDatabaseAdminClient
+	defer func() { newDatabaseAdminClient = oldFunc }()
+
+	newDatabaseAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*database.DatabaseAdminClient, error) {
+		return &database.DatabaseAdminClient{}, nil
+	}
+	c, err := GetOrCreateClient(ctx)
+	assert.NotNil(t, c)
+	assert.Nil(t, err)
+	// Explicitly set the client to nil. Running GetOrCreateClient should not create a
+	// new client since sync would already be executed.
+	spannerAdminClient = nil
+	newDatabaseAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*database.DatabaseAdminClient, error) {
+		return nil, fmt.Errorf("test error")
+	}
+	c, err = GetOrCreateClient(ctx)
+	assert.Nil(t, c)
+	assert.Nil(t, err)
+}
+
+func TestGetOrCreateClient_OnlyOnceViaIf(t *testing.T) {
+	resetTest()
+	ctx := context.Background()
+	oldFunc := newDatabaseAdminClient
+	defer func() { newDatabaseAdminClient = oldFunc }()
+
+	newDatabaseAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*database.DatabaseAdminClient, error) {
+		return &database.DatabaseAdminClient{}, nil
+	}
+	oldC, err := GetOrCreateClient(ctx)
+	assert.NotNil(t, oldC)
+	assert.Nil(t, err)
+
+	// Explicitly reset once. Running GetOrCreateClient should not create a
+	// new client the if condition should prevent it.
+	once = sync.Once{}
+	newDatabaseAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*database.DatabaseAdminClient, error) {
+		return nil, fmt.Errorf("test error")
+	}
+	newC, err := GetOrCreateClient(ctx)
+	assert.Equal(t, oldC, newC)
+	assert.Nil(t, err)
+}
+
+func TestGetOrCreateClient_Error(t *testing.T) {
+	resetTest()
+	ctx := context.Background()
+	oldFunc := newDatabaseAdminClient
+	defer func() { newDatabaseAdminClient = oldFunc }()
+
+	newDatabaseAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*database.DatabaseAdminClient, error) {
+		return nil, fmt.Errorf("test error")
+	}
+	c, err := GetOrCreateClient(ctx)
+	assert.Nil(t, c)
+	assert.NotNil(t, err)
+}