From 0124606b294dceb7240f04d91a8ce38d855f226f Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Tue, 30 Jan 2024 10:50:21 +0530 Subject: [PATCH 01/25] [feat] RR Create API 1: Add dataflow accessor (#745) * Add dataflow accessor * Add enable streaming engine struct tag Mofe Unmarshall Method to acc2 due ot storage dependency * Moved dataflow utils to accessor and creates types.go * Create dataflowutils package * Renamed testing package for dataflow util * Added unit tests * Added empty test files for clients * Move test to same package * Add tests for dataflow client * Update fake for client test * Make dataflow accessor interface and struct to make it testable * Remove interface from accessor package * Add dataflow accessor interface * Add comments to dataflow client and comments on unit tests * Move all dataflow dependencies to accessors and remove dataflow utils * Create dataflow client interface for accessor method to make it unit testable --- accessors/clients/dataflow/dataflow_client.go | 49 +++ .../clients/dataflow/dataflow_client_test.go | 116 +++++++ accessors/dataflow/dataflow_accessor.go | 178 ++++++++++ accessors/dataflow/dataflow_accessor_test.go | 310 ++++++++++++++++++ accessors/dataflow/dataflow_types.go | 32 ++ common/metrics/dashboard_components.go | 69 ++-- common/metrics/queries.go | 9 +- common/utils/dataflow_utils.go | 111 ------- streaming/streaming.go | 4 +- testing/common/utils/dataflow_utils_test.go | 117 ------- 10 files changed, 724 insertions(+), 271 deletions(-) create mode 100644 accessors/clients/dataflow/dataflow_client.go create mode 100644 accessors/clients/dataflow/dataflow_client_test.go create mode 100644 accessors/dataflow/dataflow_accessor.go create mode 100644 accessors/dataflow/dataflow_accessor_test.go create mode 100644 accessors/dataflow/dataflow_types.go delete mode 100644 common/utils/dataflow_utils.go delete mode 100644 testing/common/utils/dataflow_utils_test.go diff --git a/accessors/clients/dataflow/dataflow_client.go b/accessors/clients/dataflow/dataflow_client.go new file mode 100644 index 0000000000..cc76c34fda --- /dev/null +++ b/accessors/clients/dataflow/dataflow_client.go @@ -0,0 +1,49 @@ +// 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" + "cloud.google.com/go/dataflow/apiv1beta3/dataflowpb" + "github.com/googleapis/gax-go/v2" +) + +type DataflowClient interface { + LaunchFlexTemplate(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) +} + +var once sync.Once +var dfClient *dataflow.FlexTemplatesClient + +// This function is declared as a global variable to make it testable. The unit +// tests edit 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 +} diff --git a/accessors/clients/dataflow/dataflow_client_test.go b/accessors/clients/dataflow/dataflow_client_test.go new file mode 100644 index 0000000000..a76cc98fc4 --- /dev/null +++ b/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) +} diff --git a/accessors/dataflow/dataflow_accessor.go b/accessors/dataflow/dataflow_accessor.go new file mode 100644 index 0000000000..aa7f1bea49 --- /dev/null +++ b/accessors/dataflow/dataflow_accessor.go @@ -0,0 +1,178 @@ +// 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 dataflowaccessor + +import ( + "context" + "fmt" + "sort" + "strings" + + "cloud.google.com/go/dataflow/apiv1beta3/dataflowpb" + dataflowclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/dataflow" + "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" + "golang.org/x/exp/maps" +) + +type DataflowAccessor interface { + // This function takes the template parameters (@parameters) and runtime environment config (@cfg) as input, and returns + // the generated jobId, equivalentGcloudCommand and error if any. + LaunchFlexTemplate(ctx context.Context, c dataflowclient.DataflowClient, parameters map[string]string, cfg DataflowTuningConfig) (string, string, error) +} + +type DataflowAccessorImpl struct{} + +func (dfA *DataflowAccessorImpl) LaunchDataflowTemplate(ctx context.Context, c dataflowclient.DataflowClient, parameters map[string]string, cfg DataflowTuningConfig) (string, string, error) { + req, err := getDataflowLaunchRequest(parameters, cfg) + if err != nil { + return "", "", err + } + respDf, err := c.LaunchFlexTemplate(ctx, req) + if err != nil { + logger.Log.Error(fmt.Sprintf("flexTemplateRequest: %+v\n", req)) + return "", "", fmt.Errorf("error launching dataflow template: %v", err) + } + gCloudCmd := GetGcloudDataflowCommandFromRequest(req) + return respDf.Job.Id, gCloudCmd, nil +} + +func getDataflowLaunchRequest(parameters map[string]string, cfg DataflowTuningConfig) (*dataflowpb.LaunchFlexTemplateRequest, error) { + // If custom network is not selected, use public IP. Typical for internal testing flow. + vpcSubnetwork := "" + workerIpAddressConfig := dataflowpb.WorkerIPAddressConfiguration_WORKER_IP_PUBLIC + if cfg.Network != "" || cfg.Subnetwork != "" { + workerIpAddressConfig = dataflowpb.WorkerIPAddressConfiguration_WORKER_IP_PRIVATE + // If subnetwork is not provided, assume network has auto subnet configuration. + if cfg.Subnetwork != "" { + if cfg.VpcHostProjectId == "" || cfg.Location == "" { + return nil, fmt.Errorf("vpc host project id and location must be specified when specifying subnetwork") + } + vpcSubnetwork = fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/regions/%s/subnetworks/%s", cfg.VpcHostProjectId, cfg.Location, cfg.Subnetwork) + } + } + // Dataflow does not accept upper case letters in the name. + cfg.JobName = strings.ToLower(cfg.JobName) + request := &dataflowpb.LaunchFlexTemplateRequest{ + ProjectId: cfg.ProjectId, + LaunchParameter: &dataflowpb.LaunchFlexTemplateParameter{ + JobName: cfg.JobName, + Template: &dataflowpb.LaunchFlexTemplateParameter_ContainerSpecGcsPath{ContainerSpecGcsPath: cfg.GcsTemplatePath}, + Parameters: parameters, + Environment: &dataflowpb.FlexTemplateRuntimeEnvironment{ + MaxWorkers: cfg.MaxWorkers, + NumWorkers: cfg.NumWorkers, + ServiceAccountEmail: cfg.ServiceAccountEmail, + MachineType: cfg.MachineType, + AdditionalUserLabels: cfg.AdditionalUserLabels, + KmsKeyName: cfg.KmsKeyName, + Network: cfg.Network, + Subnetwork: vpcSubnetwork, + IpConfiguration: workerIpAddressConfig, + AdditionalExperiments: cfg.AdditionalExperiments, + EnableStreamingEngine: cfg.EnableStreamingEngine, + }, + }, + Location: cfg.Location, + } + logger.Log.Debug(fmt.Sprintf("Flex Template request generated: %+v", request)) + return request, nil +} + +// Generate the equivalent gCloud CLI command to launch a dataflow job with the same parameters and environment flags +// as the input body. +func GetGcloudDataflowCommandFromRequest(req *dataflowpb.LaunchFlexTemplateRequest) string { + lp := req.LaunchParameter + templatePath := lp.Template.(*dataflowpb.LaunchFlexTemplateParameter_ContainerSpecGcsPath).ContainerSpecGcsPath + cmd := fmt.Sprintf("gcloud dataflow flex-template run %s --project=%s --region=%s --template-file-gcs-location=%s %s %s", + lp.JobName, req.ProjectId, req.Location, templatePath, getEnvironmentFlags(lp.Environment), getParametersFlag(lp.Parameters)) + return strings.Trim(cmd, " ") +} + +// Generate the equivalent parameter flag string, returning empty string if none are specified. +func getParametersFlag(parameters map[string]string) string { + if len(parameters) == 0 { + return "" + } + params := "" + keys := maps.Keys(parameters) + sort.Strings(keys) + for _, k := range keys { + params = params + k + "=" + parameters[k] + "," + } + params = strings.TrimSuffix(params, ",") + return fmt.Sprintf("--parameters %s", params) +} + +// We don't populate all flags in the API because certain flags (like AutoscalingAlgorithm, DumpHeapOnOom etc.) +// are not supported in gCloud. +func getEnvironmentFlags(environment *dataflowpb.FlexTemplateRuntimeEnvironment) string { + flag := "" + if environment.NumWorkers != 0 { + flag += fmt.Sprintf("--num-workers %d ", environment.NumWorkers) + } + if environment.MaxWorkers != 0 { + flag += fmt.Sprintf("--max-workers %d ", environment.MaxWorkers) + } + if environment.ServiceAccountEmail != "" { + flag += fmt.Sprintf("--service-account-email %s ", environment.ServiceAccountEmail) + } + if environment.TempLocation != "" { + flag += fmt.Sprintf("--temp-location %s ", environment.TempLocation) + } + if environment.MachineType != "" { + flag += fmt.Sprintf("--worker-machine-type %s ", environment.MachineType) + } + if environment.AdditionalExperiments != nil && len(environment.AdditionalExperiments) > 0 { + flag += fmt.Sprintf("--additional-experiments %s ", strings.Join(environment.AdditionalExperiments, ",")) + } + if environment.Network != "" { + flag += fmt.Sprintf("--network %s ", environment.Network) + } + if environment.Subnetwork != "" { + flag += fmt.Sprintf("--subnetwork %s ", environment.Subnetwork) + } + if environment.AdditionalUserLabels != nil && len(environment.AdditionalUserLabels) > 0 { + flag += fmt.Sprintf("--additional-user-labels %s ", formatAdditionalUserLabels(environment.AdditionalUserLabels)) + } + if environment.KmsKeyName != "" { + flag += fmt.Sprintf("--dataflow-kms-key %s ", environment.KmsKeyName) + } + if environment.IpConfiguration == dataflowpb.WorkerIPAddressConfiguration_WORKER_IP_PRIVATE { + flag += "--disable-public-ips " + } + if environment.WorkerRegion != "" { + flag += fmt.Sprintf("--worker-region %s ", environment.WorkerRegion) + } + if environment.WorkerZone != "" { + flag += fmt.Sprintf("--worker-zone %s ", environment.WorkerZone) + } + if environment.EnableStreamingEngine { + flag += "--enable-streaming-engine " + } + if environment.FlexrsGoal != dataflowpb.FlexResourceSchedulingGoal_FLEXRS_UNSPECIFIED { + flag += fmt.Sprintf("--flexrs-goal %s ", environment.FlexrsGoal) + } + if environment.StagingLocation != "" { + flag += fmt.Sprintf("--staging-location %s ", environment.StagingLocation) + } + return strings.Trim(flag, " ") +} + +func formatAdditionalUserLabels(labels map[string]string) string { + res := []string{} + for key, value := range labels { + res = append(res, fmt.Sprintf("%s=%s", key, value)) + } + return strings.Join(res, ",") +} diff --git a/accessors/dataflow/dataflow_accessor_test.go b/accessors/dataflow/dataflow_accessor_test.go new file mode 100644 index 0000000000..ffb8b6537b --- /dev/null +++ b/accessors/dataflow/dataflow_accessor_test.go @@ -0,0 +1,310 @@ +// 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 dataflowaccessor + +import ( + "context" + "fmt" + "os" + "testing" + + "cloud.google.com/go/dataflow/apiv1beta3/dataflowpb" + "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" + "github.com/google/go-cmp/cmp" + "github.com/googleapis/gax-go/v2" + "github.com/stretchr/testify/assert" + "go.uber.org/zap" +) + +func init() { + logger.Log = zap.NewNop() +} + +func TestMain(m *testing.M) { + res := m.Run() + os.Exit(res) +} + +func getParameters() map[string]string { + return map[string]string{ + "inputFilePattern": "gs://inputFilePattern", + "streamName": "my-stream", + "instanceId": "my-instance", + "databaseId": "my-dbName", + "sessionFilePath": "gs://session.json", + "deadLetterQueueDirectory": "gs://dlq", + "transformationContextFilePath": "gs://transformationContext.json", + "directoryWatchDurationInMinutes": "480", // Setting directory watch timeout to 8 hours + } +} + +func getTuningConfig() DataflowTuningConfig { + return DataflowTuningConfig{ + ProjectId: "test-project", + JobName: "test-job", + Location: "us-central1", + VpcHostProjectId: "host-project", + Network: "my-network", + Subnetwork: "my-subnetwork", + MaxWorkers: 50, + NumWorkers: 10, + ServiceAccountEmail: "svc-account@google.com", + MachineType: "n2-standard-64", + AdditionalUserLabels: map[string]string{"name": "wrench"}, + KmsKeyName: "sample-kms-key", + GcsTemplatePath: "gs://template/Cloud_Datastream_to_Spanner", + AdditionalExperiments: []string{"use_runner_V2", "test-experiment"}, + EnableStreamingEngine: true, + } +} + +func getTemplateDfRequest1() *dataflowpb.LaunchFlexTemplateRequest { + return &dataflowpb.LaunchFlexTemplateRequest{ + ProjectId: "test-project", + Location: "us-central1", + LaunchParameter: &dataflowpb.LaunchFlexTemplateParameter{ + JobName: "test-job", + Template: &dataflowpb.LaunchFlexTemplateParameter_ContainerSpecGcsPath{ContainerSpecGcsPath: "gs://template/Cloud_Datastream_to_Spanner"}, + Parameters: getParameters(), + Environment: &dataflowpb.FlexTemplateRuntimeEnvironment{ + MaxWorkers: 50, + NumWorkers: 10, + ServiceAccountEmail: "svc-account@google.com", + MachineType: "n2-standard-64", + AdditionalUserLabels: map[string]string{"name": "wrench"}, + KmsKeyName: "sample-kms-key", + Network: "my-network", + Subnetwork: "https://www.googleapis.com/compute/v1/projects/host-project/regions/us-central1/subnetworks/my-subnetwork", + IpConfiguration: dataflowpb.WorkerIPAddressConfiguration_WORKER_IP_PRIVATE, + AdditionalExperiments: []string{"use_runner_V2", "test-experiment"}, + EnableStreamingEngine: true, + }, + }, + } +} + +func getExpectedGcloudCmd1() string { + return "gcloud dataflow flex-template run test-job " + + "--project=test-project --region=us-central1 " + + "--template-file-gcs-location=gs://template/Cloud_Datastream_to_Spanner " + + "--num-workers 10 --max-workers 50 --service-account-email svc-account@google.com " + + "--worker-machine-type n2-standard-64 " + + "--additional-experiments use_runner_V2,test-experiment --network my-network " + + "--subnetwork https://www.googleapis.com/compute/v1/projects/host-project/regions/us-central1/subnetworks/my-subnetwork --additional-user-labels name=wrench " + + "--dataflow-kms-key sample-kms-key --disable-public-ips " + + "--enable-streaming-engine " + + "--parameters databaseId=my-dbName,deadLetterQueueDirectory=gs://dlq," + + "directoryWatchDurationInMinutes=480,inputFilePattern=gs://inputFilePattern," + + "instanceId=my-instance,sessionFilePath=gs://session.json,streamName=my-stream," + + "transformationContextFilePath=gs://transformationContext.json" +} + +func getTemplateDfRequest2() *dataflowpb.LaunchFlexTemplateRequest { + return &dataflowpb.LaunchFlexTemplateRequest{ + ProjectId: "test-project", + Location: "us-central1", + LaunchParameter: &dataflowpb.LaunchFlexTemplateParameter{ + JobName: "test-job", + Template: &dataflowpb.LaunchFlexTemplateParameter_ContainerSpecGcsPath{ContainerSpecGcsPath: "gs://template/Cloud_Datastream_to_Spanner"}, + Parameters: getParameters(), + Environment: &dataflowpb.FlexTemplateRuntimeEnvironment{ + MaxWorkers: 50, + NumWorkers: 10, + ServiceAccountEmail: "svc-account@google.com", + MachineType: "n2-standard-64", + AdditionalUserLabels: map[string]string{"name": "wrench"}, + KmsKeyName: "sample-kms-key", + Network: "my-network", + Subnetwork: "https://www.googleapis.com/compute/v1/projects/host-project/regions/us-central1/subnetworks/my-subnetwork", + IpConfiguration: dataflowpb.WorkerIPAddressConfiguration_WORKER_IP_PRIVATE, + AdditionalExperiments: []string{"use_runner_V2", "test-experiment"}, + EnableStreamingEngine: true, + TempLocation: "gs://temp-location", + WorkerRegion: "test-worker-region", + WorkerZone: "test-worker-zone", + FlexrsGoal: 1, + StagingLocation: "gs://staging-location", + }, + }, + } +} + +func getExpectedGcloudCmd2() string { + return "gcloud dataflow flex-template run test-job " + + "--project=test-project --region=us-central1 " + + "--template-file-gcs-location=gs://template/Cloud_Datastream_to_Spanner " + + "--num-workers 10 --max-workers 50 --service-account-email svc-account@google.com " + + "--temp-location gs://temp-location --worker-machine-type n2-standard-64 " + + "--additional-experiments use_runner_V2,test-experiment --network my-network " + + "--subnetwork https://www.googleapis.com/compute/v1/projects/host-project/regions/us-central1/subnetworks/my-subnetwork --additional-user-labels name=wrench " + + "--dataflow-kms-key sample-kms-key --disable-public-ips --worker-region test-worker-region " + + "--worker-zone test-worker-zone --enable-streaming-engine " + + "--flexrs-goal FLEXRS_SPEED_OPTIMIZED --staging-location gs://staging-location " + + "--parameters databaseId=my-dbName,deadLetterQueueDirectory=gs://dlq," + + "directoryWatchDurationInMinutes=480,inputFilePattern=gs://inputFilePattern," + + "instanceId=my-instance,sessionFilePath=gs://session.json,streamName=my-stream," + + "transformationContextFilePath=gs://transformationContext.json" +} + +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...) +} + +func TestLaunchDataflowTemplate(t *testing.T) { + ctx := context.Background() + da := DataflowAccessorImpl{} + testCases := []struct { + name string + params map[string]string + cfg DataflowTuningConfig + dcm DataflowClientMock + expectError bool + expectedJobId string + expectedGcloudCmd string + }{ + { + name: "Basic Correct", + params: getParameters(), + cfg: getTuningConfig(), + dcm: DataflowClientMock{ + LaunchFlexTemplateMock: func(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) { + return &dataflowpb.LaunchFlexTemplateResponse{Job: &dataflowpb.Job{Id: "1234"}}, nil + }, + }, + expectError: false, + expectedJobId: "1234", + expectedGcloudCmd: getExpectedGcloudCmd1(), + }, + { + name: "Request builder error", + params: getParameters(), + cfg: DataflowTuningConfig{Subnetwork: "test"}, + dcm: DataflowClientMock{ + LaunchFlexTemplateMock: func(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) { + return &dataflowpb.LaunchFlexTemplateResponse{Job: &dataflowpb.Job{Id: "1234"}}, nil + }, + }, + expectError: true, + expectedJobId: "", + expectedGcloudCmd: "", + }, + { + name: "Launch flex template throws error", + params: getParameters(), + cfg: getTuningConfig(), + dcm: DataflowClientMock{ + LaunchFlexTemplateMock: func(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) { + return nil, fmt.Errorf("test error") + }, + }, + expectError: true, + expectedJobId: "", + expectedGcloudCmd: "", + }, + } + for _, tc := range testCases { + jobId, gcloudCmd, err := da.LaunchDataflowTemplate(ctx, &tc.dcm, tc.params, tc.cfg) + assert.Equal(t, tc.expectError, err != nil) + assert.Equal(t, tc.expectedJobId, jobId) + assert.Equal(t, tc.expectedGcloudCmd, gcloudCmd) + } +} + +func TestGetDataflowLaunchRequestBasic(t *testing.T) { + params := getParameters() + cfg := getTuningConfig() + actual, err := getDataflowLaunchRequest(params, cfg) + if err != nil { + t.Fail() + } + expected := getTemplateDfRequest1() + assert.True(t, EquateLaunchFlexTemplateRequest(expected, actual)) +} + +func TestGetDataflowLaunchRequestMissingVpcHost(t *testing.T) { + params := getParameters() + cfg := getTuningConfig() + cfg.VpcHostProjectId = "" + _, err := getDataflowLaunchRequest(params, cfg) + assert.True(t, err != nil) +} + +func TestGetDataflowLaunchRequestNameToLowerCase(t *testing.T) { + params := getParameters() + cfg := getTuningConfig() + cfg.JobName = "CAPITalJobName" + actual, err := getDataflowLaunchRequest(params, cfg) + if err != nil { + t.Fail() + } + expected := getTemplateDfRequest1() + expected.LaunchParameter.JobName = "capitaljobname" + assert.True(t, EquateLaunchFlexTemplateRequest(expected, actual)) +} + +func TestGcloudCmdWithAllParams(t *testing.T) { + + req := getTemplateDfRequest2() + expectedCmd := getExpectedGcloudCmd2() + assert.Equal(t, expectedCmd, GetGcloudDataflowCommandFromRequest(req)) +} + +func TestGcloudCmdWithPartialParams(t *testing.T) { + + req := getTemplateDfRequest2() + req.LaunchParameter.Parameters = make(map[string]string) + req.LaunchParameter.Environment.FlexrsGoal = 0 + req.LaunchParameter.Environment.IpConfiguration = 0 + req.LaunchParameter.Environment.EnableStreamingEngine = false + req.LaunchParameter.Environment.AdditionalExperiments = []string{} + req.LaunchParameter.Environment.AdditionalUserLabels = make(map[string]string) + req.LaunchParameter.Environment.WorkerRegion = "" + req.LaunchParameter.Environment.NumWorkers = 0 + req.LaunchParameter.Environment.Network = "" + req.LaunchParameter.Environment.Subnetwork = "" + + expectedCmd := "gcloud dataflow flex-template run test-job " + + "--project=test-project --region=us-central1 " + + "--template-file-gcs-location=gs://template/Cloud_Datastream_to_Spanner " + + "--max-workers 50 --service-account-email svc-account@google.com " + + "--temp-location gs://temp-location --worker-machine-type n2-standard-64 " + + "--dataflow-kms-key sample-kms-key " + + "--worker-zone test-worker-zone " + + "--staging-location gs://staging-location" + assert.Equal(t, expectedCmd, GetGcloudDataflowCommandFromRequest(req)) +} + +func EquateLaunchFlexTemplateRequest(df1 *dataflowpb.LaunchFlexTemplateRequest, df2 *dataflowpb.LaunchFlexTemplateRequest) bool { + lp1 := df1.LaunchParameter + lp2 := df2.LaunchParameter + return (df1.ProjectId == df2.ProjectId && + df1.Location == df2.Location && + lp1.JobName == lp2.JobName && + lp1.Environment.MaxWorkers == lp2.Environment.MaxWorkers && + lp1.Environment.NumWorkers == lp2.Environment.NumWorkers && + lp1.Environment.ServiceAccountEmail == lp2.Environment.ServiceAccountEmail && + lp1.Environment.MachineType == lp2.Environment.MachineType && + lp1.Environment.KmsKeyName == lp2.Environment.KmsKeyName && + lp1.Environment.Network == lp2.Environment.Network && + lp1.Environment.Subnetwork == lp2.Environment.Subnetwork && + lp1.Environment.GetIpConfiguration().String() == lp2.Environment.GetIpConfiguration().String() && + lp1.Environment.EnableStreamingEngine == lp2.Environment.EnableStreamingEngine && + cmp.Equal(lp1.Environment.AdditionalUserLabels, lp2.Environment.AdditionalUserLabels) && + cmp.Equal(lp1.Environment.AdditionalExperiments, lp2.Environment.AdditionalExperiments) && + lp1.GetContainerSpecGcsPath() == lp2.GetContainerSpecGcsPath()) +} diff --git a/accessors/dataflow/dataflow_types.go b/accessors/dataflow/dataflow_types.go new file mode 100644 index 0000000000..d7b8f355ad --- /dev/null +++ b/accessors/dataflow/dataflow_types.go @@ -0,0 +1,32 @@ +// 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 dataflowaccessor + +type DataflowTuningConfig struct { + ProjectId string `json:"projectId"` + JobName string `json:"jobName"` + Location string `json:"location"` + VpcHostProjectId string `json:"hostProjectId"` + Network string `json:"network"` + Subnetwork string `json:"subnetwork"` + MaxWorkers int32 `json:"maxWorkers"` + NumWorkers int32 `json:"numWorkers"` + ServiceAccountEmail string `json:"serviceAccountEmail"` + MachineType string `json:"machineType"` + AdditionalUserLabels map[string]string `json:"additionalUserLabels"` + KmsKeyName string `json:"kmsKeyName"` + GcsTemplatePath string `json:"gcsTemplatePath"` + AdditionalExperiments []string `json:"additionalExperiments"` + EnableStreamingEngine bool `json:"enableStreamingEngine"` +} diff --git a/common/metrics/dashboard_components.go b/common/metrics/dashboard_components.go index dcd55b3358..b0e01a530b 100644 --- a/common/metrics/dashboard_components.go +++ b/common/metrics/dashboard_components.go @@ -4,16 +4,13 @@ // 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 +// 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 utils contains common helper functions used across multiple other packages. -// Utils should not import any Spanner migration tool packages. package metrics import ( @@ -42,22 +39,22 @@ var dashboardClient *dashboard.DashboardsClient // MonitoringMetricsResources contains information required to create the monitoring dashboard type MonitoringMetricsResources struct { - ProjectId string - DataflowJobId string - DatastreamId string - JobMetadataGcsBucket string - PubsubSubscriptionId string - SpannerInstanceId string - SpannerDatabaseId string - ShardToShardResourcesMap map[string]internal.ShardResources - ShardId string + ProjectId string + DataflowJobId string + DatastreamId string + JobMetadataGcsBucket string + PubsubSubscriptionId string + SpannerInstanceId string + SpannerDatabaseId string + ShardToShardResourcesMap map[string]internal.ShardResources + ShardId string MigrationRequestId string } type TileInfo struct { Title string TimeSeriesQueries map[string]string // Map of legend template and their corresponding queries - TextContent string // string for text input + TextContent string // string for text input } type MosaicGroup struct { @@ -95,10 +92,10 @@ func createShardDataflowMetrics(resourceIds MonitoringMetricsResources) []*dashb TileInfo{ Title: "Dataflow Workers Memory Utilization", TimeSeriesQueries: map[string]string{ - "p50 worker": fmt.Sprintf(dataflowMemoryUtilPercentileQuery, resourceIds.DataflowJobId, "50"), - "p90 worker": fmt.Sprintf(dataflowMemoryUtilPercentileQuery, resourceIds.DataflowJobId, "90"), - "Max worker": fmt.Sprintf(dataflowMemoryUtilMaxQuery, resourceIds.DataflowJobId), - }}.createXYChartTile(), + "p50 worker": fmt.Sprintf(dataflowMemoryUtilPercentileQuery, resourceIds.DataflowJobId, "50"), + "p90 worker": fmt.Sprintf(dataflowMemoryUtilPercentileQuery, resourceIds.DataflowJobId, "90"), + "Max worker": fmt.Sprintf(dataflowMemoryUtilMaxQuery, resourceIds.DataflowJobId), + }}.createXYChartTile(), TileInfo{Title: "Dataflow Workers Max Backlog Time Seconds", TimeSeriesQueries: map[string]string{"": fmt.Sprintf(dataflowBacklogTimeQuery, resourceIds.DataflowJobId)}}.createXYChartTile(), } return dataflowTiles @@ -107,7 +104,7 @@ func createShardDataflowMetrics(resourceIds MonitoringMetricsResources) []*dashb func createShardDatastreamMetrics(resourceIds MonitoringMetricsResources) []*dashboardpb.MosaicLayout_Tile { datastreamTiles := []*dashboardpb.MosaicLayout_Tile{ TileInfo{ - Title: "Datastream Total Latency", + Title: "Datastream Total Latency", TimeSeriesQueries: map[string]string{"p50 " + resourceIds.DatastreamId: fmt.Sprintf(datastreamTotalLatencyQuery, resourceIds.DatastreamId, "50"), "p90 " + resourceIds.DatastreamId: fmt.Sprintf(datastreamTotalLatencyQuery, resourceIds.DatastreamId, "90")}}.createXYChartTile(), TileInfo{Title: "Datastream Throughput", TimeSeriesQueries: map[string]string{resourceIds.DatastreamId: fmt.Sprintf(datastreamThroughputQuery, resourceIds.DatastreamId)}}.createXYChartTile(), TileInfo{Title: "Datastream Unsupported Events", TimeSeriesQueries: map[string]string{resourceIds.DatastreamId: fmt.Sprintf(datastreamUnsupportedEventsQuery, resourceIds.DatastreamId)}}.createXYChartTile(), @@ -147,8 +144,8 @@ func createShardIndependentTopMetrics(resourceIds MonitoringMetricsResources) [] TileInfo{Title: "Datastream Unsupported Events", TimeSeriesQueries: map[string]string{resourceIds.DatastreamId: fmt.Sprintf(datastreamUnsupportedEventsQuery, resourceIds.DatastreamId)}}.createXYChartTile(), TileInfo{Title: "Pubsub Age of Oldest Unacknowledged Message", TimeSeriesQueries: map[string]string{resourceIds.PubsubSubscriptionId: fmt.Sprintf(pubsubOldestUnackedMessageAgeQuery, resourceIds.PubsubSubscriptionId)}}.createXYChartTile(), } - spannerMetrics:=createSpannerMetrics(resourceIds) - independentTopMetricsTiles=append(independentTopMetricsTiles,spannerMetrics...) + spannerMetrics := createSpannerMetrics(resourceIds) + independentTopMetricsTiles = append(independentTopMetricsTiles, spannerMetrics...) return independentTopMetricsTiles } @@ -178,12 +175,12 @@ func createAggDataflowMetrics(resourceIds MonitoringMetricsResources) []*dashboa "Max shard": fmt.Sprintf(dataflowAggCpuUtilMaxQuery, createAggFilterCondition("metadata.user_labels.dataflow_job_id", dataflowJobs)), }}.createXYChartTile(), TileInfo{ - Title: "Dataflow Workers Memory Utilization", + Title: "Dataflow Workers Memory Utilization", TimeSeriesQueries: map[string]string{ - "p50 shard": fmt.Sprintf(dataflowAggMemoryUtilPercentileQuery, createAggFilterCondition("metadata.user_labels.dataflow_job_id", dataflowJobs), "50"), - "p90 shard": fmt.Sprintf(dataflowAggMemoryUtilPercentileQuery, createAggFilterCondition("metadata.user_labels.dataflow_job_id", dataflowJobs), "90"), - "Max shard": fmt.Sprintf(dataflowAggMemoryUtilMaxQuery, createAggFilterCondition("metadata.user_labels.dataflow_job_id", dataflowJobs)), - }}.createXYChartTile(), + "p50 shard": fmt.Sprintf(dataflowAggMemoryUtilPercentileQuery, createAggFilterCondition("metadata.user_labels.dataflow_job_id", dataflowJobs), "50"), + "p90 shard": fmt.Sprintf(dataflowAggMemoryUtilPercentileQuery, createAggFilterCondition("metadata.user_labels.dataflow_job_id", dataflowJobs), "90"), + "Max shard": fmt.Sprintf(dataflowAggMemoryUtilMaxQuery, createAggFilterCondition("metadata.user_labels.dataflow_job_id", dataflowJobs)), + }}.createXYChartTile(), TileInfo{Title: "Dataflow Workers Max Backlog Time Seconds", TimeSeriesQueries: map[string]string{"Dataflow Backlog Time Seconds": fmt.Sprintf(dataflowAggBacklogTimeQuery, createAggFilterCondition("metric.job_id", dataflowJobs))}}.createXYChartTile(), TileInfo{Title: "Dataflow Per Shard Median CPU Utilization", TimeSeriesQueries: map[string]string{"": fmt.Sprintf(dataflowAggPerShardCpuUtil, createAggFilterCondition("metadata.user_labels.dataflow_job_id", dataflowJobs))}}.createXYChartTile(), } @@ -197,7 +194,7 @@ func createAggDatastreamMetrics(resourceIds MonitoringMetricsResources) []*dashb } datastreamTiles := []*dashboardpb.MosaicLayout_Tile{ TileInfo{ - Title: "Datastream Total Latency", + Title: "Datastream Total Latency", TimeSeriesQueries: map[string]string{"p50 Datastream Latency": fmt.Sprintf(datastreamAggTotalLatencyQuery, createAggFilterCondition("resource.stream_id", datastreamJobs), "50", "50"), "p90 Datastream Latency": fmt.Sprintf(datastreamAggTotalLatencyQuery, createAggFilterCondition("resource.stream_id", datastreamJobs), "90", "90")}}.createXYChartTile(), TileInfo{Title: "Total Datastream Throughput", TimeSeriesQueries: map[string]string{"Datastream Total Throughput": fmt.Sprintf(datastreamAggThroughputQuery, createAggFilterCondition("resource.stream_id", datastreamJobs))}}.createXYChartTile(), TileInfo{Title: "Total Datastream Unsupported Events", TimeSeriesQueries: map[string]string{"Datastream Total Unsupported Events": fmt.Sprintf(datastreamAggUnsupportedEventsQuery, createAggFilterCondition("resource.stream_id", datastreamJobs))}}.createXYChartTile(), @@ -252,9 +249,9 @@ func createAggIndependentTopMetrics(resourceIds MonitoringMetricsResources) []*d TileInfo{Title: "Total Datastream Throughput", TimeSeriesQueries: map[string]string{"Datastream Throughput": fmt.Sprintf(datastreamAggThroughputQuery, createAggFilterCondition("resource.stream_id", datastreamJobs))}}.createXYChartTile(), TileInfo{Title: "Total Datastream Unsupported Events", TimeSeriesQueries: map[string]string{"Datastream Unsupported Events": fmt.Sprintf(datastreamAggUnsupportedEventsQuery, createAggFilterCondition("resource.stream_id", datastreamJobs))}}.createXYChartTile(), TileInfo{Title: "Pubsub Age of Oldest Unacknowledged Message", TimeSeriesQueries: map[string]string{"Pubsub Age of Oldest Unacknowledged Message": fmt.Sprintf(pubsubAggOldestUnackedMessageAgeQuery, createAggFilterCondition("resource.subscription_id", pubsubSubs))}}.createXYChartTile(), - } - spannerMetrics:=createSpannerMetrics(resourceIds) - independentTopMetricsTiles=append(independentTopMetricsTiles,spannerMetrics...) + } + spannerMetrics := createSpannerMetrics(resourceIds) + independentTopMetricsTiles = append(independentTopMetricsTiles, spannerMetrics...) return independentTopMetricsTiles } @@ -263,7 +260,7 @@ func createAggIndependentBottomMetrics(resourceIds MonitoringMetricsResources) [ for shardId, shardResource := range resourceIds.ShardToShardResourcesMap { shardUrl := fmt.Sprintf("https://console.cloud.google.com/monitoring/dashboards/builder/%v?project=%v", shardResource.MonitoringResources.DashboardName, resourceIds.ProjectId) shardString := fmt.Sprintf("Shard [%s](%s)", shardId, shardUrl) - if(shardToDashboardMappingText == ""){ + if shardToDashboardMappingText == "" { shardToDashboardMappingText = shardString } else { shardToDashboardMappingText += " \\\n" + shardString @@ -271,7 +268,7 @@ func createAggIndependentBottomMetrics(resourceIds MonitoringMetricsResources) [ } independentBottomMetricsTiles := []*dashboardpb.MosaicLayout_Tile{ TileInfo{ - Title: "Shard Dashboards", + Title: "Shard Dashboards", TextContent: shardToDashboardMappingText, }.createTextTile(), } @@ -332,14 +329,14 @@ func (tileInfo TileInfo) createCollapsibleGroupTile(tiles []*dashboardpb.MosaicL return &groupTile, heightOffset + groupTileHeight } -func (tileInfo TileInfo) createTextTile() (*dashboardpb.MosaicLayout_Tile){ - textTile := dashboardpb.MosaicLayout_Tile{ +func (tileInfo TileInfo) createTextTile() *dashboardpb.MosaicLayout_Tile { + textTile := dashboardpb.MosaicLayout_Tile{ Widget: &dashboardpb.Widget{ Title: tileInfo.Title, Content: &dashboardpb.Widget_Text{ Text: &dashboardpb.Text{ Content: tileInfo.TextContent, - Format: dashboardpb.Text_MARKDOWN, + Format: dashboardpb.Text_MARKDOWN, }, }, }, @@ -382,7 +379,7 @@ func getCreateMonitoringDashboardRequest( } // create bottom independent metrics tiles - if createAggIndependentBottomMetrics!= nil{ + if createAggIndependentBottomMetrics != nil { independentBottomMetricsTiles := createAggIndependentBottomMetrics(resourceIds) heightOffset += setWidgetPositions(independentBottomMetricsTiles, heightOffset) mosaicLayoutTiles = append(mosaicLayoutTiles, independentBottomMetricsTiles...) diff --git a/common/metrics/queries.go b/common/metrics/queries.go index f5f124d44e..4c7dd6ea39 100644 --- a/common/metrics/queries.go +++ b/common/metrics/queries.go @@ -4,16 +4,13 @@ // 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 +// 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 utils contains common helper functions used across multiple other packages. -// Utils should not import any Spanner migration tool packages. package metrics // Defines queries for Monitoring Dashboard Metrics @@ -90,8 +87,8 @@ const ( "filter && (%s) | group_by 1m, [value_estimated_backlog_processing_time_mean: " + "mean(value.estimated_backlog_processing_time)] | every 1m | group_by [], [value_estimated_backlog_processing_time_mean_mean: " + "mean(value_estimated_backlog_processing_time_mean)]" - dataflowAggPerShardCpuUtil = "fetch gce_instance | metric 'compute.googleapis.com/instance/cpu/utilization' | filter (%s) " + - "| group_by 1m, [value_utilization_mean: mean(value.utilization)] | every 1m | group_by [metadata.user_labels.dataflow_job_id]," + + dataflowAggPerShardCpuUtil = "fetch gce_instance | metric 'compute.googleapis.com/instance/cpu/utilization' | filter (%s) " + + "| group_by 1m, [value_utilization_mean: mean(value.utilization)] | every 1m | group_by [metadata.user_labels.dataflow_job_id]," + " [value_utilization_mean_percentile: percentile(value_utilization_mean, 50)]" datastreamAggThroughputQuery = "fetch datastream.googleapis.com/Stream | metric 'datastream.googleapis.com/stream/event_count' | " + "filter (%s) | align rate(1m) | every 1m | group_by [], [value_event_count_aggregate: aggregate(value.event_count)]" diff --git a/common/utils/dataflow_utils.go b/common/utils/dataflow_utils.go deleted file mode 100644 index a5d4ac09f9..0000000000 --- a/common/utils/dataflow_utils.go +++ /dev/null @@ -1,111 +0,0 @@ -// Copyright 2023 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 utils contains common helper functions used across multiple other packages. -// Utils should not import any Spanner migration tool packages. -package utils - -import ( - "encoding/json" - "fmt" - "sort" - "strings" - - "cloud.google.com/go/dataflow/apiv1beta3/dataflowpb" - "golang.org/x/exp/maps" -) - -// Generate the equivalent gCloud CLI command to launch a dataflow job with the same parameters and environment flags -// as the input body. -func GetGcloudDataflowCommand(req *dataflowpb.LaunchFlexTemplateRequest) string { - lp := req.LaunchParameter - templatePath := lp.Template.(*dataflowpb.LaunchFlexTemplateParameter_ContainerSpecGcsPath).ContainerSpecGcsPath - cmd := fmt.Sprintf("gcloud dataflow flex-template run %s --project=%s --region=%s --template-file-gcs-location=%s %s %s", - lp.JobName, req.ProjectId, req.Location, templatePath, getEnvironmentFlags(lp.Environment), getParametersFlag(lp.Parameters)) - return strings.Trim(cmd, " ") -} - -// Generate the equivalent parameter flag string, returning empty string if none are specified. -func getParametersFlag(parameters map[string]string) string { - if len(parameters) == 0 { - return "" - } - params := "" - keys := maps.Keys(parameters) - sort.Strings(keys) - for _, k := range keys { - params = params + k + "=" + parameters[k] + "," - } - params = strings.TrimSuffix(params, ",") - return fmt.Sprintf("--parameters %s", params) -} - -// We don't populate all flags in the API because certain flags (like AutoscalingAlgorithm, DumpHeapOnOom etc.) -// are not supported in gCloud. -func getEnvironmentFlags(environment *dataflowpb.FlexTemplateRuntimeEnvironment) string { - flag := "" - if environment.NumWorkers != 0 { - flag += fmt.Sprintf("--num-workers %d ", environment.NumWorkers) - } - if environment.MaxWorkers != 0 { - flag += fmt.Sprintf("--max-workers %d ", environment.MaxWorkers) - } - if environment.ServiceAccountEmail != "" { - flag += fmt.Sprintf("--service-account-email %s ", environment.ServiceAccountEmail) - } - if environment.TempLocation != "" { - flag += fmt.Sprintf("--temp-location %s ", environment.TempLocation) - } - if environment.MachineType != "" { - flag += fmt.Sprintf("--worker-machine-type %s ", environment.MachineType) - } - if environment.AdditionalExperiments != nil && len(environment.AdditionalExperiments) > 0 { - flag += fmt.Sprintf("--additional-experiments %s ", strings.Join(environment.AdditionalExperiments, ",")) - } - if environment.Network != "" { - flag += fmt.Sprintf("--network %s ", environment.Network) - } - if environment.Subnetwork != "" { - flag += fmt.Sprintf("--subnetwork %s ", environment.Subnetwork) - } - if environment.AdditionalUserLabels != nil && len(environment.AdditionalUserLabels) > 0 { - jsonByteStr, err := json.Marshal(environment.AdditionalUserLabels) - // If error is not nil, omit this flag and move on. We don't need error handling here. - if err == nil { - flag += fmt.Sprintf("--additional-user-labels %s ", string(jsonByteStr)) - } - } - if environment.KmsKeyName != "" { - flag += fmt.Sprintf("--dataflow-kms-key %s ", environment.KmsKeyName) - } - if environment.IpConfiguration == dataflowpb.WorkerIPAddressConfiguration_WORKER_IP_PRIVATE { - flag += "--disable-public-ips " - } - if environment.WorkerRegion != "" { - flag += fmt.Sprintf("--worker-region %s ", environment.WorkerRegion) - } - if environment.WorkerZone != "" { - flag += fmt.Sprintf("--worker-zone %s ", environment.WorkerZone) - } - if environment.EnableStreamingEngine { - flag += "--enable-streaming-engine " - } - if environment.FlexrsGoal != dataflowpb.FlexResourceSchedulingGoal_FLEXRS_UNSPECIFIED { - flag += fmt.Sprintf("--flexrs-goal %s ", environment.FlexrsGoal) - } - if environment.StagingLocation != "" { - flag += fmt.Sprintf("--staging-location %s ", environment.StagingLocation) - } - return strings.Trim(flag, " ") -} diff --git a/streaming/streaming.go b/streaming/streaming.go index 481770b255..1d95521993 100644 --- a/streaming/streaming.go +++ b/streaming/streaming.go @@ -33,6 +33,7 @@ import ( resourcemanager "cloud.google.com/go/resourcemanager/apiv3" resourcemanagerpb "cloud.google.com/go/resourcemanager/apiv3/resourcemanagerpb" + dataflowaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/dataflow" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" "github.com/GoogleCloudPlatform/spanner-migration-tool/internal" @@ -722,7 +723,8 @@ func LaunchDataflowJob(ctx context.Context, targetProfile profiles.TargetProfile fmt.Printf("flexTemplateRequest: %+v\n", req) return internal.DataflowOutput{}, fmt.Errorf("unable to launch template: %v", err) } - gcloudDfCmd := utils.GetGcloudDataflowCommand(req) + // Refactor to use accessor return value. + gcloudDfCmd := dataflowaccessor.GetGcloudDataflowCommandFromRequest(req) logger.Log.Debug(fmt.Sprintf("\nEquivalent gCloud command for job %s:\n%s\n\n", req.LaunchParameter.JobName, gcloudDfCmd)) return internal.DataflowOutput{JobID: respDf.Job.Id, GCloudCmd: gcloudDfCmd}, nil } diff --git a/testing/common/utils/dataflow_utils_test.go b/testing/common/utils/dataflow_utils_test.go deleted file mode 100644 index 948748c46f..0000000000 --- a/testing/common/utils/dataflow_utils_test.go +++ /dev/null @@ -1,117 +0,0 @@ -// Copyright 2023 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. - -// TODO: Refactor this file and other integration tests by moving all common code -// to remove redundancy. - -package utils_test - -import ( - "os" - "testing" - - "cloud.google.com/go/dataflow/apiv1beta3/dataflowpb" - "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" - "github.com/stretchr/testify/assert" -) - -func TestMain(m *testing.M) { - res := m.Run() - os.Exit(res) -} - -func getTemplateDfRequest() *dataflowpb.LaunchFlexTemplateRequest { - launchParameters := &dataflowpb.LaunchFlexTemplateParameter{ - JobName: "test-job", - Template: &dataflowpb.LaunchFlexTemplateParameter_ContainerSpecGcsPath{ContainerSpecGcsPath: "gs://template/Cloud_Datastream_to_Spanner"}, - Parameters: map[string]string{ - "inputFilePattern": "gs://inputFilePattern", - "streamName": "my-stream", - "instanceId": "my-instance", - "databaseId": "my-dbName", - "sessionFilePath": "gs://session.json", - "deadLetterQueueDirectory": "gs://dlq", - "transformationContextFilePath": "gs://transformationContext.json", - "directoryWatchDurationInMinutes": "480", // Setting directory watch timeout to 8 hours - }, - Environment: &dataflowpb.FlexTemplateRuntimeEnvironment{ - MaxWorkers: 50, - NumWorkers: 10, - ServiceAccountEmail: "svc-account@google.com", - TempLocation: "gs://temp-location", - MachineType: "n2-standard-16", - AdditionalExperiments: []string{"use_runner_V2", "test-experiment"}, - Network: "my-network", - Subnetwork: "my-subnetwork", - AdditionalUserLabels: map[string]string{"name": "wrench", "count": "3"}, - KmsKeyName: "sample-kms-key", - IpConfiguration: dataflowpb.WorkerIPAddressConfiguration_WORKER_IP_PRIVATE, - WorkerRegion: "test-worker-region", - WorkerZone: "test-worker-zone", - EnableStreamingEngine: true, - FlexrsGoal: 1, - StagingLocation: "gs://staging-location", - }, - } - req := &dataflowpb.LaunchFlexTemplateRequest{ - ProjectId: "test-project", - LaunchParameter: launchParameters, - Location: "us-central1", - } - return req -} - -func TestGcloudCmdWithAllParams(t *testing.T) { - - req := getTemplateDfRequest() - expectedCmd := "gcloud dataflow flex-template run test-job " + - "--project=test-project --region=us-central1 " + - "--template-file-gcs-location=gs://template/Cloud_Datastream_to_Spanner " + - "--num-workers 10 --max-workers 50 --service-account-email svc-account@google.com " + - "--temp-location gs://temp-location --worker-machine-type n2-standard-16 " + - "--additional-experiments use_runner_V2,test-experiment --network my-network " + - "--subnetwork my-subnetwork --additional-user-labels {\"count\":\"3\",\"name\":\"wrench\"} " + - "--dataflow-kms-key sample-kms-key --disable-public-ips --worker-region test-worker-region " + - "--worker-zone test-worker-zone --enable-streaming-engine " + - "--flexrs-goal FLEXRS_SPEED_OPTIMIZED --staging-location gs://staging-location " + - "--parameters databaseId=my-dbName,deadLetterQueueDirectory=gs://dlq," + - "directoryWatchDurationInMinutes=480,inputFilePattern=gs://inputFilePattern," + - "instanceId=my-instance,sessionFilePath=gs://session.json,streamName=my-stream," + - "transformationContextFilePath=gs://transformationContext.json" - assert.Equal(t, expectedCmd, utils.GetGcloudDataflowCommand(req)) -} - -func TestGcloudCmdWithPartialParams(t *testing.T) { - - req := getTemplateDfRequest() - req.LaunchParameter.Parameters = make(map[string]string) - req.LaunchParameter.Environment.FlexrsGoal = 0 - req.LaunchParameter.Environment.IpConfiguration = 0 - req.LaunchParameter.Environment.EnableStreamingEngine = false - req.LaunchParameter.Environment.AdditionalExperiments = []string{} - req.LaunchParameter.Environment.AdditionalUserLabels = make(map[string]string) - req.LaunchParameter.Environment.WorkerRegion = "" - req.LaunchParameter.Environment.NumWorkers = 0 - - expectedCmd := "gcloud dataflow flex-template run test-job " + - "--project=test-project --region=us-central1 " + - "--template-file-gcs-location=gs://template/Cloud_Datastream_to_Spanner " + - "--max-workers 50 --service-account-email svc-account@google.com " + - "--temp-location gs://temp-location --worker-machine-type n2-standard-16 " + - "--network my-network --subnetwork my-subnetwork " + - "--dataflow-kms-key sample-kms-key " + - "--worker-zone test-worker-zone " + - "--staging-location gs://staging-location" - assert.Equal(t, expectedCmd, utils.GetGcloudDataflowCommand(req)) -} From fca5dcf9daabb14b5ac8d4e9837650a7f3d92903 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Tue, 9 Jan 2024 14:16:18 +0530 Subject: [PATCH 02/25] Add accessors for storage and spanner. --- .../clients/spanner/admin/admin_client.go | 39 ++++ .../clients/spanner/client/spanner_client.go | 39 ++++ .../instanceadmin/spanner_instance_admin.go | 39 ++++ accessors/clients/storage/storage_client.go | 39 ++++ accessors/spanner/spanner_accessor.go | 203 ++++++++++++++++++ accessors/storage/storage_accessor.go | 192 +++++++++++++++++ cmd/data.go | 3 +- common/utils/storage_utils.go | 41 ++++ common/utils/utils.go | 87 +------- conversion/conversion.go | 43 +--- streaming/streaming.go | 46 +--- .../spanner/spanner_accessor_test.go | 129 +++++++++++ testing/conversion/conversion_test.go | 23 -- webv2/helpers/helpers.go | 4 +- webv2/profile/profile.go | 3 +- webv2/session/session_service.go | 4 +- webv2/web.go | 13 +- 17 files changed, 758 insertions(+), 189 deletions(-) create mode 100644 accessors/clients/spanner/admin/admin_client.go create mode 100644 accessors/clients/spanner/client/spanner_client.go create mode 100644 accessors/clients/spanner/instanceadmin/spanner_instance_admin.go create mode 100644 accessors/clients/storage/storage_client.go create mode 100644 accessors/spanner/spanner_accessor.go create mode 100644 accessors/storage/storage_accessor.go create mode 100644 common/utils/storage_utils.go create mode 100644 testing/accessors/spanner/spanner_accessor_test.go diff --git a/accessors/clients/spanner/admin/admin_client.go b/accessors/clients/spanner/admin/admin_client.go new file mode 100644 index 0000000000..c0ba8aed5f --- /dev/null +++ b/accessors/clients/spanner/admin/admin_client.go @@ -0,0 +1,39 @@ +// 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 + +func GetOrCreateClient(ctx context.Context) (*database.DatabaseAdminClient, error) { + var err error + if spannerAdminClient == nil { + once.Do(func() { + spannerAdminClient, err = database.NewDatabaseAdminClient(ctx) + }) + if err != nil { + return nil, fmt.Errorf("failed to create spanner admin client: %v", err) + } + return spannerAdminClient, nil + } + return spannerAdminClient, nil +} diff --git a/accessors/clients/spanner/client/spanner_client.go b/accessors/clients/spanner/client/spanner_client.go new file mode 100644 index 0000000000..bbb3e252c1 --- /dev/null +++ b/accessors/clients/spanner/client/spanner_client.go @@ -0,0 +1,39 @@ +// 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 spannerclient + +import ( + "context" + "fmt" + "sync" + + sp "cloud.google.com/go/spanner" +) + +var once sync.Once +var spannerClient *sp.Client + +func GetOrCreateClient(ctx context.Context, dbURI string) (*sp.Client, error) { + var err error + if spannerClient == nil || spannerClient.DatabaseName() != dbURI { + once.Do(func() { + spannerClient, err = sp.NewClient(ctx, dbURI) + }) + if err != nil { + return nil, fmt.Errorf("failed to create spanner database client: %v", err) + } + return spannerClient, nil + } + return spannerClient, nil +} diff --git a/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go b/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go new file mode 100644 index 0000000000..324da4ac32 --- /dev/null +++ b/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go @@ -0,0 +1,39 @@ +// 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 spinstanceadmin + +import ( + "context" + "fmt" + "sync" + + instance "cloud.google.com/go/spanner/admin/instance/apiv1" +) + +var once sync.Once +var instanceAdminClient *instance.InstanceAdminClient + +func GetOrCreateClient(ctx context.Context) (*instance.InstanceAdminClient, error) { + var err error + if instanceAdminClient == nil { + once.Do(func() { + instanceAdminClient, err = instance.NewInstanceAdminClient(ctx) + }) + if err != nil { + return nil, fmt.Errorf("failed to create spanner instance admin client: %v", err) + } + return instanceAdminClient, nil + } + return instanceAdminClient, nil +} diff --git a/accessors/clients/storage/storage_client.go b/accessors/clients/storage/storage_client.go new file mode 100644 index 0000000000..569841d299 --- /dev/null +++ b/accessors/clients/storage/storage_client.go @@ -0,0 +1,39 @@ +// 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 storageclient + +import ( + "context" + "fmt" + "sync" + + "cloud.google.com/go/storage" +) + +var once sync.Once +var gcsClient *storage.Client + +func GetOrCreateClient(ctx context.Context) (*storage.Client, error) { + var err error + if gcsClient == nil { + once.Do(func() { + gcsClient, err = storage.NewClient(ctx) + }) + if err != nil { + return nil, fmt.Errorf("failed to create storage client: %v", err) + } + return gcsClient, nil + } + return gcsClient, nil +} diff --git a/accessors/spanner/spanner_accessor.go b/accessors/spanner/spanner_accessor.go new file mode 100644 index 0000000000..101a9dd93e --- /dev/null +++ b/accessors/spanner/spanner_accessor.go @@ -0,0 +1,203 @@ +// 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 spanneracc + +import ( + "context" + "fmt" + "strings" + "time" + + "cloud.google.com/go/spanner" + "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" + "cloud.google.com/go/spanner/admin/instance/apiv1/instancepb" + spanneradmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/admin" + spannerclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/client" + spinstanceadmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/instanceadmin" + "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" + "google.golang.org/api/iterator" +) + +func GetDatabase(ctx context.Context, dbURI string) (*databasepb.Database, error) { + adminClient, err := spanneradmin.GetOrCreateClient(ctx) + if err != nil { + return nil, err + } + return adminClient.GetDatabase(ctx, &databasepb.GetDatabaseRequest{Name: dbURI}) +} + +func GetDatabaseDialect(ctx context.Context, dbURI string) (string, error) { + result, err := GetDatabase(ctx, dbURI) + if err != nil { + return "", fmt.Errorf("cannot connect to database: %v", err) + } + return strings.ToLower(result.DatabaseDialect.String()), nil +} + +// CheckExistingDb checks whether the database with dbURI exists or not. +// If API call doesn't respond then user is informed after every 5 minutes on command line. +func CheckExistingDb(ctx context.Context, dbURI string) (bool, error) { + gotResponse := make(chan bool) + var err error + adminClient, err := spanneradmin.GetOrCreateClient(ctx) + if err != nil { + return false, err + } + go func() { + _, err = adminClient.GetDatabase(ctx, &databasepb.GetDatabaseRequest{Name: dbURI}) + gotResponse <- true + }() + for { + select { + case <-time.After(5 * time.Minute): + fmt.Println("WARNING! API call not responding: make sure that spanner api endpoint is configured properly") + case <-gotResponse: + if err != nil { + if utils.ContainsAny(strings.ToLower(err.Error()), []string{"database not found"}) { + return false, nil + } + return false, fmt.Errorf("can't get database info: %s", err) + } + return true, nil + } + } +} + +func CreateEmptyDatabase(ctx context.Context, dbURI string) error { + adminClient, err := spanneradmin.GetOrCreateClient(ctx) + if err != nil { + return err + } + project, instance, dbName := utils.ParseDbURI(dbURI) + req := &databasepb.CreateDatabaseRequest{ + Parent: fmt.Sprintf("projects/%s/instances/%s", project, instance), + CreateStatement: "CREATE DATABASE `" + dbName + "`", + } + op, err := adminClient.CreateDatabase(ctx, req) + if err != nil { + return fmt.Errorf("can't build CreateDatabaseRequest: %w", utils.AnalyzeError(err, dbURI)) + } + if _, err := op.Wait(ctx); err != nil { + return fmt.Errorf("createDatabase call failed: %w", utils.AnalyzeError(err, dbURI)) + } + return nil +} + +func GetSpannerLeaderLocation(ctx context.Context, instanceURI string) (string, error) { + instanceClient, err := spinstanceadmin.GetOrCreateClient(ctx) + if err != nil { + return "", err + } + instanceInfo, err := instanceClient.GetInstance(ctx, &instancepb.GetInstanceRequest{Name: instanceURI}) + if err != nil { + return "", err + } + instanceConfig, err := instanceClient.GetInstanceConfig(ctx, &instancepb.GetInstanceConfigRequest{Name: instanceInfo.Config}) + if err != nil { + return "", err + + } + for _, replica := range instanceConfig.Replicas { + if replica.DefaultLeaderLocation { + return replica.Location, nil + } + } + return "", fmt.Errorf("no leader found for spanner instance %s while trying fetch location", instanceURI) +} + +func CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI string) (bool, error) { + spClient, err := spannerclient.GetOrCreateClient(ctx, dbURI) + if err != nil { + return false, err + } + stmt := spanner.Statement{ + SQL: `SELECT * FROM information_schema.change_streams`, + } + iter := spClient.Single().Query(ctx, stmt) + defer iter.Stop() + var cs_catalog, cs_schema, cs_name string + var coversAll bool + csExists := false + for { + row, err := iter.Next() + if err == iterator.Done { + break + } + if err != nil { + return false, fmt.Errorf("couldn't read row from change_streams table: %w", err) + } + err = row.Columns(&cs_catalog, &cs_schema, &cs_name, &coversAll) + if err != nil { + return false, fmt.Errorf("can't scan row from change_streams table: %v", err) + } + if cs_name == changeStreamName { + csExists = true + break + } + } + return csExists, nil +} + +func ValidateChangeStreamOptions(ctx context.Context, changeStreamName, dbURI string) error { + spClient, err := spannerclient.GetOrCreateClient(ctx, dbURI) + if err != nil { + return err + } + // Validate if change stream options are set correctly. + stmt := spanner.Statement{ + SQL: `SELECT option_value FROM information_schema.change_stream_options + WHERE change_stream_name = @p1 AND option_name = 'value_capture_type'`, + Params: map[string]interface{}{ + "p1": changeStreamName, + }, + } + iter := spClient.Single().Query(ctx, stmt) + defer iter.Stop() + var option_value string + for { + row, err := iter.Next() + if err == iterator.Done { + break + } + if err != nil { + return fmt.Errorf("couldn't read row from change_stream_options table: %w", err) + } + err = row.Columns(&option_value) + if err != nil { + return fmt.Errorf("can't scan row from change_stream_options table: %v", err) + } + if option_value != "NEW_ROW" { + return fmt.Errorf("VALUE_CAPTURE_TYPE for changestream %s is not NEW_ROW. Please update the changestream option or create a new one", changeStreamName) + } + } + return nil +} + +func CreateChangeStream(ctx context.Context, changeStreamName, dbURI string) error { + spClient, _ := spanneradmin.GetOrCreateClient(ctx) + op, err := spClient.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{ + Database: dbURI, + // TODO: create change stream for only the tables present in Spanner. + Statements: []string{fmt.Sprintf("CREATE CHANGE STREAM %s FOR ALL OPTIONS (value_capture_type = 'NEW_ROW')", changeStreamName)}, + }) + if err != nil { + return fmt.Errorf("cannot submit request create change stream request: %v", err) + } + if err := op.Wait(ctx); err != nil { + return fmt.Errorf("could not update database ddl: %v", err) + } else { + fmt.Println("Successfully created changestream", changeStreamName) + } + return nil +} diff --git a/accessors/storage/storage_accessor.go b/accessors/storage/storage_accessor.go new file mode 100644 index 0000000000..0d880b978a --- /dev/null +++ b/accessors/storage/storage_accessor.go @@ -0,0 +1,192 @@ +// 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 storageacc + +import ( + "context" + "fmt" + "io" + "os" + "strings" + + "cloud.google.com/go/storage" + storageclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/storage" + "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" + "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" + "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" + "google.golang.org/api/googleapi" +) + +func CreateGCSBucket(ctx context.Context, bucketName, projectID, location string) error { + return createGCSBucketUtil(ctx, bucketName, projectID, location, nil, 0) +} + +func CreateGCSBucketWithLifecycle(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { + return createGCSBucketUtil(ctx, bucketName, projectID, location, matchesPrefix, ttl) +} + +func createGCSBucketUtil(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { + client, err := storageclient.GetOrCreateClient(ctx) + if err != nil { + return err + } + bucket := client.Bucket(bucketName) + attrs := storage.BucketAttrs{ + Location: location, + } + if ttl > 0 { + attrs.Lifecycle = storage.Lifecycle{ + Rules: []storage.LifecycleRule{ + { + Action: storage.LifecycleAction{Type: "Delete"}, + Condition: storage.LifecycleCondition{ + AgeInDays: ttl, + // The prefixes should not contain the bucket names and starting slash. + // For object gs://my_bucket/pictures/paris_2022.jpg, + // you would use a condition such as "matchesPrefix":["pictures/paris_"]. + MatchesPrefix: matchesPrefix, + }, + }, + }, + } + } + + if err := bucket.Create(ctx, projectID, &attrs); err != nil { + if e, ok := err.(*googleapi.Error); ok { + // Ignoring the bucket already exists error. + if e.Code != 409 { + return fmt.Errorf("failed to create bucket: %v", err) + } else { + fmt.Printf("Using the existing bucket: %v \n", bucketName) + } + } else { + return fmt.Errorf("failed to create bucket: %v", err) + } + + } else { + fmt.Printf("Created new GCS bucket: %v\n", bucketName) + } + return nil +} + +// Applies the bucket lifecycle with delete rule. Only accepts the Age and +// prefix rule conditions as it is only used for the Datastream destination +// bucket currently. +func EnableBucketLifecycleDeleteRule(ctx context.Context, bucketName string, matchesPrefix []string, ttl int64) error { + client, err := storageclient.GetOrCreateClient(ctx) + if err != nil { + return fmt.Errorf("could not create client while enabling lifecycle: %w", err) + } + + for i, str := range matchesPrefix { + matchesPrefix[i] = strings.TrimPrefix(str, "/") + } + bucket := client.Bucket(bucketName) + bucketAttrsToUpdate := storage.BucketAttrsToUpdate{ + Lifecycle: &storage.Lifecycle{ + Rules: []storage.LifecycleRule{ + { + Action: storage.LifecycleAction{Type: "Delete"}, + Condition: storage.LifecycleCondition{ + AgeInDays: ttl, + // The prefixes should not contain the bucket names and starting slash. + // For object gs://my_bucket/pictures/paris_2022.jpg, + // you would use a condition such as "matchesPrefix":["pictures/paris_"]. + MatchesPrefix: matchesPrefix, + }, + }, + }, + }, + } + + attrs, err := bucket.Update(ctx, bucketAttrsToUpdate) + if err != nil { + return fmt.Errorf("could not bucket with lifecycle: %w", err) + } + logger.Log.Info(fmt.Sprintf("Added lifecycle rule to bucket %v\n. Rule Action: %v\t Rule Condition: %v\n", + bucketName, attrs.Lifecycle.Rules[0].Action, attrs.Lifecycle.Rules[0].Condition)) + return nil +} + +// UploadLocalFileToGCS uploads an object. +func UploadLocalFileToGCS(ctx context.Context, filePath, fileName, localFilePath string) error { + data, err := os.ReadFile(localFilePath) + if err != nil { + return fmt.Errorf("could not read file %s: %w", localFilePath, err) + } + return WriteDataToGCS(ctx, filePath, fileName, string(data)) +} + +func WriteDataToGCS(ctx context.Context, filePath, fileName, data string) error { + client, err := storageclient.GetOrCreateClient(ctx) + if err != nil { + return fmt.Errorf("could not create client while uploading to GCS: %w", err) + } + + u, err := utils.ParseGCSFilePath(filePath) + if err != nil { + return fmt.Errorf("parseFilePath: unable to parse file path: %v", err) + } + bucketName := u.Host + bucket := client.Bucket(bucketName) + obj := bucket.Object(u.Path[1:] + fileName) + + w := obj.NewWriter(ctx) + if _, err := fmt.Fprint(w, data); err != nil { + fmt.Printf("Failed to write to Cloud Storage: %s", filePath) + return err + } + if err := w.Close(); err != nil { + fmt.Printf("Failed to close GCS file: %s", filePath) + return err + } + return nil +} + +func ReadGcsFile(ctx context.Context, filePath string) (string, error) { + client, err := storageclient.GetOrCreateClient(ctx) + if err != nil { + return "", fmt.Errorf("could not create client: %w", err) + } + + u, err := utils.ParseGCSFilePath(filePath) + if err != nil { + return "", fmt.Errorf("unable to parse file path: %v", err) + } + bucketName := u.Host + bucket := client.Bucket(bucketName) + obj := bucket.Object(u.Path[1:]) + + rc, err := obj.NewReader(ctx) + if err != nil { + return "", err + } + defer rc.Close() + buf := new(strings.Builder) + if _, err := io.Copy(buf, rc); err != nil { + return "", err + } + return buf.String(), nil +} + +func ReadAnyFile(ctx context.Context, filePath string) (string, error) { + if strings.HasPrefix(filePath, constants.GCS_FILE_PREFIX) { + return ReadGcsFile(ctx, filePath) + } + buf, err := os.ReadFile(filePath) + if err != nil { + return "", err + } + return string(buf), nil +} diff --git a/cmd/data.go b/cmd/data.go index c908f8ad94..6190ed3696 100644 --- a/cmd/data.go +++ b/cmd/data.go @@ -26,6 +26,7 @@ import ( sp "cloud.google.com/go/spanner" database "cloud.google.com/go/spanner/admin/database/apiv1" + spanneracc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" "github.com/GoogleCloudPlatform/spanner-migration-tool/conversion" @@ -178,7 +179,7 @@ func (cmd *DataCmd) Execute(ctx context.Context, f *flag.FlagSet, _ ...interface // validateExistingDb validates that the existing spanner schema is in accordance with the one specified in the session file. func validateExistingDb(ctx context.Context, spDialect, dbURI string, adminClient *database.DatabaseAdminClient, client *sp.Client, conv *internal.Conv) error { - dbExists, err := conversion.CheckExistingDb(ctx, adminClient, dbURI) + dbExists, err := spanneracc.CheckExistingDb(ctx, dbURI) if err != nil { err = fmt.Errorf("can't verify target database: %v", err) return err diff --git a/common/utils/storage_utils.go b/common/utils/storage_utils.go new file mode 100644 index 0000000000..3429c6d293 --- /dev/null +++ b/common/utils/storage_utils.go @@ -0,0 +1,41 @@ +// 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 utils contains common helper functions used across multiple other packages. +// Utils should not import any Spanner migration tool packages. +package utils + +import ( + "fmt" + "net/url" + + "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" +) + +func ParseGCSFilePath(filePath string) (*url.URL, error) { + if len(filePath) == 0 { + return nil, fmt.Errorf("found empty GCS path") + } + if filePath[len(filePath)-1] != '/' { + filePath = filePath + "/" + } + u, err := url.Parse(filePath) + if err != nil { + return nil, fmt.Errorf("parseFilePath: unable to parse file path %s", filePath) + } + if u.Scheme != constants.GCS_SCHEME { + return nil, fmt.Errorf("not a valid GCS path: %s, should start with 'gs'", filePath) + } + return u, nil +} diff --git a/common/utils/utils.go b/common/utils/utils.go index 70ec53f882..932da34f11 100644 --- a/common/utils/utils.go +++ b/common/utils/utils.go @@ -19,11 +19,11 @@ package utils import ( "bufio" "context" + "crypto/rand" "fmt" "io" "io/ioutil" "log" - "math/rand" "net/url" "os" "os/exec" @@ -37,6 +37,7 @@ import ( sp "cloud.google.com/go/spanner" database "cloud.google.com/go/spanner/admin/database/apiv1" instance "cloud.google.com/go/spanner/admin/instance/apiv1" + "cloud.google.com/go/spanner/admin/instance/apiv1/instancepb" "cloud.google.com/go/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/internal" @@ -44,10 +45,8 @@ import ( "github.com/GoogleCloudPlatform/spanner-migration-tool/sources/spanner" "github.com/GoogleCloudPlatform/spanner-migration-tool/spanner/ddl" "golang.org/x/crypto/ssh/terminal" - "google.golang.org/api/googleapi" "google.golang.org/api/iterator" "google.golang.org/api/option" - instancepb "google.golang.org/genproto/googleapis/spanner/admin/instance/v1" ) // IOStreams is a struct that contains the file descriptor for dumpFile. @@ -173,82 +172,6 @@ func PreloadGCSFiles(tables []ManifestTable) ([]ManifestTable, error) { return tables, nil } -func ParseGCSFilePath(filePath string) (*url.URL, error) { - if len(filePath) == 0 { - return nil, fmt.Errorf("found empty GCS path") - } - if filePath[len(filePath)-1] != '/' { - filePath = filePath + "/" - } - u, err := url.Parse(filePath) - if err != nil { - return nil, fmt.Errorf("parseFilePath: unable to parse file path %s", filePath) - } - if u.Scheme != constants.GCS_SCHEME { - return nil, fmt.Errorf("not a valid GCS path: %s, should start with 'gs'", filePath) - } - return u, nil -} - -func WriteToGCS(filePath, fileName, data string) error { - ctx := context.Background() - - client, err := storage.NewClient(ctx) - if err != nil { - fmt.Printf("Failed to create GCS client") - return err - } - defer client.Close() - u, err := ParseGCSFilePath(filePath) - if err != nil { - return fmt.Errorf("parseFilePath: unable to parse file path: %v", err) - } - bucketName := u.Host - bucket := client.Bucket(bucketName) - obj := bucket.Object(u.Path[1:] + fileName) - - w := obj.NewWriter(ctx) - if _, err := fmt.Fprint(w, data); err != nil { - fmt.Printf("Failed to write to Cloud Storage: %s", filePath) - return err - } - if err := w.Close(); err != nil { - fmt.Printf("Failed to close GCS file: %s", filePath) - return err - } - return nil -} - -func CreateGCSBucket(bucketName, projectID, location string) error { - ctx := context.Background() - - client, err := storage.NewClient(ctx) - if err != nil { - return fmt.Errorf("failed to create GCS client: %v", err) - } - defer client.Close() - bucket := client.Bucket(bucketName) - attrs := storage.BucketAttrs{ - Location: location, - } - if err := bucket.Create(ctx, projectID, &attrs); err != nil { - if e, ok := err.(*googleapi.Error); ok { - // Ignoring the bucket already exists error. - if e.Code != 409 { - return fmt.Errorf("failed to create bucket: %v", err) - } else { - fmt.Printf("Using the existing bucket: %v \n", bucketName) - } - } else { - return fmt.Errorf("failed to create bucket: %v", err) - } - - } else { - fmt.Printf("Created new GCS bucket: %v\n", bucketName) - } - return nil -} - // GetProject returns the cloud project we should use for accessing Spanner. // Use environment variable GCLOUD_PROJECT if it is set. // Otherwise, use the default project returned from gcloud. @@ -347,6 +270,12 @@ func GenerateName(prefix string) (string, error) { return fmt.Sprintf("%s_%x-%x", prefix, b[0:2], b[2:4]), nil } +func GenerateHashStr() string { + b := make([]byte, 4) + rand.Read(b) + return fmt.Sprintf("%x-%x", b[0:2], b[2:4]) +} + // parseURI parses an unknown URI string that could be a database, instance or project URI. func parseURI(URI string) (project, instance, dbName string) { project, instance, dbName = "", "", "" diff --git a/conversion/conversion.go b/conversion/conversion.go index 20216dfea4..1608f98054 100644 --- a/conversion/conversion.go +++ b/conversion/conversion.go @@ -44,6 +44,8 @@ import ( datastream "cloud.google.com/go/datastream/apiv1" sp "cloud.google.com/go/spanner" database "cloud.google.com/go/spanner/admin/database/apiv1" + spanneracc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" + storageacc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/metrics" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" @@ -355,7 +357,7 @@ func dataFromDatabase(ctx context.Context, sourceProfile profiles.SourceProfile, // Try to apply lifecycle rule to Datastream destination bucket. gcsConfig := streamingCfg.GcsCfg if gcsConfig.TtlInDaysSet { - err = streaming.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = storageacc.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") @@ -476,7 +478,7 @@ func dataFromDatabaseForDataflowMigration(targetProfile profiles.TargetProfile, // Try to apply lifecycle rule to Datastream destination bucket. gcsConfig := streamingCfg.GcsCfg if gcsConfig.TtlInDaysSet { - err = streaming.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = storageacc.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") @@ -520,11 +522,11 @@ func dataFromDatabaseForDataflowMigration(targetProfile profiles.TargetProfile, // create monitoring aggregated dashboard for sharded migration aggMonitoringResources := metrics.MonitoringMetricsResources{ - ProjectId: targetProfile.Conn.Sp.Project, - SpannerInstanceId: targetProfile.Conn.Sp.Instance, - SpannerDatabaseId: targetProfile.Conn.Sp.Dbname, - ShardToShardResourcesMap: conv.Audit.StreamingStats.ShardToShardResourcesMap, - MigrationRequestId: conv.Audit.MigrationRequestId, + ProjectId: targetProfile.Conn.Sp.Project, + SpannerInstanceId: targetProfile.Conn.Sp.Instance, + SpannerDatabaseId: targetProfile.Conn.Sp.Dbname, + ShardToShardResourcesMap: conv.Audit.StreamingStats.ShardToShardResourcesMap, + MigrationRequestId: conv.Audit.MigrationRequestId, } aggRespDash, dashboardErr := aggMonitoringResources.CreateDataflowAggMonitoringDashboard(ctx) if dashboardErr != nil { @@ -798,7 +800,7 @@ func getSeekable(f *os.File) (*os.File, int64, error) { // VerifyDb checks whether the db exists and if it does, verifies if the schema is what we currently support. func VerifyDb(ctx context.Context, adminClient *database.DatabaseAdminClient, dbURI string) (dbExists bool, err error) { - dbExists, err = CheckExistingDb(ctx, adminClient, dbURI) + dbExists, err = spanneracc.CheckExistingDb(ctx, dbURI) if err != nil { return dbExists, err } @@ -808,31 +810,6 @@ func VerifyDb(ctx context.Context, adminClient *database.DatabaseAdminClient, db return dbExists, err } -// CheckExistingDb checks whether the database with dbURI exists or not. -// If API call doesn't respond then user is informed after every 5 minutes on command line. -func CheckExistingDb(ctx context.Context, adminClient *database.DatabaseAdminClient, dbURI string) (bool, error) { - gotResponse := make(chan bool) - var err error - go func() { - _, err = adminClient.GetDatabase(ctx, &adminpb.GetDatabaseRequest{Name: dbURI}) - gotResponse <- true - }() - for { - select { - case <-time.After(5 * time.Minute): - fmt.Println("WARNING! API call not responding: make sure that spanner api endpoint is configured properly") - case <-gotResponse: - if err != nil { - if utils.ContainsAny(strings.ToLower(err.Error()), []string{"database not found"}) { - return false, nil - } - return false, fmt.Errorf("can't get database info: %s", err) - } - return true, nil - } - } -} - // ValidateTables validates that all the tables in the database are empty. // It returns the name of the first non-empty table if found, and an empty string otherwise. func ValidateTables(ctx context.Context, client *sp.Client, spDialect string) (string, error) { diff --git a/streaming/streaming.go b/streaming/streaming.go index 1d95521993..7a96839d4f 100644 --- a/streaming/streaming.go +++ b/streaming/streaming.go @@ -34,10 +34,12 @@ import ( resourcemanager "cloud.google.com/go/resourcemanager/apiv3" resourcemanagerpb "cloud.google.com/go/resourcemanager/apiv3/resourcemanagerpb" dataflowaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/dataflow" + storageacc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" "github.com/GoogleCloudPlatform/spanner-migration-tool/internal" "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" + "github.com/GoogleCloudPlatform/spanner-migration-tool/profiles" "github.com/google/uuid" "github.com/googleapis/gax-go/v2" @@ -864,7 +866,7 @@ func StartDataflow(ctx context.Context, targetProfile profiles.TargetProfile, st if err != nil { return internal.DataflowOutput{}, fmt.Errorf("can't encode session state to JSON: %v", err) } - err = utils.WriteToGCS(streamingCfg.TmpDir, "session.json", string(convJSON)) + err = storageacc.WriteDataToGCS(ctx, streamingCfg.TmpDir, "session.json", string(convJSON)) if err != nil { return internal.DataflowOutput{}, fmt.Errorf("error while writing to GCS: %v", err) } @@ -875,7 +877,7 @@ func StartDataflow(ctx context.Context, targetProfile profiles.TargetProfile, st if err != nil { return internal.DataflowOutput{}, fmt.Errorf("failed to compute transformation context: %s", err.Error()) } - err = utils.WriteToGCS(streamingCfg.TmpDir, "transformationContext.json", string(transformationContext)) + err = storageacc.WriteDataToGCS(ctx, streamingCfg.TmpDir, "transformationContext.json", string(transformationContext)) if err != nil { return internal.DataflowOutput{}, fmt.Errorf("error while writing to GCS: %v", err) } @@ -885,43 +887,3 @@ func StartDataflow(ctx context.Context, targetProfile profiles.TargetProfile, st } return dfOutput, nil } - -// Applies the bucket lifecycle with delete rule. Only accepts the Age and -// prefix rule conditions as it is only used for the Datastream destination -// bucket currently. -func EnableBucketLifecycleDeleteRule(ctx context.Context, bucketName string, matchesPrefix []string, ttl int64) error { - client, err := storage.NewClient(ctx) - if err != nil { - return fmt.Errorf("could not create client while enabling lifecycle: %w", err) - } - defer client.Close() - - for i, str := range matchesPrefix { - matchesPrefix[i] = strings.TrimPrefix(str, "/") - } - bucket := client.Bucket(bucketName) - bucketAttrsToUpdate := storage.BucketAttrsToUpdate{ - Lifecycle: &storage.Lifecycle{ - Rules: []storage.LifecycleRule{ - { - Action: storage.LifecycleAction{Type: "Delete"}, - Condition: storage.LifecycleCondition{ - AgeInDays: ttl, - // The prefixes should not contain the bucket names and starting slash. - // For object gs://my_bucket/pictures/paris_2022.jpg, - // you would use a condition such as "matchesPrefix":["pictures/paris_"]. - MatchesPrefix: matchesPrefix, - }, - }, - }, - }, - } - - attrs, err := bucket.Update(ctx, bucketAttrsToUpdate) - if err != nil { - return fmt.Errorf("could not bucket with lifecycle: %w", err) - } - logger.Log.Info(fmt.Sprintf("Added lifecycle rule to bucket %v\n. Rule Action: %v\t Rule Condition: %v\n", - bucketName, attrs.Lifecycle.Rules[0].Action, attrs.Lifecycle.Rules[0].Condition)) - return nil -} diff --git a/testing/accessors/spanner/spanner_accessor_test.go b/testing/accessors/spanner/spanner_accessor_test.go new file mode 100644 index 0000000000..1d7e77bd72 --- /dev/null +++ b/testing/accessors/spanner/spanner_accessor_test.go @@ -0,0 +1,129 @@ +// 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. + +// TODO: Refactor this file and other integration tests by moving all common code +// to remove redundancy. + +package utils_test + +import ( + "context" + "flag" + "fmt" + "log" + "os" + "testing" + "time" + + database "cloud.google.com/go/spanner/admin/database/apiv1" + "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" + spanneracc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" + "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" + "github.com/GoogleCloudPlatform/spanner-migration-tool/conversion" + "github.com/GoogleCloudPlatform/spanner-migration-tool/internal" + "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" + "github.com/stretchr/testify/assert" + "go.uber.org/zap" +) + +var ( + projectID string + instanceID string + + ctx context.Context + databaseAdmin *database.DatabaseAdminClient +) + +func TestMain(m *testing.M) { + cleanup := initTests() + res := m.Run() + cleanup() + os.Exit(res) +} + +func init() { + logger.Log = zap.NewNop() +} + +func initTests() (cleanup func()) { + projectID = os.Getenv("SPANNER_MIGRATION_TOOL_TESTS_GCLOUD_PROJECT_ID") + instanceID = os.Getenv("SPANNER_MIGRATION_TOOL_TESTS_GCLOUD_INSTANCE_ID") + + ctx = context.Background() + flag.Parse() // Needed for testing.Short(). + noop := func() {} + + if testing.Short() { + log.Println("Unit test for UpdateDDLForeignKeys skipped in -short mode.") + return noop + } + + if projectID == "" { + log.Println("Unit test for UpdateDDLForeignKeys skipped: SPANNER_MIGRATION_TOOL_TESTS_GCLOUD_PROJECT_ID is missing") + return noop + } + + if instanceID == "" { + log.Println("Unit test for UpdateDDLForeignKeys skipped: SPANNER_MIGRATION_TOOL_TESTS_GCLOUD_INSTANCE_ID is missing") + return noop + } + + var err error + databaseAdmin, err = database.NewDatabaseAdminClient(ctx) + if err != nil { + log.Fatalf("cannot create databaseAdmin client: %v", err) + } + + return func() { + databaseAdmin.Close() + } +} + +func dropDatabase(t *testing.T, dbPath string) { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) + defer cancel() + // Drop the testing database. + if err := databaseAdmin.DropDatabase(ctx, &databasepb.DropDatabaseRequest{Database: dbPath}); err != nil { + t.Fatalf("failed to drop testing database %v: %v", dbPath, err) + } +} + +func TestCheckExistingDb(t *testing.T) { + onlyRunForEmulatorTest(t) + dbURI := fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectID, instanceID, "check-db-exists") + err := conversion.CreateDatabase(ctx, databaseAdmin, dbURI, internal.MakeConv(), os.Stdout, "", constants.BULK_MIGRATION) + if err != nil { + t.Fatal(err) + } + defer dropDatabase(t, dbURI) + testCases := []struct { + dbName string + dbExists bool + }{ + {"check-db-exists", true}, + {"check-db-does-not-exist", false}, + } + + for _, tc := range testCases { + dbExists, err := spanneracc.CheckExistingDb(ctx, fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectID, instanceID, tc.dbName)) + assert.Nil(t, err) + assert.Equal(t, tc.dbExists, dbExists) + } +} + +func onlyRunForEmulatorTest(t *testing.T) { + if os.Getenv("SPANNER_EMULATOR_HOST") == "" { + t.Skip("Skipping tests only running against the emulator.") + } +} diff --git a/testing/conversion/conversion_test.go b/testing/conversion/conversion_test.go index ff9a8e49ef..ba08519a71 100644 --- a/testing/conversion/conversion_test.go +++ b/testing/conversion/conversion_test.go @@ -248,29 +248,6 @@ func TestVerifyDb(t *testing.T) { } } -func TestCheckExistingDb(t *testing.T) { - onlyRunForEmulatorTest(t) - dbURI := fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectID, instanceID, "check-db-exists") - err := conversion.CreateDatabase(ctx, databaseAdmin, dbURI, internal.MakeConv(), os.Stdout, "", constants.BULK_MIGRATION) - if err != nil { - t.Fatal(err) - } - defer dropDatabase(t, dbURI) - testCases := []struct { - dbName string - dbExists bool - }{ - {"check-db-exists", true}, - {"check-db-does-not-exist", false}, - } - - for _, tc := range testCases { - dbExists, err := conversion.CheckExistingDb(ctx, databaseAdmin, fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectID, instanceID, tc.dbName)) - assert.Nil(t, err) - assert.Equal(t, tc.dbExists, dbExists) - } -} - func TestValidateDDL(t *testing.T) { onlyRunForEmulatorTest(t) diff --git a/webv2/helpers/helpers.go b/webv2/helpers/helpers.go index d0dbea31f9..7495d9324a 100644 --- a/webv2/helpers/helpers.go +++ b/webv2/helpers/helpers.go @@ -21,8 +21,8 @@ import ( "strings" database "cloud.google.com/go/spanner/admin/database/apiv1" + spanneracc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" - "github.com/GoogleCloudPlatform/spanner-migration-tool/conversion" adminpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1" ) @@ -160,7 +160,7 @@ func CheckOrCreateMetadataDb(projectId string, instanceId string) bool { } defer adminClient.Close() - dbExists, err := conversion.CheckExistingDb(ctx, adminClient, uri) + dbExists, err := spanneracc.CheckExistingDb(ctx, uri) if err != nil { fmt.Println(err) return false diff --git a/webv2/profile/profile.go b/webv2/profile/profile.go index 9808a594f6..e46a78ed34 100644 --- a/webv2/profile/profile.go +++ b/webv2/profile/profile.go @@ -11,6 +11,7 @@ import ( "strings" datastream "cloud.google.com/go/datastream/apiv1" + storageacc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" "github.com/GoogleCloudPlatform/spanner-migration-tool/streaming" @@ -160,7 +161,7 @@ func CreateConnectionProfile(w http.ResponseWriter, r *http.Request) { } else { bucketName = strings.ToLower(sessionState.Conv.Audit.MigrationRequestId) } - err = utils.CreateGCSBucket(bucketName, sessionState.GCPProjectID, sessionState.Region) + err = storageacc.CreateGCSBucket(ctx, bucketName, sessionState.GCPProjectID, sessionState.Region) if err != nil { http.Error(w, fmt.Sprintf("Error while creating bucket: %v", err), http.StatusBadRequest) return diff --git a/webv2/session/session_service.go b/webv2/session/session_service.go index df190eac4c..d6a26f8685 100644 --- a/webv2/session/session_service.go +++ b/webv2/session/session_service.go @@ -7,7 +7,7 @@ import ( "cloud.google.com/go/spanner" database "cloud.google.com/go/spanner/admin/database/apiv1" "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" - "github.com/GoogleCloudPlatform/spanner-migration-tool/conversion" + spanneracc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" helpers "github.com/GoogleCloudPlatform/spanner-migration-tool/webv2/helpers" ) @@ -88,7 +88,7 @@ func migrateMetadataDb(projectId, instanceId string) { defer adminClient.Close() oldMetadataDbUri := getOldMetadataDbUri(projectId, instanceId) - oldMetadataDBExists, err := conversion.CheckExistingDb(ctx, adminClient, oldMetadataDbUri) + oldMetadataDBExists, err := spanneracc.CheckExistingDb(ctx, oldMetadataDbUri) if err != nil { fmt.Printf("could not check if oldMetadataDB exists. error=%v\n", err) return diff --git a/webv2/web.go b/webv2/web.go index 7b8c84b8cc..0ae24e2f33 100644 --- a/webv2/web.go +++ b/webv2/web.go @@ -36,6 +36,7 @@ import ( "time" instance "cloud.google.com/go/spanner/admin/instance/apiv1" + storageacc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/cmd" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" @@ -2263,7 +2264,7 @@ func migrate(w http.ResponseWriter, r *http.Request) { http.Error(w, fmt.Sprintf("Can't get source and target profiles: %v", err), http.StatusBadRequest) return } - err = writeSessionFile(sessionState) + err = writeSessionFile(ctx, sessionState) if err != nil { log.Println("can't write session file") http.Error(w, fmt.Sprintf("Can't write session file to GCS: %v", err), http.StatusBadRequest) @@ -2484,9 +2485,9 @@ func createConfigFileForShardedBulkMigration(sessionState *session.SessionState, return nil } -func writeSessionFile(sessionState *session.SessionState) error { +func writeSessionFile(ctx context.Context, sessionState *session.SessionState) error { - err := utils.CreateGCSBucket(sessionState.Bucket, sessionState.GCPProjectID, sessionState.Region) + err := storageacc.CreateGCSBucket(ctx, sessionState.Bucket, sessionState.GCPProjectID, sessionState.Region) if err != nil { return fmt.Errorf("error while creating bucket: %v", err) } @@ -2495,7 +2496,7 @@ func writeSessionFile(sessionState *session.SessionState) error { if err != nil { return fmt.Errorf("can't encode session state to JSON: %v", err) } - err = utils.WriteToGCS("gs://"+sessionState.Bucket+sessionState.RootPath, "session.json", string(convJSON)) + err = storageacc.WriteDataToGCS(ctx, "gs://"+sessionState.Bucket+sessionState.RootPath, "session.json", string(convJSON)) if err != nil { return fmt.Errorf("error while writing to GCS: %v", err) } @@ -3031,7 +3032,7 @@ type ResourceDetails struct { ResourceType string `json:"ResourceType"` ResourceName string `json:"ResourceName"` ResourceUrl string `json:"ResourceUrl"` - GcloudCmd string `json:"GcloudCmd"` + GcloudCmd string `json:"GcloudCmd"` } type GeneratedResources struct { MigrationJobId string `json:"MigrationJobId"` @@ -3054,7 +3055,7 @@ type GeneratedResources struct { AggMonitoringDashboardName string `json:"AggMonitoringDashboardName"` AggMonitoringDashboardUrl string `json:"AggMonitoringDashboardUrl"` //Used for sharded migration flow - ShardToShardResourcesMap map[string][]ResourceDetails `json:"ShardToShardResourcesMap"` + ShardToShardResourcesMap map[string][]ResourceDetails `json:"ShardToShardResourcesMap"` } func addTypeToList(convertedType string, spType string, issues []internal.SchemaIssue, l []typeIssue) []typeIssue { From b452efc8174c07427d86ab99cf88388efe886629 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Tue, 9 Jan 2024 14:45:55 +0530 Subject: [PATCH 03/25] Add Unmarshall method --- accessors/dataflow/dataflow_accessor.go | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/accessors/dataflow/dataflow_accessor.go b/accessors/dataflow/dataflow_accessor.go index aa7f1bea49..277da6aa1e 100644 --- a/accessors/dataflow/dataflow_accessor.go +++ b/accessors/dataflow/dataflow_accessor.go @@ -15,12 +15,14 @@ package dataflowaccessor import ( "context" + "encoding/json" "fmt" "sort" "strings" "cloud.google.com/go/dataflow/apiv1beta3/dataflowpb" dataflowclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/dataflow" + storageacc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" "golang.org/x/exp/maps" ) @@ -176,3 +178,16 @@ func formatAdditionalUserLabels(labels map[string]string) string { } return strings.Join(res, ",") } + +func UnmarshalDataflowTuningConfig(ctx context.Context, filePath string) (DataflowTuningConfig, error) { + jsonStr, err := storageacc.ReadAnyFile(ctx, filePath) + if err != nil { + return DataflowTuningConfig{}, err + } + tuningCfg := DataflowTuningConfig{} + err = json.Unmarshal([]byte(jsonStr), &tuningCfg) + if err != nil { + return DataflowTuningConfig{}, err + } + return tuningCfg, nil +} From f7600882ebb8d403539d7ea1fbee92f5eab7a967 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Tue, 9 Jan 2024 22:45:17 +0530 Subject: [PATCH 04/25] Rename storageacc and spanner acc to storageaccessor and spanneraccessor --- accessors/dataflow/dataflow_accessor.go | 15 ---- accessors/spanner/spanner_accessor.go | 2 +- accessors/storage/storage_accessor.go | 2 +- accessors/utils/dataflow/dataflow_utils.go | 35 +++++++++ cmd/data.go | 4 +- common/constants/constants.go | 3 +- conversion/conversion.go | 72 +++++++++---------- streaming/streaming.go | 6 +- .../spanner/spanner_accessor_test.go | 6 +- webv2/helpers/helpers.go | 4 +- webv2/profile/profile.go | 4 +- webv2/session/session_service.go | 4 +- webv2/web.go | 6 +- 13 files changed, 92 insertions(+), 71 deletions(-) create mode 100644 accessors/utils/dataflow/dataflow_utils.go diff --git a/accessors/dataflow/dataflow_accessor.go b/accessors/dataflow/dataflow_accessor.go index 277da6aa1e..aa7f1bea49 100644 --- a/accessors/dataflow/dataflow_accessor.go +++ b/accessors/dataflow/dataflow_accessor.go @@ -15,14 +15,12 @@ package dataflowaccessor import ( "context" - "encoding/json" "fmt" "sort" "strings" "cloud.google.com/go/dataflow/apiv1beta3/dataflowpb" dataflowclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/dataflow" - storageacc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" "golang.org/x/exp/maps" ) @@ -178,16 +176,3 @@ func formatAdditionalUserLabels(labels map[string]string) string { } return strings.Join(res, ",") } - -func UnmarshalDataflowTuningConfig(ctx context.Context, filePath string) (DataflowTuningConfig, error) { - jsonStr, err := storageacc.ReadAnyFile(ctx, filePath) - if err != nil { - return DataflowTuningConfig{}, err - } - tuningCfg := DataflowTuningConfig{} - err = json.Unmarshal([]byte(jsonStr), &tuningCfg) - if err != nil { - return DataflowTuningConfig{}, err - } - return tuningCfg, nil -} diff --git a/accessors/spanner/spanner_accessor.go b/accessors/spanner/spanner_accessor.go index 101a9dd93e..1b1594c1e6 100644 --- a/accessors/spanner/spanner_accessor.go +++ b/accessors/spanner/spanner_accessor.go @@ -11,7 +11,7 @@ // 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 spanneracc +package spanneraccessor import ( "context" diff --git a/accessors/storage/storage_accessor.go b/accessors/storage/storage_accessor.go index 0d880b978a..f587b32627 100644 --- a/accessors/storage/storage_accessor.go +++ b/accessors/storage/storage_accessor.go @@ -11,7 +11,7 @@ // 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 storageacc +package storageaccessor import ( "context" diff --git a/accessors/utils/dataflow/dataflow_utils.go b/accessors/utils/dataflow/dataflow_utils.go new file mode 100644 index 0000000000..6b8978ef3e --- /dev/null +++ b/accessors/utils/dataflow/dataflow_utils.go @@ -0,0 +1,35 @@ +// 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 dataflowutils + +import ( + "context" + "encoding/json" + + dataflowaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/dataflow" + storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" +) + +func UnmarshalDataflowTuningConfig(ctx context.Context, filePath string) (dataflowaccessor.DataflowTuningConfig, error) { + jsonStr, err := storageaccessor.ReadAnyFile(ctx, filePath) + if err != nil { + return dataflowaccessor.DataflowTuningConfig{}, err + } + tuningCfg := dataflowaccessor.DataflowTuningConfig{} + err = json.Unmarshal([]byte(jsonStr), &tuningCfg) + if err != nil { + return dataflowaccessor.DataflowTuningConfig{}, err + } + return tuningCfg, nil +} diff --git a/cmd/data.go b/cmd/data.go index 6190ed3696..c739859bf7 100644 --- a/cmd/data.go +++ b/cmd/data.go @@ -26,7 +26,7 @@ import ( sp "cloud.google.com/go/spanner" database "cloud.google.com/go/spanner/admin/database/apiv1" - spanneracc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" + spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" "github.com/GoogleCloudPlatform/spanner-migration-tool/conversion" @@ -179,7 +179,7 @@ func (cmd *DataCmd) Execute(ctx context.Context, f *flag.FlagSet, _ ...interface // validateExistingDb validates that the existing spanner schema is in accordance with the one specified in the session file. func validateExistingDb(ctx context.Context, spDialect, dbURI string, adminClient *database.DatabaseAdminClient, client *sp.Client, conv *internal.Conv) error { - dbExists, err := spanneracc.CheckExistingDb(ctx, dbURI) + dbExists, err := spanneraccessor.CheckExistingDb(ctx, dbURI) if err != nil { err = fmt.Errorf("can't verify target database: %v", err) return err diff --git a/common/constants/constants.go b/common/constants/constants.go index e98855a81d..c0eaa87685 100644 --- a/common/constants/constants.go +++ b/common/constants/constants.go @@ -64,7 +64,8 @@ const ( MigrationMetadataKey string = "cloud-spanner-migration-metadata" // Scheme used for GCS paths - GCS_SCHEME string = "gs" + GCS_SCHEME string = "gs" + GCS_FILE_PREFIX string = "gs://" // File upload prefix for dump and session load. UPLOAD_FILE_DIR string = "upload-file" diff --git a/conversion/conversion.go b/conversion/conversion.go index 1608f98054..c6376f0215 100644 --- a/conversion/conversion.go +++ b/conversion/conversion.go @@ -44,8 +44,8 @@ import ( datastream "cloud.google.com/go/datastream/apiv1" sp "cloud.google.com/go/spanner" database "cloud.google.com/go/spanner/admin/database/apiv1" - spanneracc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" - storageacc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" + spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" + storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/metrics" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" @@ -69,12 +69,12 @@ import ( dydb "github.com/aws/aws-sdk-go/service/dynamodb" "github.com/aws/aws-sdk-go/service/dynamodbstreams" mysqldriver "github.com/go-sql-driver/mysql" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/stdlib" "go.uber.org/zap" adminpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1" "google.golang.org/grpc/metadata" "google.golang.org/protobuf/proto" - "github.com/jackc/pgx/v5" - "github.com/jackc/pgx/v5/stdlib" ) var ( @@ -357,7 +357,7 @@ func dataFromDatabase(ctx context.Context, sourceProfile profiles.SourceProfile, // Try to apply lifecycle rule to Datastream destination bucket. gcsConfig := streamingCfg.GcsCfg if gcsConfig.TtlInDaysSet { - err = storageacc.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = storageaccessor.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") @@ -478,7 +478,7 @@ func dataFromDatabaseForDataflowMigration(targetProfile profiles.TargetProfile, // Try to apply lifecycle rule to Datastream destination bucket. gcsConfig := streamingCfg.GcsCfg if gcsConfig.TtlInDaysSet { - err = storageacc.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = storageaccessor.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") @@ -800,7 +800,7 @@ func getSeekable(f *os.File) (*os.File, int64, error) { // VerifyDb checks whether the db exists and if it does, verifies if the schema is what we currently support. func VerifyDb(ctx context.Context, adminClient *database.DatabaseAdminClient, dbURI string) (dbExists bool, err error) { - dbExists, err = spanneracc.CheckExistingDb(ctx, dbURI) + dbExists, err = spanneraccessor.CheckExistingDb(ctx, dbURI) if err != nil { return dbExists, err } @@ -1282,20 +1282,20 @@ func GetInfoSchemaFromCloudSQL(sourceProfile profiles.SourceProfile, targetProfi switch driver { case constants.MYSQL: d, err := cloudsqlconn.NewDialer(context.Background(), cloudsqlconn.WithIAMAuthN()) - if err != nil { - return nil, fmt.Errorf("cloudsqlconn.NewDialer: %w", err) - } - var opts []cloudsqlconn.DialOption + if err != nil { + return nil, fmt.Errorf("cloudsqlconn.NewDialer: %w", err) + } + var opts []cloudsqlconn.DialOption instanceName := fmt.Sprintf("%s:%s:%s", sourceProfile.ConnCloudSQL.Mysql.Project, sourceProfile.ConnCloudSQL.Mysql.Region, sourceProfile.ConnCloudSQL.Mysql.InstanceName) - mysqldriver.RegisterDialContext("cloudsqlconn", - func(ctx context.Context, addr string) (net.Conn, error) { - return d.Dial(ctx, instanceName, opts...) - }) + mysqldriver.RegisterDialContext("cloudsqlconn", + func(ctx context.Context, addr string) (net.Conn, error) { + return d.Dial(ctx, instanceName, opts...) + }) - dbURI := fmt.Sprintf("%s:empty@cloudsqlconn(localhost:3306)/%s?parseTime=true", - sourceProfile.ConnCloudSQL.Mysql.User, sourceProfile.ConnCloudSQL.Mysql.Db) + dbURI := fmt.Sprintf("%s:empty@cloudsqlconn(localhost:3306)/%s?parseTime=true", + sourceProfile.ConnCloudSQL.Mysql.User, sourceProfile.ConnCloudSQL.Mysql.Db) - db, err := sql.Open("mysql", dbURI) + db, err := sql.Open("mysql", dbURI) if err != nil { return nil, fmt.Errorf("sql.Open: %w", err) } @@ -1307,25 +1307,25 @@ func GetInfoSchemaFromCloudSQL(sourceProfile profiles.SourceProfile, targetProfi }, nil case constants.POSTGRES: d, err := cloudsqlconn.NewDialer(context.Background(), cloudsqlconn.WithIAMAuthN()) - if err != nil { - return nil, fmt.Errorf("cloudsqlconn.NewDialer: %w", err) - } - var opts []cloudsqlconn.DialOption - - dsn := fmt.Sprintf("user=%s database=%s", sourceProfile.ConnCloudSQL.Pg.User, sourceProfile.ConnCloudSQL.Pg.Db) - config, err := pgx.ParseConfig(dsn) - if err != nil { - return nil, err - } + if err != nil { + return nil, fmt.Errorf("cloudsqlconn.NewDialer: %w", err) + } + var opts []cloudsqlconn.DialOption + + dsn := fmt.Sprintf("user=%s database=%s", sourceProfile.ConnCloudSQL.Pg.User, sourceProfile.ConnCloudSQL.Pg.Db) + config, err := pgx.ParseConfig(dsn) + if err != nil { + return nil, err + } instanceName := fmt.Sprintf("%s:%s:%s", sourceProfile.ConnCloudSQL.Pg.Project, sourceProfile.ConnCloudSQL.Pg.Region, sourceProfile.ConnCloudSQL.Pg.InstanceName) - config.DialFunc = func(ctx context.Context, network, instance string) (net.Conn, error) { - return d.Dial(ctx, instanceName, opts...) - } - dbURI := stdlib.RegisterConnConfig(config) - db, err := sql.Open("pgx", dbURI) - if err != nil { - return nil, fmt.Errorf("sql.Open: %w", err) - } + config.DialFunc = func(ctx context.Context, network, instance string) (net.Conn, error) { + return d.Dial(ctx, instanceName, opts...) + } + dbURI := stdlib.RegisterConnConfig(config) + db, err := sql.Open("pgx", dbURI) + if err != nil { + return nil, fmt.Errorf("sql.Open: %w", err) + } temp := false return postgres.InfoSchemaImpl{ Db: db, diff --git a/streaming/streaming.go b/streaming/streaming.go index 7a96839d4f..8fa21dfa01 100644 --- a/streaming/streaming.go +++ b/streaming/streaming.go @@ -34,7 +34,7 @@ import ( resourcemanager "cloud.google.com/go/resourcemanager/apiv3" resourcemanagerpb "cloud.google.com/go/resourcemanager/apiv3/resourcemanagerpb" dataflowaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/dataflow" - storageacc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" + storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" "github.com/GoogleCloudPlatform/spanner-migration-tool/internal" @@ -866,7 +866,7 @@ func StartDataflow(ctx context.Context, targetProfile profiles.TargetProfile, st if err != nil { return internal.DataflowOutput{}, fmt.Errorf("can't encode session state to JSON: %v", err) } - err = storageacc.WriteDataToGCS(ctx, streamingCfg.TmpDir, "session.json", string(convJSON)) + err = storageaccessor.WriteDataToGCS(ctx, streamingCfg.TmpDir, "session.json", string(convJSON)) if err != nil { return internal.DataflowOutput{}, fmt.Errorf("error while writing to GCS: %v", err) } @@ -877,7 +877,7 @@ func StartDataflow(ctx context.Context, targetProfile profiles.TargetProfile, st if err != nil { return internal.DataflowOutput{}, fmt.Errorf("failed to compute transformation context: %s", err.Error()) } - err = storageacc.WriteDataToGCS(ctx, streamingCfg.TmpDir, "transformationContext.json", string(transformationContext)) + err = storageaccessor.WriteDataToGCS(ctx, streamingCfg.TmpDir, "transformationContext.json", string(transformationContext)) if err != nil { return internal.DataflowOutput{}, fmt.Errorf("error while writing to GCS: %v", err) } diff --git a/testing/accessors/spanner/spanner_accessor_test.go b/testing/accessors/spanner/spanner_accessor_test.go index 1d7e77bd72..eb7775fcca 100644 --- a/testing/accessors/spanner/spanner_accessor_test.go +++ b/testing/accessors/spanner/spanner_accessor_test.go @@ -15,7 +15,7 @@ // TODO: Refactor this file and other integration tests by moving all common code // to remove redundancy. -package utils_test +package spanneraccessor_test import ( "context" @@ -28,7 +28,7 @@ import ( database "cloud.google.com/go/spanner/admin/database/apiv1" "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" - spanneracc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" + spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/conversion" "github.com/GoogleCloudPlatform/spanner-migration-tool/internal" @@ -116,7 +116,7 @@ func TestCheckExistingDb(t *testing.T) { } for _, tc := range testCases { - dbExists, err := spanneracc.CheckExistingDb(ctx, fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectID, instanceID, tc.dbName)) + dbExists, err := spanneraccessor.CheckExistingDb(ctx, fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectID, instanceID, tc.dbName)) assert.Nil(t, err) assert.Equal(t, tc.dbExists, dbExists) } diff --git a/webv2/helpers/helpers.go b/webv2/helpers/helpers.go index 7495d9324a..7ad584746b 100644 --- a/webv2/helpers/helpers.go +++ b/webv2/helpers/helpers.go @@ -21,7 +21,7 @@ import ( "strings" database "cloud.google.com/go/spanner/admin/database/apiv1" - spanneracc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" + spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" adminpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1" ) @@ -160,7 +160,7 @@ func CheckOrCreateMetadataDb(projectId string, instanceId string) bool { } defer adminClient.Close() - dbExists, err := spanneracc.CheckExistingDb(ctx, uri) + dbExists, err := spanneraccessor.CheckExistingDb(ctx, uri) if err != nil { fmt.Println(err) return false diff --git a/webv2/profile/profile.go b/webv2/profile/profile.go index e46a78ed34..634cd7a651 100644 --- a/webv2/profile/profile.go +++ b/webv2/profile/profile.go @@ -11,7 +11,7 @@ import ( "strings" datastream "cloud.google.com/go/datastream/apiv1" - storageacc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" + storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" "github.com/GoogleCloudPlatform/spanner-migration-tool/streaming" @@ -161,7 +161,7 @@ func CreateConnectionProfile(w http.ResponseWriter, r *http.Request) { } else { bucketName = strings.ToLower(sessionState.Conv.Audit.MigrationRequestId) } - err = storageacc.CreateGCSBucket(ctx, bucketName, sessionState.GCPProjectID, sessionState.Region) + err = storageaccessor.CreateGCSBucket(ctx, bucketName, sessionState.GCPProjectID, sessionState.Region) if err != nil { http.Error(w, fmt.Sprintf("Error while creating bucket: %v", err), http.StatusBadRequest) return diff --git a/webv2/session/session_service.go b/webv2/session/session_service.go index d6a26f8685..e579dd1ec6 100644 --- a/webv2/session/session_service.go +++ b/webv2/session/session_service.go @@ -7,7 +7,7 @@ import ( "cloud.google.com/go/spanner" database "cloud.google.com/go/spanner/admin/database/apiv1" "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" - spanneracc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" + spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" helpers "github.com/GoogleCloudPlatform/spanner-migration-tool/webv2/helpers" ) @@ -88,7 +88,7 @@ func migrateMetadataDb(projectId, instanceId string) { defer adminClient.Close() oldMetadataDbUri := getOldMetadataDbUri(projectId, instanceId) - oldMetadataDBExists, err := spanneracc.CheckExistingDb(ctx, oldMetadataDbUri) + oldMetadataDBExists, err := spanneraccessor.CheckExistingDb(ctx, oldMetadataDbUri) if err != nil { fmt.Printf("could not check if oldMetadataDB exists. error=%v\n", err) return diff --git a/webv2/web.go b/webv2/web.go index 0ae24e2f33..a0cde1399d 100644 --- a/webv2/web.go +++ b/webv2/web.go @@ -36,7 +36,7 @@ import ( "time" instance "cloud.google.com/go/spanner/admin/instance/apiv1" - storageacc "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" + storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/cmd" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" @@ -2487,7 +2487,7 @@ func createConfigFileForShardedBulkMigration(sessionState *session.SessionState, func writeSessionFile(ctx context.Context, sessionState *session.SessionState) error { - err := storageacc.CreateGCSBucket(ctx, sessionState.Bucket, sessionState.GCPProjectID, sessionState.Region) + err := storageaccessor.CreateGCSBucket(ctx, sessionState.Bucket, sessionState.GCPProjectID, sessionState.Region) if err != nil { return fmt.Errorf("error while creating bucket: %v", err) } @@ -2496,7 +2496,7 @@ func writeSessionFile(ctx context.Context, sessionState *session.SessionState) e if err != nil { return fmt.Errorf("can't encode session state to JSON: %v", err) } - err = storageacc.WriteDataToGCS(ctx, "gs://"+sessionState.Bucket+sessionState.RootPath, "session.json", string(convJSON)) + err = storageaccessor.WriteDataToGCS(ctx, "gs://"+sessionState.Bucket+sessionState.RootPath, "session.json", string(convJSON)) if err != nil { return fmt.Errorf("error while writing to GCS: %v", err) } From 5bde597a2f45fbd55b63ec97edba848d754985db Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 10 Jan 2024 20:36:05 +0530 Subject: [PATCH 05/25] Add empty test files --- .../clients/spanner/admin/admin_client_test.go | 14 ++++++++++++++ .../clients/spanner/client/spanner_client_test.go | 14 ++++++++++++++ .../instanceadmin/spanner_instance_admin_test.go | 14 ++++++++++++++ accessors/clients/storage/storage_client_test.go | 14 ++++++++++++++ accessors/spanner/spanner_accessor_test.go | 14 ++++++++++++++ accessors/storage/storage_accessor_test.go | 14 ++++++++++++++ common/utils/storage_utils.go | 6 ++++-- testing/accessors/spanner/spanner_accessor_test.go | 1 + 8 files changed, 89 insertions(+), 2 deletions(-) create mode 100644 accessors/clients/spanner/admin/admin_client_test.go create mode 100644 accessors/clients/spanner/client/spanner_client_test.go create mode 100644 accessors/clients/spanner/instanceadmin/spanner_instance_admin_test.go create mode 100644 accessors/clients/storage/storage_client_test.go create mode 100644 accessors/spanner/spanner_accessor_test.go create mode 100644 accessors/storage/storage_accessor_test.go diff --git a/accessors/clients/spanner/admin/admin_client_test.go b/accessors/clients/spanner/admin/admin_client_test.go new file mode 100644 index 0000000000..eafb726e1a --- /dev/null +++ b/accessors/clients/spanner/admin/admin_client_test.go @@ -0,0 +1,14 @@ +// 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 diff --git a/accessors/clients/spanner/client/spanner_client_test.go b/accessors/clients/spanner/client/spanner_client_test.go new file mode 100644 index 0000000000..b675039f44 --- /dev/null +++ b/accessors/clients/spanner/client/spanner_client_test.go @@ -0,0 +1,14 @@ +// 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 spannerclient diff --git a/accessors/clients/spanner/instanceadmin/spanner_instance_admin_test.go b/accessors/clients/spanner/instanceadmin/spanner_instance_admin_test.go new file mode 100644 index 0000000000..5bee8f5d97 --- /dev/null +++ b/accessors/clients/spanner/instanceadmin/spanner_instance_admin_test.go @@ -0,0 +1,14 @@ +// 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 spinstanceadmin diff --git a/accessors/clients/storage/storage_client_test.go b/accessors/clients/storage/storage_client_test.go new file mode 100644 index 0000000000..02bff2f3ac --- /dev/null +++ b/accessors/clients/storage/storage_client_test.go @@ -0,0 +1,14 @@ +// 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 storageclient diff --git a/accessors/spanner/spanner_accessor_test.go b/accessors/spanner/spanner_accessor_test.go new file mode 100644 index 0000000000..5f0d2f60b9 --- /dev/null +++ b/accessors/spanner/spanner_accessor_test.go @@ -0,0 +1,14 @@ +// 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 spanneraccessor diff --git a/accessors/storage/storage_accessor_test.go b/accessors/storage/storage_accessor_test.go new file mode 100644 index 0000000000..1519db317b --- /dev/null +++ b/accessors/storage/storage_accessor_test.go @@ -0,0 +1,14 @@ +// 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 storageaccessor diff --git a/common/utils/storage_utils.go b/common/utils/storage_utils.go index 3429c6d293..3968b1731e 100644 --- a/common/utils/storage_utils.go +++ b/common/utils/storage_utils.go @@ -12,8 +12,10 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Package utils contains common helper functions used across multiple other packages. -// Utils should not import any Spanner migration tool packages. +/* +Package utils contains common helper functions used across multiple other packages. +Utils should not import any Spanner migration tool packages. +*/ package utils import ( diff --git a/testing/accessors/spanner/spanner_accessor_test.go b/testing/accessors/spanner/spanner_accessor_test.go index eb7775fcca..a4c440a0c3 100644 --- a/testing/accessors/spanner/spanner_accessor_test.go +++ b/testing/accessors/spanner/spanner_accessor_test.go @@ -45,6 +45,7 @@ var ( databaseAdmin *database.DatabaseAdminClient ) +// This test should move as a mock unit test inside accessors itself. func TestMain(m *testing.M) { cleanup := initTests() res := m.Run() From 66b67bc30252931438cd6e291e65d58776ae52ae Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Tue, 23 Jan 2024 13:08:54 +0530 Subject: [PATCH 06/25] Increade version retention period Add log statements to storage accessor functions --- accessors/spanner/spanner_accessor.go | 15 +++++---------- accessors/storage/storage_accessor.go | 13 ++++++++++--- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/accessors/spanner/spanner_accessor.go b/accessors/spanner/spanner_accessor.go index 1b1594c1e6..aa03a3884e 100644 --- a/accessors/spanner/spanner_accessor.go +++ b/accessors/spanner/spanner_accessor.go @@ -50,12 +50,8 @@ func GetDatabaseDialect(ctx context.Context, dbURI string) (string, error) { func CheckExistingDb(ctx context.Context, dbURI string) (bool, error) { gotResponse := make(chan bool) var err error - adminClient, err := spanneradmin.GetOrCreateClient(ctx) - if err != nil { - return false, err - } go func() { - _, err = adminClient.GetDatabase(ctx, &databasepb.GetDatabaseRequest{Name: dbURI}) + _, err = GetDatabase(ctx, dbURI) gotResponse <- true }() for { @@ -122,12 +118,11 @@ func CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI stri return false, err } stmt := spanner.Statement{ - SQL: `SELECT * FROM information_schema.change_streams`, + SQL: `SELECT CHANGE_STREAM_NAME FROM information_schema.change_streams`, } iter := spClient.Single().Query(ctx, stmt) defer iter.Stop() - var cs_catalog, cs_schema, cs_name string - var coversAll bool + var cs_name string csExists := false for { row, err := iter.Next() @@ -137,7 +132,7 @@ func CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI stri if err != nil { return false, fmt.Errorf("couldn't read row from change_streams table: %w", err) } - err = row.Columns(&cs_catalog, &cs_schema, &cs_name, &coversAll) + err = row.Columns(&cs_name) if err != nil { return false, fmt.Errorf("can't scan row from change_streams table: %v", err) } @@ -189,7 +184,7 @@ func CreateChangeStream(ctx context.Context, changeStreamName, dbURI string) err op, err := spClient.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{ Database: dbURI, // TODO: create change stream for only the tables present in Spanner. - Statements: []string{fmt.Sprintf("CREATE CHANGE STREAM %s FOR ALL OPTIONS (value_capture_type = 'NEW_ROW')", changeStreamName)}, + Statements: []string{fmt.Sprintf("CREATE CHANGE STREAM %s FOR ALL OPTIONS (value_capture_type = 'NEW_ROW', retention_period = '7d')", changeStreamName)}, }) if err != nil { return fmt.Errorf("cannot submit request create change stream request: %v", err) diff --git a/accessors/storage/storage_accessor.go b/accessors/storage/storage_accessor.go index f587b32627..25c0828906 100644 --- a/accessors/storage/storage_accessor.go +++ b/accessors/storage/storage_accessor.go @@ -75,7 +75,7 @@ func createGCSBucketUtil(ctx context.Context, bucketName, projectID, location st } } else { - fmt.Printf("Created new GCS bucket: %v\n", bucketName) + logger.Log.Info(fmt.Sprintf("Created new GCS bucket: %v\n", bucketName)) } return nil } @@ -143,10 +143,14 @@ func WriteDataToGCS(ctx context.Context, filePath, fileName, data string) error obj := bucket.Object(u.Path[1:] + fileName) w := obj.NewWriter(ctx) - if _, err := fmt.Fprint(w, data); err != nil { + logger.Log.Info(fmt.Sprintf("Writing data to %s", filePath)) + n, err := fmt.Fprint(w, data) + if err != nil { fmt.Printf("Failed to write to Cloud Storage: %s", filePath) return err } + logger.Log.Info(fmt.Sprintf("Wrote %d bytes to GCS", n)) + if err := w.Close(); err != nil { fmt.Printf("Failed to close GCS file: %s", filePath) return err @@ -174,9 +178,12 @@ func ReadGcsFile(ctx context.Context, filePath string) (string, error) { } defer rc.Close() buf := new(strings.Builder) - if _, err := io.Copy(buf, rc); err != nil { + logger.Log.Info(fmt.Sprintf("Reading from %s", filePath)) + n, err := io.Copy(buf, rc) + if err != nil { return "", err } + logger.Log.Info(fmt.Sprintf("Read %d bytes", n)) return buf.String(), nil } From 418094240c7568285580d4a71baaf7a7b02d1fb3 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 24 Jan 2024 13:40:27 +0530 Subject: [PATCH 07/25] Add storage accessor interface and impl --- accessors/storage/storage_accessor.go | 36 ++++++++++++++++++--------- conversion/conversion.go | 6 +++-- streaming/streaming.go | 6 ++--- webv2/profile/profile.go | 3 ++- webv2/web.go | 6 ++--- 5 files changed, 36 insertions(+), 21 deletions(-) diff --git a/accessors/storage/storage_accessor.go b/accessors/storage/storage_accessor.go index 25c0828906..59d29e6f9c 100644 --- a/accessors/storage/storage_accessor.go +++ b/accessors/storage/storage_accessor.go @@ -28,15 +28,27 @@ import ( "google.golang.org/api/googleapi" ) -func CreateGCSBucket(ctx context.Context, bucketName, projectID, location string) error { - return createGCSBucketUtil(ctx, bucketName, projectID, location, nil, 0) +type StorageAccessor interface { + CreateGCSBucket(ctx context.Context, bucketName, projectID, location string) error + CreateGCSBucketWithLifecycle(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error + EnableBucketLifecycleDeleteRule(ctx context.Context, bucketName string, matchesPrefix []string, ttl int64) error + UploadLocalFileToGCS(ctx context.Context, filePath, fileName, localFilePath string) error + WriteDataToGCS(ctx context.Context, filePath, fileName, data string) error + ReadGcsFile(ctx context.Context, filePath string) (string, error) + ReadAnyFile(ctx context.Context, filePath string) (string, error) } -func CreateGCSBucketWithLifecycle(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { - return createGCSBucketUtil(ctx, bucketName, projectID, location, matchesPrefix, ttl) +type StorageAccessorImpl struct{} + +func (sa StorageAccessorImpl) CreateGCSBucket(ctx context.Context, bucketName, projectID, location string) error { + return sa.createGCSBucketUtil(ctx, bucketName, projectID, location, nil, 0) +} + +func (sa StorageAccessorImpl) CreateGCSBucketWithLifecycle(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { + return sa.createGCSBucketUtil(ctx, bucketName, projectID, location, matchesPrefix, ttl) } -func createGCSBucketUtil(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { +func (sa StorageAccessorImpl) createGCSBucketUtil(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { client, err := storageclient.GetOrCreateClient(ctx) if err != nil { return err @@ -83,7 +95,7 @@ func createGCSBucketUtil(ctx context.Context, bucketName, projectID, location st // Applies the bucket lifecycle with delete rule. Only accepts the Age and // prefix rule conditions as it is only used for the Datastream destination // bucket currently. -func EnableBucketLifecycleDeleteRule(ctx context.Context, bucketName string, matchesPrefix []string, ttl int64) error { +func (sa StorageAccessorImpl) EnableBucketLifecycleDeleteRule(ctx context.Context, bucketName string, matchesPrefix []string, ttl int64) error { client, err := storageclient.GetOrCreateClient(ctx) if err != nil { return fmt.Errorf("could not create client while enabling lifecycle: %w", err) @@ -120,15 +132,15 @@ func EnableBucketLifecycleDeleteRule(ctx context.Context, bucketName string, mat } // UploadLocalFileToGCS uploads an object. -func UploadLocalFileToGCS(ctx context.Context, filePath, fileName, localFilePath string) error { +func (sa StorageAccessorImpl) UploadLocalFileToGCS(ctx context.Context, filePath, fileName, localFilePath string) error { data, err := os.ReadFile(localFilePath) if err != nil { return fmt.Errorf("could not read file %s: %w", localFilePath, err) } - return WriteDataToGCS(ctx, filePath, fileName, string(data)) + return sa.WriteDataToGCS(ctx, filePath, fileName, string(data)) } -func WriteDataToGCS(ctx context.Context, filePath, fileName, data string) error { +func (sa StorageAccessorImpl) WriteDataToGCS(ctx context.Context, filePath, fileName, data string) error { client, err := storageclient.GetOrCreateClient(ctx) if err != nil { return fmt.Errorf("could not create client while uploading to GCS: %w", err) @@ -158,7 +170,7 @@ func WriteDataToGCS(ctx context.Context, filePath, fileName, data string) error return nil } -func ReadGcsFile(ctx context.Context, filePath string) (string, error) { +func (sa StorageAccessorImpl) ReadGcsFile(ctx context.Context, filePath string) (string, error) { client, err := storageclient.GetOrCreateClient(ctx) if err != nil { return "", fmt.Errorf("could not create client: %w", err) @@ -187,9 +199,9 @@ func ReadGcsFile(ctx context.Context, filePath string) (string, error) { return buf.String(), nil } -func ReadAnyFile(ctx context.Context, filePath string) (string, error) { +func (sa StorageAccessorImpl) ReadAnyFile(ctx context.Context, filePath string) (string, error) { if strings.HasPrefix(filePath, constants.GCS_FILE_PREFIX) { - return ReadGcsFile(ctx, filePath) + return sa.ReadGcsFile(ctx, filePath) } buf, err := os.ReadFile(filePath) if err != nil { diff --git a/conversion/conversion.go b/conversion/conversion.go index c6376f0215..23610c0653 100644 --- a/conversion/conversion.go +++ b/conversion/conversion.go @@ -356,8 +356,9 @@ func dataFromDatabase(ctx context.Context, sourceProfile profiles.SourceProfile, // Try to apply lifecycle rule to Datastream destination bucket. gcsConfig := streamingCfg.GcsCfg + sa := storageaccessor.StorageAccessorImpl{} if gcsConfig.TtlInDaysSet { - err = storageaccessor.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = sa.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") @@ -477,8 +478,9 @@ func dataFromDatabaseForDataflowMigration(targetProfile profiles.TargetProfile, // Try to apply lifecycle rule to Datastream destination bucket. gcsConfig := streamingCfg.GcsCfg + sa := storageaccessor.StorageAccessorImpl{} if gcsConfig.TtlInDaysSet { - err = storageaccessor.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = sa.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") diff --git a/streaming/streaming.go b/streaming/streaming.go index 8fa21dfa01..f00af44198 100644 --- a/streaming/streaming.go +++ b/streaming/streaming.go @@ -861,12 +861,12 @@ func StartDatastream(ctx context.Context, streamingCfg StreamingCfg, sourceProfi } func StartDataflow(ctx context.Context, targetProfile profiles.TargetProfile, streamingCfg StreamingCfg, conv *internal.Conv) (internal.DataflowOutput, error) { - + sa := storageaccessor.StorageAccessorImpl{} convJSON, err := json.MarshalIndent(conv, "", " ") if err != nil { return internal.DataflowOutput{}, fmt.Errorf("can't encode session state to JSON: %v", err) } - err = storageaccessor.WriteDataToGCS(ctx, streamingCfg.TmpDir, "session.json", string(convJSON)) + err = sa.WriteDataToGCS(ctx, streamingCfg.TmpDir, "session.json", string(convJSON)) if err != nil { return internal.DataflowOutput{}, fmt.Errorf("error while writing to GCS: %v", err) } @@ -877,7 +877,7 @@ func StartDataflow(ctx context.Context, targetProfile profiles.TargetProfile, st if err != nil { return internal.DataflowOutput{}, fmt.Errorf("failed to compute transformation context: %s", err.Error()) } - err = storageaccessor.WriteDataToGCS(ctx, streamingCfg.TmpDir, "transformationContext.json", string(transformationContext)) + err = sa.WriteDataToGCS(ctx, streamingCfg.TmpDir, "transformationContext.json", string(transformationContext)) if err != nil { return internal.DataflowOutput{}, fmt.Errorf("error while writing to GCS: %v", err) } diff --git a/webv2/profile/profile.go b/webv2/profile/profile.go index 634cd7a651..923bb5d84d 100644 --- a/webv2/profile/profile.go +++ b/webv2/profile/profile.go @@ -154,6 +154,7 @@ func CreateConnectionProfile(w http.ResponseWriter, r *http.Request) { ValidateOnly: details.ValidateOnly, } var bucketName string + sa := storageaccessor.StorageAccessorImpl{} if !details.IsSource { if sessionState.IsSharded { @@ -161,7 +162,7 @@ func CreateConnectionProfile(w http.ResponseWriter, r *http.Request) { } else { bucketName = strings.ToLower(sessionState.Conv.Audit.MigrationRequestId) } - err = storageaccessor.CreateGCSBucket(ctx, bucketName, sessionState.GCPProjectID, sessionState.Region) + err = sa.CreateGCSBucket(ctx, bucketName, sessionState.GCPProjectID, sessionState.Region) if err != nil { http.Error(w, fmt.Sprintf("Error while creating bucket: %v", err), http.StatusBadRequest) return diff --git a/webv2/web.go b/webv2/web.go index a0cde1399d..c45dab3514 100644 --- a/webv2/web.go +++ b/webv2/web.go @@ -2486,8 +2486,8 @@ func createConfigFileForShardedBulkMigration(sessionState *session.SessionState, } func writeSessionFile(ctx context.Context, sessionState *session.SessionState) error { - - err := storageaccessor.CreateGCSBucket(ctx, sessionState.Bucket, sessionState.GCPProjectID, sessionState.Region) + sa := storageaccessor.StorageAccessorImpl{} + err := sa.CreateGCSBucket(ctx, sessionState.Bucket, sessionState.GCPProjectID, sessionState.Region) if err != nil { return fmt.Errorf("error while creating bucket: %v", err) } @@ -2496,7 +2496,7 @@ func writeSessionFile(ctx context.Context, sessionState *session.SessionState) e if err != nil { return fmt.Errorf("can't encode session state to JSON: %v", err) } - err = storageaccessor.WriteDataToGCS(ctx, "gs://"+sessionState.Bucket+sessionState.RootPath, "session.json", string(convJSON)) + err = sa.WriteDataToGCS(ctx, "gs://"+sessionState.Bucket+sessionState.RootPath, "session.json", string(convJSON)) if err != nil { return fmt.Errorf("error while writing to GCS: %v", err) } From 9fedfa31c6384a0b2f77fa21cb8eb2ba6c0aacc9 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 24 Jan 2024 14:00:44 +0530 Subject: [PATCH 08/25] Add storage client unit tests --- accessors/clients/storage/storage_client.go | 6 +- .../clients/storage/storage_client_test.go | 103 ++++++++++++++++++ 2 files changed, 108 insertions(+), 1 deletion(-) diff --git a/accessors/clients/storage/storage_client.go b/accessors/clients/storage/storage_client.go index 569841d299..28c62b868f 100644 --- a/accessors/clients/storage/storage_client.go +++ b/accessors/clients/storage/storage_client.go @@ -24,11 +24,15 @@ import ( var once sync.Once var gcsClient *storage.Client +// This function is declared as a global variable to make it testable. The unit +// tests edit this function, acting like a double. +var newClient = storage.NewClient + func GetOrCreateClient(ctx context.Context) (*storage.Client, error) { var err error if gcsClient == nil { once.Do(func() { - gcsClient, err = storage.NewClient(ctx) + gcsClient, err = newClient(ctx) }) if err != nil { return nil, fmt.Errorf("failed to create storage client: %v", err) diff --git a/accessors/clients/storage/storage_client_test.go b/accessors/clients/storage/storage_client_test.go index 02bff2f3ac..73bd6b873a 100644 --- a/accessors/clients/storage/storage_client_test.go +++ b/accessors/clients/storage/storage_client_test.go @@ -12,3 +12,106 @@ // See the License for the specific language governing permissions and // limitations under the License. package storageclient + +import ( + "context" + "fmt" + "os" + "sync" + "testing" + + "cloud.google.com/go/storage" + "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() { + gcsClient = nil + once = sync.Once{} +} + +func TestGetOrCreateClient_Basic(t *testing.T) { + resetTest() + ctx := context.Background() + oldFunc := newClient + defer func() { newClient = oldFunc }() + newClient = func(ctx context.Context, opts ...option.ClientOption) (*storage.Client, error) { + return &storage.Client{}, nil + } + c, err := GetOrCreateClient(ctx) + assert.NotNil(t, c) + assert.Nil(t, err) +} + +func TestGetOrCreateClient_OnlyOnceViaSync(t *testing.T) { + resetTest() + ctx := context.Background() + oldFunc := newClient + defer func() { newClient = oldFunc }() + + newClient = func(ctx context.Context, opts ...option.ClientOption) (*storage.Client, error) { + return &storage.Client{}, 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. + gcsClient = nil + + newClient = func(ctx context.Context, opts ...option.ClientOption) (*storage.Client, 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 := newClient + defer func() { newClient = oldFunc }() + + newClient = func(ctx context.Context, opts ...option.ClientOption) (*storage.Client, error) { + return &storage.Client{}, 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{} + newClient = func(ctx context.Context, opts ...option.ClientOption) (*storage.Client, 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 := newClient + defer func() { newClient = oldFunc }() + + newClient = func(ctx context.Context, opts ...option.ClientOption) (*storage.Client, error) { + return nil, fmt.Errorf("test error") + } + c, err := GetOrCreateClient(ctx) + assert.Nil(t, c) + assert.NotNil(t, err) +} From 31aad0d9204c4c2fd5467b1da0aafa1bf80e7870 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 24 Jan 2024 14:10:11 +0530 Subject: [PATCH 09/25] Add spanner admin client unit tests --- .../clients/spanner/admin/admin_client.go | 6 +- .../spanner/admin/admin_client_test.go | 102 ++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) diff --git a/accessors/clients/spanner/admin/admin_client.go b/accessors/clients/spanner/admin/admin_client.go index c0ba8aed5f..220621ee09 100644 --- a/accessors/clients/spanner/admin/admin_client.go +++ b/accessors/clients/spanner/admin/admin_client.go @@ -24,11 +24,15 @@ import ( var once sync.Once var spannerAdminClient *database.DatabaseAdminClient +// This function is declared as a global variable to make it testable. The unit +// tests edit 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 = database.NewDatabaseAdminClient(ctx) + spannerAdminClient, err = newDatabaseAdminClient(ctx) }) if err != nil { return nil, fmt.Errorf("failed to create spanner admin client: %v", err) diff --git a/accessors/clients/spanner/admin/admin_client_test.go b/accessors/clients/spanner/admin/admin_client_test.go index eafb726e1a..7c911ee096 100644 --- a/accessors/clients/spanner/admin/admin_client_test.go +++ b/accessors/clients/spanner/admin/admin_client_test.go @@ -12,3 +12,105 @@ // 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) +} From 8dff86315cf24bfc902b3fcdb5899f134bd93701 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 24 Jan 2024 14:13:18 +0530 Subject: [PATCH 10/25] Add spanner instance admin client unit tests --- .../instanceadmin/spanner_instance_admin.go | 6 +- .../spanner_instance_admin_test.go | 102 ++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) diff --git a/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go b/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go index 324da4ac32..8e6529bfc7 100644 --- a/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go +++ b/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go @@ -24,11 +24,15 @@ import ( var once sync.Once var instanceAdminClient *instance.InstanceAdminClient +// This function is declared as a global variable to make it testable. The unit +// tests edit this function, acting like a double. +var newInstanceAdminClient = instance.NewInstanceAdminClient + func GetOrCreateClient(ctx context.Context) (*instance.InstanceAdminClient, error) { var err error if instanceAdminClient == nil { once.Do(func() { - instanceAdminClient, err = instance.NewInstanceAdminClient(ctx) + instanceAdminClient, err = newInstanceAdminClient(ctx) }) if err != nil { return nil, fmt.Errorf("failed to create spanner instance admin client: %v", err) diff --git a/accessors/clients/spanner/instanceadmin/spanner_instance_admin_test.go b/accessors/clients/spanner/instanceadmin/spanner_instance_admin_test.go index 5bee8f5d97..2792d03e82 100644 --- a/accessors/clients/spanner/instanceadmin/spanner_instance_admin_test.go +++ b/accessors/clients/spanner/instanceadmin/spanner_instance_admin_test.go @@ -12,3 +12,105 @@ // See the License for the specific language governing permissions and // limitations under the License. package spinstanceadmin + +import ( + "context" + "fmt" + "os" + "sync" + "testing" + + instance "cloud.google.com/go/spanner/admin/instance/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() { + instanceAdminClient = nil + once = sync.Once{} +} + +func TestGetOrCreateClient_Basic(t *testing.T) { + resetTest() + ctx := context.Background() + oldFunc := newInstanceAdminClient + defer func() { newInstanceAdminClient = oldFunc }() + newInstanceAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*instance.InstanceAdminClient, error) { + return &instance.InstanceAdminClient{}, nil + } + c, err := GetOrCreateClient(ctx) + assert.NotNil(t, c) + assert.Nil(t, err) +} + +func TestGetOrCreateClient_OnlyOnceViaSync(t *testing.T) { + resetTest() + ctx := context.Background() + oldFunc := newInstanceAdminClient + defer func() { newInstanceAdminClient = oldFunc }() + + newInstanceAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*instance.InstanceAdminClient, error) { + return &instance.InstanceAdminClient{}, 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. + instanceAdminClient = nil + newInstanceAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*instance.InstanceAdminClient, 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 := newInstanceAdminClient + defer func() { newInstanceAdminClient = oldFunc }() + + newInstanceAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*instance.InstanceAdminClient, error) { + return &instance.InstanceAdminClient{}, 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{} + newInstanceAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*instance.InstanceAdminClient, 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 := newInstanceAdminClient + defer func() { newInstanceAdminClient = oldFunc }() + + newInstanceAdminClient = func(ctx context.Context, opts ...option.ClientOption) (*instance.InstanceAdminClient, error) { + return nil, fmt.Errorf("test error") + } + c, err := GetOrCreateClient(ctx) + assert.Nil(t, c) + assert.NotNil(t, err) +} From 372e03d8472947fa4cc000756d2ff6c310d79c78 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 24 Jan 2024 14:19:58 +0530 Subject: [PATCH 11/25] Add spanner client unit tests --- .../clients/spanner/client/spanner_client.go | 8 +- .../spanner/client/spanner_client_test.go | 102 ++++++++++++++++++ 2 files changed, 108 insertions(+), 2 deletions(-) diff --git a/accessors/clients/spanner/client/spanner_client.go b/accessors/clients/spanner/client/spanner_client.go index bbb3e252c1..d6edd81ece 100644 --- a/accessors/clients/spanner/client/spanner_client.go +++ b/accessors/clients/spanner/client/spanner_client.go @@ -24,11 +24,15 @@ import ( var once sync.Once var spannerClient *sp.Client +// This function is declared as a global variable to make it testable. The unit +// tests edit this function, acting like a double. +var newClient = sp.NewClient + func GetOrCreateClient(ctx context.Context, dbURI string) (*sp.Client, error) { var err error - if spannerClient == nil || spannerClient.DatabaseName() != dbURI { + if spannerClient == nil { once.Do(func() { - spannerClient, err = sp.NewClient(ctx, dbURI) + spannerClient, err = newClient(ctx, dbURI) }) if err != nil { return nil, fmt.Errorf("failed to create spanner database client: %v", err) diff --git a/accessors/clients/spanner/client/spanner_client_test.go b/accessors/clients/spanner/client/spanner_client_test.go index b675039f44..66f5059591 100644 --- a/accessors/clients/spanner/client/spanner_client_test.go +++ b/accessors/clients/spanner/client/spanner_client_test.go @@ -12,3 +12,105 @@ // See the License for the specific language governing permissions and // limitations under the License. package spannerclient + +import ( + "context" + "fmt" + "os" + "sync" + "testing" + + sp "cloud.google.com/go/spanner" + "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() { + spannerClient = nil + once = sync.Once{} +} + +func TestGetOrCreateClient_Basic(t *testing.T) { + resetTest() + ctx := context.Background() + oldFunc := newClient + defer func() { newClient = oldFunc }() + newClient = func(ctx context.Context, database string, opts ...option.ClientOption) (*sp.Client, error) { + return &sp.Client{}, nil + } + c, err := GetOrCreateClient(ctx, "testURI") + assert.NotNil(t, c) + assert.Nil(t, err) +} + +func TestGetOrCreateClient_OnlyOnceViaSync(t *testing.T) { + resetTest() + ctx := context.Background() + oldFunc := newClient + defer func() { newClient = oldFunc }() + + newClient = func(ctx context.Context, database string, opts ...option.ClientOption) (*sp.Client, error) { + return &sp.Client{}, nil + } + c, err := GetOrCreateClient(ctx, "testURI") + 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. + spannerClient = nil + newClient = func(ctx context.Context, database string, opts ...option.ClientOption) (*sp.Client, error) { + return nil, fmt.Errorf("test error") + } + c, err = GetOrCreateClient(ctx, "testURI") + assert.Nil(t, c) + assert.Nil(t, err) +} + +func TestGetOrCreateClient_OnlyOnceViaIf(t *testing.T) { + resetTest() + ctx := context.Background() + oldFunc := newClient + defer func() { newClient = oldFunc }() + + newClient = func(ctx context.Context, database string, opts ...option.ClientOption) (*sp.Client, error) { + return &sp.Client{}, nil + } + oldC, err := GetOrCreateClient(ctx, "testURI") + 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{} + newClient = func(ctx context.Context, database string, opts ...option.ClientOption) (*sp.Client, error) { + return nil, fmt.Errorf("test error") + } + newC, err := GetOrCreateClient(ctx, "testURI") + assert.Equal(t, oldC, newC) + assert.Nil(t, err) +} + +func TestGetOrCreateClient_Error(t *testing.T) { + resetTest() + ctx := context.Background() + oldFunc := newClient + defer func() { newClient = oldFunc }() + + newClient = func(ctx context.Context, database string, opts ...option.ClientOption) (*sp.Client, error) { + return nil, fmt.Errorf("test error") + } + c, err := GetOrCreateClient(ctx, "testURI") + assert.Nil(t, c) + assert.NotNil(t, err) +} From 00cea899d70fd10e65ff4417e34a00e956bc0966 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 24 Jan 2024 16:02:25 +0530 Subject: [PATCH 12/25] Add interface and implementor for Spanner Accessor --- accessors/spanner/spanner_accessor.go | 33 +++++++++++++------ cmd/data.go | 3 +- conversion/conversion.go | 3 +- .../spanner/spanner_accessor_test.go | 4 +-- webv2/helpers/helpers.go | 3 +- webv2/session/session_service.go | 3 +- 6 files changed, 33 insertions(+), 16 deletions(-) diff --git a/accessors/spanner/spanner_accessor.go b/accessors/spanner/spanner_accessor.go index aa03a3884e..93aea4555f 100644 --- a/accessors/spanner/spanner_accessor.go +++ b/accessors/spanner/spanner_accessor.go @@ -29,7 +29,20 @@ import ( "google.golang.org/api/iterator" ) -func GetDatabase(ctx context.Context, dbURI string) (*databasepb.Database, error) { +type SpannerAccessor interface { + GetDatabase(ctx context.Context, dbURI string) (*databasepb.Database, error) + GetDatabaseDialect(ctx context.Context, dbURI string) (string, error) + CheckExistingDb(ctx context.Context, dbURI string) (bool, error) + CreateEmptyDatabase(ctx context.Context, dbURI string) error + GetSpannerLeaderLocation(ctx context.Context, instanceURI string) (string, error) + CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI string) (bool, error) + ValidateChangeStreamOptions(ctx context.Context, changeStreamName, dbURI string) error + CreateChangeStream(ctx context.Context, changeStreamName, dbURI string) error +} + +type SpannerAccessorImpl struct{} + +func (sp SpannerAccessorImpl) GetDatabase(ctx context.Context, dbURI string) (*databasepb.Database, error) { adminClient, err := spanneradmin.GetOrCreateClient(ctx) if err != nil { return nil, err @@ -37,8 +50,8 @@ func GetDatabase(ctx context.Context, dbURI string) (*databasepb.Database, error return adminClient.GetDatabase(ctx, &databasepb.GetDatabaseRequest{Name: dbURI}) } -func GetDatabaseDialect(ctx context.Context, dbURI string) (string, error) { - result, err := GetDatabase(ctx, dbURI) +func (sp SpannerAccessorImpl) GetDatabaseDialect(ctx context.Context, dbURI string) (string, error) { + result, err := sp.GetDatabase(ctx, dbURI) if err != nil { return "", fmt.Errorf("cannot connect to database: %v", err) } @@ -47,11 +60,11 @@ func GetDatabaseDialect(ctx context.Context, dbURI string) (string, error) { // CheckExistingDb checks whether the database with dbURI exists or not. // If API call doesn't respond then user is informed after every 5 minutes on command line. -func CheckExistingDb(ctx context.Context, dbURI string) (bool, error) { +func (sp SpannerAccessorImpl) CheckExistingDb(ctx context.Context, dbURI string) (bool, error) { gotResponse := make(chan bool) var err error go func() { - _, err = GetDatabase(ctx, dbURI) + _, err = sp.GetDatabase(ctx, dbURI) gotResponse <- true }() for { @@ -70,7 +83,7 @@ func CheckExistingDb(ctx context.Context, dbURI string) (bool, error) { } } -func CreateEmptyDatabase(ctx context.Context, dbURI string) error { +func (sp SpannerAccessorImpl) CreateEmptyDatabase(ctx context.Context, dbURI string) error { adminClient, err := spanneradmin.GetOrCreateClient(ctx) if err != nil { return err @@ -90,7 +103,7 @@ func CreateEmptyDatabase(ctx context.Context, dbURI string) error { return nil } -func GetSpannerLeaderLocation(ctx context.Context, instanceURI string) (string, error) { +func (sp SpannerAccessorImpl) GetSpannerLeaderLocation(ctx context.Context, instanceURI string) (string, error) { instanceClient, err := spinstanceadmin.GetOrCreateClient(ctx) if err != nil { return "", err @@ -112,7 +125,7 @@ func GetSpannerLeaderLocation(ctx context.Context, instanceURI string) (string, return "", fmt.Errorf("no leader found for spanner instance %s while trying fetch location", instanceURI) } -func CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI string) (bool, error) { +func (sp SpannerAccessorImpl) CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI string) (bool, error) { spClient, err := spannerclient.GetOrCreateClient(ctx, dbURI) if err != nil { return false, err @@ -144,7 +157,7 @@ func CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI stri return csExists, nil } -func ValidateChangeStreamOptions(ctx context.Context, changeStreamName, dbURI string) error { +func (sp SpannerAccessorImpl) ValidateChangeStreamOptions(ctx context.Context, changeStreamName, dbURI string) error { spClient, err := spannerclient.GetOrCreateClient(ctx, dbURI) if err != nil { return err @@ -179,7 +192,7 @@ func ValidateChangeStreamOptions(ctx context.Context, changeStreamName, dbURI st return nil } -func CreateChangeStream(ctx context.Context, changeStreamName, dbURI string) error { +func (sp SpannerAccessorImpl) CreateChangeStream(ctx context.Context, changeStreamName, dbURI string) error { spClient, _ := spanneradmin.GetOrCreateClient(ctx) op, err := spClient.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{ Database: dbURI, diff --git a/cmd/data.go b/cmd/data.go index c739859bf7..b31f041269 100644 --- a/cmd/data.go +++ b/cmd/data.go @@ -179,7 +179,8 @@ func (cmd *DataCmd) Execute(ctx context.Context, f *flag.FlagSet, _ ...interface // validateExistingDb validates that the existing spanner schema is in accordance with the one specified in the session file. func validateExistingDb(ctx context.Context, spDialect, dbURI string, adminClient *database.DatabaseAdminClient, client *sp.Client, conv *internal.Conv) error { - dbExists, err := spanneraccessor.CheckExistingDb(ctx, dbURI) + spA := spanneraccessor.SpannerAccessorImpl{} + dbExists, err := spA.CheckExistingDb(ctx, dbURI) if err != nil { err = fmt.Errorf("can't verify target database: %v", err) return err diff --git a/conversion/conversion.go b/conversion/conversion.go index 23610c0653..bf02976781 100644 --- a/conversion/conversion.go +++ b/conversion/conversion.go @@ -802,7 +802,8 @@ func getSeekable(f *os.File) (*os.File, int64, error) { // VerifyDb checks whether the db exists and if it does, verifies if the schema is what we currently support. func VerifyDb(ctx context.Context, adminClient *database.DatabaseAdminClient, dbURI string) (dbExists bool, err error) { - dbExists, err = spanneraccessor.CheckExistingDb(ctx, dbURI) + spA := spanneraccessor.SpannerAccessorImpl{} + dbExists, err = spA.CheckExistingDb(ctx, dbURI) if err != nil { return dbExists, err } diff --git a/testing/accessors/spanner/spanner_accessor_test.go b/testing/accessors/spanner/spanner_accessor_test.go index a4c440a0c3..0dd0baa07f 100644 --- a/testing/accessors/spanner/spanner_accessor_test.go +++ b/testing/accessors/spanner/spanner_accessor_test.go @@ -115,9 +115,9 @@ func TestCheckExistingDb(t *testing.T) { {"check-db-exists", true}, {"check-db-does-not-exist", false}, } - + spA := spanneraccessor.SpannerAccessorImpl{} for _, tc := range testCases { - dbExists, err := spanneraccessor.CheckExistingDb(ctx, fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectID, instanceID, tc.dbName)) + dbExists, err := spA.CheckExistingDb(ctx, fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectID, instanceID, tc.dbName)) assert.Nil(t, err) assert.Equal(t, tc.dbExists, dbExists) } diff --git a/webv2/helpers/helpers.go b/webv2/helpers/helpers.go index 7ad584746b..0511070cb3 100644 --- a/webv2/helpers/helpers.go +++ b/webv2/helpers/helpers.go @@ -160,7 +160,8 @@ func CheckOrCreateMetadataDb(projectId string, instanceId string) bool { } defer adminClient.Close() - dbExists, err := spanneraccessor.CheckExistingDb(ctx, uri) + spA := spanneraccessor.SpannerAccessorImpl{} + dbExists, err := spA.CheckExistingDb(ctx, uri) if err != nil { fmt.Println(err) return false diff --git a/webv2/session/session_service.go b/webv2/session/session_service.go index e579dd1ec6..705e2284dc 100644 --- a/webv2/session/session_service.go +++ b/webv2/session/session_service.go @@ -87,8 +87,9 @@ func migrateMetadataDb(projectId, instanceId string) { } defer adminClient.Close() + spA := spanneraccessor.SpannerAccessorImpl{} oldMetadataDbUri := getOldMetadataDbUri(projectId, instanceId) - oldMetadataDBExists, err := spanneraccessor.CheckExistingDb(ctx, oldMetadataDbUri) + oldMetadataDBExists, err := spA.CheckExistingDb(ctx, oldMetadataDbUri) if err != nil { fmt.Printf("could not check if oldMetadataDB exists. error=%v\n", err) return From d322cf4933f4c6df5e09f0fbb3249caa56072a70 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 24 Jan 2024 16:02:44 +0530 Subject: [PATCH 13/25] Add unit test for storage utils --- common/utils/storage_utils_test.go | 87 ++++++++++++++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 common/utils/storage_utils_test.go diff --git a/common/utils/storage_utils_test.go b/common/utils/storage_utils_test.go new file mode 100644 index 0000000000..c9ea7f75a5 --- /dev/null +++ b/common/utils/storage_utils_test.go @@ -0,0 +1,87 @@ +// 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 utils + +import ( + "net/url" + "os" + "testing" + + "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" + "github.com/stretchr/testify/assert" + "go.uber.org/zap" +) + +func init() { + logger.Log = zap.NewNop() +} + +func TestMain(m *testing.M) { + res := m.Run() + os.Exit(res) +} + +func TestParseGCSFilePath(t *testing.T) { + testCases := []struct { + name string + filePath string + expectError bool + want *url.URL + }{ + { + name: "Basic", + filePath: "gs://test-bucket/path/to/folder/", + expectError: false, + want: &url.URL{ + Scheme: "gs", + Host: "test-bucket", + Path: "/path/to/folder/", + }, + }, + { + name: "Append Slash", + filePath: "gs://test-bucket/path/to/folder", + expectError: false, + want: &url.URL{ + Scheme: "gs", + Host: "test-bucket", + Path: "/path/to/folder/", + }, + }, + { + name: "Empty File path", + filePath: "", + expectError: true, + want: nil, + }, + { + name: "Wrong Scheme", + filePath: "ab://testpath", + expectError: true, + want: nil, + }, + { + name: "Malformed Path", + filePath: "://path", + expectError: true, + want: nil, + }, + } + + for _, tc := range testCases { + got, err := ParseGCSFilePath(tc.filePath) + assert.Equal(t, tc.expectError, err != nil) + assert.Equal(t, tc.want, got) + } +} From 3abafe39189561a40272fecf07da0180e598cde1 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 24 Jan 2024 16:57:03 +0530 Subject: [PATCH 14/25] Add unit test for dataflow utils:UnmarshalDataflowConfig --- accessors/spanner/spanner_accessor.go | 16 +- accessors/storage/storage_accessor.go | 16 +- accessors/utils/dataflow/dataflow_utils.go | 7 +- .../utils/dataflow/dataflow_utils_test.go | 144 ++++++++++++++++++ 4 files changed, 165 insertions(+), 18 deletions(-) create mode 100644 accessors/utils/dataflow/dataflow_utils_test.go diff --git a/accessors/spanner/spanner_accessor.go b/accessors/spanner/spanner_accessor.go index 93aea4555f..a5972ae036 100644 --- a/accessors/spanner/spanner_accessor.go +++ b/accessors/spanner/spanner_accessor.go @@ -42,7 +42,7 @@ type SpannerAccessor interface { type SpannerAccessorImpl struct{} -func (sp SpannerAccessorImpl) GetDatabase(ctx context.Context, dbURI string) (*databasepb.Database, error) { +func (sp *SpannerAccessorImpl) GetDatabase(ctx context.Context, dbURI string) (*databasepb.Database, error) { adminClient, err := spanneradmin.GetOrCreateClient(ctx) if err != nil { return nil, err @@ -50,7 +50,7 @@ func (sp SpannerAccessorImpl) GetDatabase(ctx context.Context, dbURI string) (*d return adminClient.GetDatabase(ctx, &databasepb.GetDatabaseRequest{Name: dbURI}) } -func (sp SpannerAccessorImpl) GetDatabaseDialect(ctx context.Context, dbURI string) (string, error) { +func (sp *SpannerAccessorImpl) GetDatabaseDialect(ctx context.Context, dbURI string) (string, error) { result, err := sp.GetDatabase(ctx, dbURI) if err != nil { return "", fmt.Errorf("cannot connect to database: %v", err) @@ -60,7 +60,7 @@ func (sp SpannerAccessorImpl) GetDatabaseDialect(ctx context.Context, dbURI stri // CheckExistingDb checks whether the database with dbURI exists or not. // If API call doesn't respond then user is informed after every 5 minutes on command line. -func (sp SpannerAccessorImpl) CheckExistingDb(ctx context.Context, dbURI string) (bool, error) { +func (sp *SpannerAccessorImpl) CheckExistingDb(ctx context.Context, dbURI string) (bool, error) { gotResponse := make(chan bool) var err error go func() { @@ -83,7 +83,7 @@ func (sp SpannerAccessorImpl) CheckExistingDb(ctx context.Context, dbURI string) } } -func (sp SpannerAccessorImpl) CreateEmptyDatabase(ctx context.Context, dbURI string) error { +func (sp *SpannerAccessorImpl) CreateEmptyDatabase(ctx context.Context, dbURI string) error { adminClient, err := spanneradmin.GetOrCreateClient(ctx) if err != nil { return err @@ -103,7 +103,7 @@ func (sp SpannerAccessorImpl) CreateEmptyDatabase(ctx context.Context, dbURI str return nil } -func (sp SpannerAccessorImpl) GetSpannerLeaderLocation(ctx context.Context, instanceURI string) (string, error) { +func (sp *SpannerAccessorImpl) GetSpannerLeaderLocation(ctx context.Context, instanceURI string) (string, error) { instanceClient, err := spinstanceadmin.GetOrCreateClient(ctx) if err != nil { return "", err @@ -125,7 +125,7 @@ func (sp SpannerAccessorImpl) GetSpannerLeaderLocation(ctx context.Context, inst return "", fmt.Errorf("no leader found for spanner instance %s while trying fetch location", instanceURI) } -func (sp SpannerAccessorImpl) CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI string) (bool, error) { +func (sp *SpannerAccessorImpl) CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI string) (bool, error) { spClient, err := spannerclient.GetOrCreateClient(ctx, dbURI) if err != nil { return false, err @@ -157,7 +157,7 @@ func (sp SpannerAccessorImpl) CheckIfChangeStreamExists(ctx context.Context, cha return csExists, nil } -func (sp SpannerAccessorImpl) ValidateChangeStreamOptions(ctx context.Context, changeStreamName, dbURI string) error { +func (sp *SpannerAccessorImpl) ValidateChangeStreamOptions(ctx context.Context, changeStreamName, dbURI string) error { spClient, err := spannerclient.GetOrCreateClient(ctx, dbURI) if err != nil { return err @@ -192,7 +192,7 @@ func (sp SpannerAccessorImpl) ValidateChangeStreamOptions(ctx context.Context, c return nil } -func (sp SpannerAccessorImpl) CreateChangeStream(ctx context.Context, changeStreamName, dbURI string) error { +func (sp *SpannerAccessorImpl) CreateChangeStream(ctx context.Context, changeStreamName, dbURI string) error { spClient, _ := spanneradmin.GetOrCreateClient(ctx) op, err := spClient.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{ Database: dbURI, diff --git a/accessors/storage/storage_accessor.go b/accessors/storage/storage_accessor.go index 59d29e6f9c..dc0d404bb3 100644 --- a/accessors/storage/storage_accessor.go +++ b/accessors/storage/storage_accessor.go @@ -40,15 +40,15 @@ type StorageAccessor interface { type StorageAccessorImpl struct{} -func (sa StorageAccessorImpl) CreateGCSBucket(ctx context.Context, bucketName, projectID, location string) error { +func (sa *StorageAccessorImpl) CreateGCSBucket(ctx context.Context, bucketName, projectID, location string) error { return sa.createGCSBucketUtil(ctx, bucketName, projectID, location, nil, 0) } -func (sa StorageAccessorImpl) CreateGCSBucketWithLifecycle(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { +func (sa *StorageAccessorImpl) CreateGCSBucketWithLifecycle(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { return sa.createGCSBucketUtil(ctx, bucketName, projectID, location, matchesPrefix, ttl) } -func (sa StorageAccessorImpl) createGCSBucketUtil(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { +func (sa *StorageAccessorImpl) createGCSBucketUtil(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { client, err := storageclient.GetOrCreateClient(ctx) if err != nil { return err @@ -95,7 +95,7 @@ func (sa StorageAccessorImpl) createGCSBucketUtil(ctx context.Context, bucketNam // Applies the bucket lifecycle with delete rule. Only accepts the Age and // prefix rule conditions as it is only used for the Datastream destination // bucket currently. -func (sa StorageAccessorImpl) EnableBucketLifecycleDeleteRule(ctx context.Context, bucketName string, matchesPrefix []string, ttl int64) error { +func (sa *StorageAccessorImpl) EnableBucketLifecycleDeleteRule(ctx context.Context, bucketName string, matchesPrefix []string, ttl int64) error { client, err := storageclient.GetOrCreateClient(ctx) if err != nil { return fmt.Errorf("could not create client while enabling lifecycle: %w", err) @@ -132,7 +132,7 @@ func (sa StorageAccessorImpl) EnableBucketLifecycleDeleteRule(ctx context.Contex } // UploadLocalFileToGCS uploads an object. -func (sa StorageAccessorImpl) UploadLocalFileToGCS(ctx context.Context, filePath, fileName, localFilePath string) error { +func (sa *StorageAccessorImpl) UploadLocalFileToGCS(ctx context.Context, filePath, fileName, localFilePath string) error { data, err := os.ReadFile(localFilePath) if err != nil { return fmt.Errorf("could not read file %s: %w", localFilePath, err) @@ -140,7 +140,7 @@ func (sa StorageAccessorImpl) UploadLocalFileToGCS(ctx context.Context, filePath return sa.WriteDataToGCS(ctx, filePath, fileName, string(data)) } -func (sa StorageAccessorImpl) WriteDataToGCS(ctx context.Context, filePath, fileName, data string) error { +func (sa *StorageAccessorImpl) WriteDataToGCS(ctx context.Context, filePath, fileName, data string) error { client, err := storageclient.GetOrCreateClient(ctx) if err != nil { return fmt.Errorf("could not create client while uploading to GCS: %w", err) @@ -170,7 +170,7 @@ func (sa StorageAccessorImpl) WriteDataToGCS(ctx context.Context, filePath, file return nil } -func (sa StorageAccessorImpl) ReadGcsFile(ctx context.Context, filePath string) (string, error) { +func (sa *StorageAccessorImpl) ReadGcsFile(ctx context.Context, filePath string) (string, error) { client, err := storageclient.GetOrCreateClient(ctx) if err != nil { return "", fmt.Errorf("could not create client: %w", err) @@ -199,7 +199,7 @@ func (sa StorageAccessorImpl) ReadGcsFile(ctx context.Context, filePath string) return buf.String(), nil } -func (sa StorageAccessorImpl) ReadAnyFile(ctx context.Context, filePath string) (string, error) { +func (sa *StorageAccessorImpl) ReadAnyFile(ctx context.Context, filePath string) (string, error) { if strings.HasPrefix(filePath, constants.GCS_FILE_PREFIX) { return sa.ReadGcsFile(ctx, filePath) } diff --git a/accessors/utils/dataflow/dataflow_utils.go b/accessors/utils/dataflow/dataflow_utils.go index 6b8978ef3e..3d52512636 100644 --- a/accessors/utils/dataflow/dataflow_utils.go +++ b/accessors/utils/dataflow/dataflow_utils.go @@ -11,6 +11,9 @@ // 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. + +// This is a package is kept with accessors because some functions import other accessors. +// The common/utils package should not import any SMT dependency. package dataflowutils import ( @@ -21,8 +24,8 @@ import ( storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" ) -func UnmarshalDataflowTuningConfig(ctx context.Context, filePath string) (dataflowaccessor.DataflowTuningConfig, error) { - jsonStr, err := storageaccessor.ReadAnyFile(ctx, filePath) +func UnmarshalDataflowTuningConfig(ctx context.Context, sa storageaccessor.StorageAccessor, filePath string) (dataflowaccessor.DataflowTuningConfig, error) { + jsonStr, err := sa.ReadAnyFile(ctx, filePath) if err != nil { return dataflowaccessor.DataflowTuningConfig{}, err } diff --git a/accessors/utils/dataflow/dataflow_utils_test.go b/accessors/utils/dataflow/dataflow_utils_test.go new file mode 100644 index 0000000000..a5195e0dd2 --- /dev/null +++ b/accessors/utils/dataflow/dataflow_utils_test.go @@ -0,0 +1,144 @@ +// Copyright 2023 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 dataflowutils + +import ( + "context" + "fmt" + "os" + "testing" + + dataflowaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/dataflow" + storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" + "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" + "github.com/stretchr/testify/assert" + "go.uber.org/zap" +) + +func init() { + logger.Log = zap.NewNop() +} + +func TestMain(m *testing.M) { + res := m.Run() + os.Exit(res) +} + +type StorageAccessorMock struct { + storageaccessor.StorageAccessorImpl +} + +var readAnyFileMock func(ctx context.Context, filePath string) (string, error) + +func (sam StorageAccessorMock) ReadAnyFile(ctx context.Context, filePath string) (string, error) { + return readAnyFileMock(ctx, filePath) +} + +func TestUnmarshalDataflowTuningConfig(t *testing.T) { + testCases := []struct { + name string + readAnyFileMock func(ctx context.Context, filePath string) (string, error) + expectError bool + want dataflowaccessor.DataflowTuningConfig + }{ + { + name: "Basic", + readAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + return `{ + "projectId": "test-project", + "jobName": "test-job-name", + "location": "us-central1", + "network": "test-network", + "subnetwork": "test-subnetwork", + "hostProjectId": "test-host-project", + "maxWorkers": 3, + "numWorkers": 2, + "serviceAccountEmail": "abc@xyz.com", + "machineType": "n1-standard-8", + "additionalUserLabels": {"my": "label"}, + "kmsKeyName": "test-key", + "gcsTemplatePath": "gs://path", + "additionalExperiments": ["xyz","123"], + "enableStreamingEngine": true + }`, nil + }, + expectError: false, + want: dataflowaccessor.DataflowTuningConfig{ + ProjectId: "test-project", + JobName: "test-job-name", + Location: "us-central1", + Network: "test-network", + Subnetwork: "test-subnetwork", + VpcHostProjectId: "test-host-project", + MaxWorkers: 3, + NumWorkers: 2, + ServiceAccountEmail: "abc@xyz.com", + MachineType: "n1-standard-8", + AdditionalUserLabels: map[string]string{"my": "label"}, + KmsKeyName: "test-key", + GcsTemplatePath: "gs://path", + AdditionalExperiments: []string{"xyz", "123"}, + EnableStreamingEngine: true, + }, + }, + { + name: "Defaults", + readAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + return `{}`, nil + }, + expectError: false, + want: dataflowaccessor.DataflowTuningConfig{ + ProjectId: "", + JobName: "", + Location: "", + Network: "", + Subnetwork: "", + VpcHostProjectId: "", + MaxWorkers: 0, + NumWorkers: 0, + ServiceAccountEmail: "", + MachineType: "", + AdditionalUserLabels: nil, + KmsKeyName: "", + GcsTemplatePath: "", + AdditionalExperiments: nil, + EnableStreamingEngine: false, + }, + }, + { + name: "ReadAnyFile throws error", + readAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + return "", fmt.Errorf("test error") + }, + expectError: true, + want: dataflowaccessor.DataflowTuningConfig{}, + }, + { + name: "Json unmarshall throws error", + readAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + return "{\"abc\"", nil + }, + expectError: true, + want: dataflowaccessor.DataflowTuningConfig{}, + }, + } + ctx := context.Background() + saMock := StorageAccessorMock{} + for _, tc := range testCases { + readAnyFileMock = tc.readAnyFileMock + got, err := UnmarshalDataflowTuningConfig(ctx, &saMock, "unused/path/due/to/mock") + assert.Equal(t, tc.expectError, err != nil) + assert.Equal(t, tc.want, got) + } +} From c8993599e08ea665c9102a0549b7ffa0ad039cde Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Mon, 29 Jan 2024 18:11:37 +0530 Subject: [PATCH 15/25] Move mock methods inside mock struct --- .../utils/dataflow/dataflow_utils_test.go | 73 ++++++++++--------- 1 file changed, 39 insertions(+), 34 deletions(-) diff --git a/accessors/utils/dataflow/dataflow_utils_test.go b/accessors/utils/dataflow/dataflow_utils_test.go index a5195e0dd2..49638589c9 100644 --- a/accessors/utils/dataflow/dataflow_utils_test.go +++ b/accessors/utils/dataflow/dataflow_utils_test.go @@ -37,41 +37,42 @@ func TestMain(m *testing.M) { type StorageAccessorMock struct { storageaccessor.StorageAccessorImpl + ReadAnyFileMock func(ctx context.Context, filePath string) (string, error) } -var readAnyFileMock func(ctx context.Context, filePath string) (string, error) - func (sam StorageAccessorMock) ReadAnyFile(ctx context.Context, filePath string) (string, error) { - return readAnyFileMock(ctx, filePath) + return sam.ReadAnyFileMock(ctx, filePath) } func TestUnmarshalDataflowTuningConfig(t *testing.T) { testCases := []struct { - name string - readAnyFileMock func(ctx context.Context, filePath string) (string, error) - expectError bool - want dataflowaccessor.DataflowTuningConfig + name string + sam StorageAccessorMock + expectError bool + want dataflowaccessor.DataflowTuningConfig }{ { name: "Basic", - readAnyFileMock: func(ctx context.Context, filePath string) (string, error) { - return `{ - "projectId": "test-project", - "jobName": "test-job-name", - "location": "us-central1", - "network": "test-network", - "subnetwork": "test-subnetwork", - "hostProjectId": "test-host-project", - "maxWorkers": 3, - "numWorkers": 2, - "serviceAccountEmail": "abc@xyz.com", - "machineType": "n1-standard-8", - "additionalUserLabels": {"my": "label"}, - "kmsKeyName": "test-key", - "gcsTemplatePath": "gs://path", - "additionalExperiments": ["xyz","123"], - "enableStreamingEngine": true - }`, nil + sam: StorageAccessorMock{ + ReadAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + return `{ + "projectId": "test-project", + "jobName": "test-job-name", + "location": "us-central1", + "network": "test-network", + "subnetwork": "test-subnetwork", + "hostProjectId": "test-host-project", + "maxWorkers": 3, + "numWorkers": 2, + "serviceAccountEmail": "abc@xyz.com", + "machineType": "n1-standard-8", + "additionalUserLabels": {"my": "label"}, + "kmsKeyName": "test-key", + "gcsTemplatePath": "gs://path", + "additionalExperiments": ["xyz","123"], + "enableStreamingEngine": true + }`, nil + }, }, expectError: false, want: dataflowaccessor.DataflowTuningConfig{ @@ -94,8 +95,10 @@ func TestUnmarshalDataflowTuningConfig(t *testing.T) { }, { name: "Defaults", - readAnyFileMock: func(ctx context.Context, filePath string) (string, error) { - return `{}`, nil + sam: StorageAccessorMock{ + ReadAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + return `{}`, nil + }, }, expectError: false, want: dataflowaccessor.DataflowTuningConfig{ @@ -118,26 +121,28 @@ func TestUnmarshalDataflowTuningConfig(t *testing.T) { }, { name: "ReadAnyFile throws error", - readAnyFileMock: func(ctx context.Context, filePath string) (string, error) { - return "", fmt.Errorf("test error") + sam: StorageAccessorMock{ + ReadAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + return "", fmt.Errorf("test error") + }, }, expectError: true, want: dataflowaccessor.DataflowTuningConfig{}, }, { name: "Json unmarshall throws error", - readAnyFileMock: func(ctx context.Context, filePath string) (string, error) { - return "{\"abc\"", nil + sam: StorageAccessorMock{ + ReadAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + return "{\"abc\"", nil + }, }, expectError: true, want: dataflowaccessor.DataflowTuningConfig{}, }, } ctx := context.Background() - saMock := StorageAccessorMock{} for _, tc := range testCases { - readAnyFileMock = tc.readAnyFileMock - got, err := UnmarshalDataflowTuningConfig(ctx, &saMock, "unused/path/due/to/mock") + got, err := UnmarshalDataflowTuningConfig(ctx, &tc.sam, "unused/path/due/to/mock") assert.Equal(t, tc.expectError, err != nil) assert.Equal(t, tc.want, got) } From 94e0c921c3a580e12c2f31d875c3bf942ff0ed63 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Tue, 30 Jan 2024 16:50:57 +0530 Subject: [PATCH 16/25] Add wrapper for storage client --- accessors/clients/storage/mocks.go | 14 +++++ accessors/clients/storage/storage_client.go | 60 +++++++++++++++++++ accessors/storage/mocks.go | 58 ++++++++++++++++++ accessors/storage/storage_accessor.go | 59 +++++++----------- accessors/utils/dataflow/dataflow_utils.go | 5 +- .../utils/dataflow/dataflow_utils_test.go | 30 ++++------ conversion/conversion.go | 13 +++- streaming/streaming.go | 9 ++- webv2/profile/profile.go | 8 ++- webv2/web.go | 9 ++- 10 files changed, 200 insertions(+), 65 deletions(-) create mode 100644 accessors/clients/storage/mocks.go create mode 100644 accessors/storage/mocks.go diff --git a/accessors/clients/storage/mocks.go b/accessors/clients/storage/mocks.go new file mode 100644 index 0000000000..02bff2f3ac --- /dev/null +++ b/accessors/clients/storage/mocks.go @@ -0,0 +1,14 @@ +// 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 storageclient diff --git a/accessors/clients/storage/storage_client.go b/accessors/clients/storage/storage_client.go index 28c62b868f..4ce4026b83 100644 --- a/accessors/clients/storage/storage_client.go +++ b/accessors/clients/storage/storage_client.go @@ -16,6 +16,7 @@ package storageclient import ( "context" "fmt" + "io" "sync" "cloud.google.com/go/storage" @@ -41,3 +42,62 @@ func GetOrCreateClient(ctx context.Context) (*storage.Client, error) { } return gcsClient, nil } + +type StorageClient interface { + Bucket(name string) BucketHandle +} + +type BucketHandle interface { + Create(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) + Update(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) + Object(name string) ObjectHandle +} + +type ObjectHandle interface { + NewWriter(ctx context.Context) io.WriteCloser + NewReader(ctx context.Context) (io.ReadCloser, error) +} + +type StorageClientImpl struct { + client *storage.Client +} + +func NewStorageClientImpl(ctx context.Context) (*StorageClientImpl, error) { + c, err := GetOrCreateClient(ctx) + if err != nil { + return nil, err + } + return &StorageClientImpl{client: c}, nil +} + +func (c *StorageClientImpl) Bucket(name string) BucketHandle { + return &BucketHandleImpl{bucketHandle: c.client.Bucket(name)} +} + +type BucketHandleImpl struct { + bucketHandle *storage.BucketHandle +} + +func (b *BucketHandleImpl) Create(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) { + return b.bucketHandle.Create(ctx, projectID, attrs) +} + +func (b *BucketHandleImpl) Update(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) { + return b.bucketHandle.Update(ctx, uattrs) +} + +func (b *BucketHandleImpl) Object(name string) ObjectHandle { + return &ObjectHandleImpl{objectHandle: b.bucketHandle.Object(name)} +} + +type ObjectHandleImpl struct { + objectHandle *storage.ObjectHandle +} + +func (o *ObjectHandleImpl) NewWriter(ctx context.Context) io.WriteCloser { + return o.objectHandle.NewWriter(ctx) +} + +func (o *ObjectHandleImpl) NewReader(ctx context.Context) (io.ReadCloser, error) { + return o.objectHandle.NewReader(ctx) +} diff --git a/accessors/storage/mocks.go b/accessors/storage/mocks.go new file mode 100644 index 0000000000..7c3f906f2a --- /dev/null +++ b/accessors/storage/mocks.go @@ -0,0 +1,58 @@ +// 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 storageaccessor + +import ( + "context" + + storageclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/storage" +) + +type StorageAccessorMock struct { + CreateGCSBucketMock func(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string) error + CreateGCSBucketWithLifecycleMock func(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error + EnableBucketLifecycleDeleteRuleMock func(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error + UploadLocalFileToGCSMock func(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error + WriteDataToGCSMock func(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, data string) error + ReadGcsFileMock func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) + ReadAnyFileMock func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) +} + +func (sam *StorageAccessorMock) CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string) error { + return sam.CreateGCSBucketMock(ctx, sc, bucketName, projectID, location) +} + +func (sam *StorageAccessorMock) CreateGCSBucketWithLifecycle(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { + return sam.CreateGCSBucketWithLifecycleMock(ctx, sc, bucketName, projectID, location, matchesPrefix, ttl) +} + +func (sam *StorageAccessorMock) EnableBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error { + return sam.EnableBucketLifecycleDeleteRuleMock(ctx, sc, bucketName, matchesPrefix, ttl) +} + +func (sam *StorageAccessorMock) UploadLocalFileToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error { + return sam.UploadLocalFileToGCSMock(ctx, sc, filePath, fileName, localFilePath) +} + +func (sam *StorageAccessorMock) WriteDataToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, data string) error { + return sam.WriteDataToGCSMock(ctx, sc, filePath, fileName, data) +} + +func (sam *StorageAccessorMock) ReadGcsFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { + return sam.ReadGcsFileMock(ctx, sc, filePath) +} + +func (sam *StorageAccessorMock) ReadAnyFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { + return sam.ReadAnyFileMock(ctx, sc, filePath) +} diff --git a/accessors/storage/storage_accessor.go b/accessors/storage/storage_accessor.go index dc0d404bb3..ee3202db22 100644 --- a/accessors/storage/storage_accessor.go +++ b/accessors/storage/storage_accessor.go @@ -29,26 +29,26 @@ import ( ) type StorageAccessor interface { - CreateGCSBucket(ctx context.Context, bucketName, projectID, location string) error - CreateGCSBucketWithLifecycle(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error - EnableBucketLifecycleDeleteRule(ctx context.Context, bucketName string, matchesPrefix []string, ttl int64) error - UploadLocalFileToGCS(ctx context.Context, filePath, fileName, localFilePath string) error - WriteDataToGCS(ctx context.Context, filePath, fileName, data string) error - ReadGcsFile(ctx context.Context, filePath string) (string, error) - ReadAnyFile(ctx context.Context, filePath string) (string, error) + CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string) error + CreateGCSBucketWithLifecycle(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error + EnableBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error + UploadLocalFileToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error + WriteDataToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, data string) error + ReadGcsFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) + ReadAnyFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) } type StorageAccessorImpl struct{} -func (sa *StorageAccessorImpl) CreateGCSBucket(ctx context.Context, bucketName, projectID, location string) error { - return sa.createGCSBucketUtil(ctx, bucketName, projectID, location, nil, 0) +func (sa *StorageAccessorImpl) CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string) error { + return sa.createGCSBucketUtil(ctx, sc, bucketName, projectID, location, nil, 0) } -func (sa *StorageAccessorImpl) CreateGCSBucketWithLifecycle(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { - return sa.createGCSBucketUtil(ctx, bucketName, projectID, location, matchesPrefix, ttl) +func (sa *StorageAccessorImpl) CreateGCSBucketWithLifecycle(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { + return sa.createGCSBucketUtil(ctx, sc, bucketName, projectID, location, matchesPrefix, ttl) } -func (sa *StorageAccessorImpl) createGCSBucketUtil(ctx context.Context, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { +func (sa *StorageAccessorImpl) createGCSBucketUtil(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { client, err := storageclient.GetOrCreateClient(ctx) if err != nil { return err @@ -95,16 +95,11 @@ func (sa *StorageAccessorImpl) createGCSBucketUtil(ctx context.Context, bucketNa // Applies the bucket lifecycle with delete rule. Only accepts the Age and // prefix rule conditions as it is only used for the Datastream destination // bucket currently. -func (sa *StorageAccessorImpl) EnableBucketLifecycleDeleteRule(ctx context.Context, bucketName string, matchesPrefix []string, ttl int64) error { - client, err := storageclient.GetOrCreateClient(ctx) - if err != nil { - return fmt.Errorf("could not create client while enabling lifecycle: %w", err) - } - +func (sa *StorageAccessorImpl) EnableBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error { for i, str := range matchesPrefix { matchesPrefix[i] = strings.TrimPrefix(str, "/") } - bucket := client.Bucket(bucketName) + bucket := sc.Bucket(bucketName) bucketAttrsToUpdate := storage.BucketAttrsToUpdate{ Lifecycle: &storage.Lifecycle{ Rules: []storage.LifecycleRule{ @@ -132,26 +127,21 @@ func (sa *StorageAccessorImpl) EnableBucketLifecycleDeleteRule(ctx context.Conte } // UploadLocalFileToGCS uploads an object. -func (sa *StorageAccessorImpl) UploadLocalFileToGCS(ctx context.Context, filePath, fileName, localFilePath string) error { +func (sa *StorageAccessorImpl) UploadLocalFileToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error { data, err := os.ReadFile(localFilePath) if err != nil { return fmt.Errorf("could not read file %s: %w", localFilePath, err) } - return sa.WriteDataToGCS(ctx, filePath, fileName, string(data)) + return sa.WriteDataToGCS(ctx, sc, filePath, fileName, string(data)) } -func (sa *StorageAccessorImpl) WriteDataToGCS(ctx context.Context, filePath, fileName, data string) error { - client, err := storageclient.GetOrCreateClient(ctx) - if err != nil { - return fmt.Errorf("could not create client while uploading to GCS: %w", err) - } - +func (sa *StorageAccessorImpl) WriteDataToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, data string) error { u, err := utils.ParseGCSFilePath(filePath) if err != nil { return fmt.Errorf("parseFilePath: unable to parse file path: %v", err) } bucketName := u.Host - bucket := client.Bucket(bucketName) + bucket := sc.Bucket(bucketName) obj := bucket.Object(u.Path[1:] + fileName) w := obj.NewWriter(ctx) @@ -170,18 +160,13 @@ func (sa *StorageAccessorImpl) WriteDataToGCS(ctx context.Context, filePath, fil return nil } -func (sa *StorageAccessorImpl) ReadGcsFile(ctx context.Context, filePath string) (string, error) { - client, err := storageclient.GetOrCreateClient(ctx) - if err != nil { - return "", fmt.Errorf("could not create client: %w", err) - } - +func (sa *StorageAccessorImpl) ReadGcsFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { u, err := utils.ParseGCSFilePath(filePath) if err != nil { return "", fmt.Errorf("unable to parse file path: %v", err) } bucketName := u.Host - bucket := client.Bucket(bucketName) + bucket := sc.Bucket(bucketName) obj := bucket.Object(u.Path[1:]) rc, err := obj.NewReader(ctx) @@ -199,9 +184,9 @@ func (sa *StorageAccessorImpl) ReadGcsFile(ctx context.Context, filePath string) return buf.String(), nil } -func (sa *StorageAccessorImpl) ReadAnyFile(ctx context.Context, filePath string) (string, error) { +func (sa *StorageAccessorImpl) ReadAnyFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { if strings.HasPrefix(filePath, constants.GCS_FILE_PREFIX) { - return sa.ReadGcsFile(ctx, filePath) + return sa.ReadGcsFile(ctx, sc, filePath) } buf, err := os.ReadFile(filePath) if err != nil { diff --git a/accessors/utils/dataflow/dataflow_utils.go b/accessors/utils/dataflow/dataflow_utils.go index 3d52512636..476981a061 100644 --- a/accessors/utils/dataflow/dataflow_utils.go +++ b/accessors/utils/dataflow/dataflow_utils.go @@ -20,12 +20,13 @@ import ( "context" "encoding/json" + storageclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/storage" dataflowaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/dataflow" storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" ) -func UnmarshalDataflowTuningConfig(ctx context.Context, sa storageaccessor.StorageAccessor, filePath string) (dataflowaccessor.DataflowTuningConfig, error) { - jsonStr, err := sa.ReadAnyFile(ctx, filePath) +func UnmarshalDataflowTuningConfig(ctx context.Context, sc storageclient.StorageClient, sa storageaccessor.StorageAccessor, filePath string) (dataflowaccessor.DataflowTuningConfig, error) { + jsonStr, err := sa.ReadAnyFile(ctx, sc, filePath) if err != nil { return dataflowaccessor.DataflowTuningConfig{}, err } diff --git a/accessors/utils/dataflow/dataflow_utils_test.go b/accessors/utils/dataflow/dataflow_utils_test.go index 49638589c9..4b9edcc6d5 100644 --- a/accessors/utils/dataflow/dataflow_utils_test.go +++ b/accessors/utils/dataflow/dataflow_utils_test.go @@ -19,6 +19,7 @@ import ( "os" "testing" + storageclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/storage" dataflowaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/dataflow" storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" @@ -35,26 +36,17 @@ func TestMain(m *testing.M) { os.Exit(res) } -type StorageAccessorMock struct { - storageaccessor.StorageAccessorImpl - ReadAnyFileMock func(ctx context.Context, filePath string) (string, error) -} - -func (sam StorageAccessorMock) ReadAnyFile(ctx context.Context, filePath string) (string, error) { - return sam.ReadAnyFileMock(ctx, filePath) -} - func TestUnmarshalDataflowTuningConfig(t *testing.T) { testCases := []struct { name string - sam StorageAccessorMock + sam storageaccessor.StorageAccessorMock expectError bool want dataflowaccessor.DataflowTuningConfig }{ { name: "Basic", - sam: StorageAccessorMock{ - ReadAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + sam: storageaccessor.StorageAccessorMock{ + ReadAnyFileMock: func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { return `{ "projectId": "test-project", "jobName": "test-job-name", @@ -95,8 +87,8 @@ func TestUnmarshalDataflowTuningConfig(t *testing.T) { }, { name: "Defaults", - sam: StorageAccessorMock{ - ReadAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + sam: storageaccessor.StorageAccessorMock{ + ReadAnyFileMock: func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { return `{}`, nil }, }, @@ -121,8 +113,8 @@ func TestUnmarshalDataflowTuningConfig(t *testing.T) { }, { name: "ReadAnyFile throws error", - sam: StorageAccessorMock{ - ReadAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + sam: storageaccessor.StorageAccessorMock{ + ReadAnyFileMock: func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { return "", fmt.Errorf("test error") }, }, @@ -131,8 +123,8 @@ func TestUnmarshalDataflowTuningConfig(t *testing.T) { }, { name: "Json unmarshall throws error", - sam: StorageAccessorMock{ - ReadAnyFileMock: func(ctx context.Context, filePath string) (string, error) { + sam: storageaccessor.StorageAccessorMock{ + ReadAnyFileMock: func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { return "{\"abc\"", nil }, }, @@ -142,7 +134,7 @@ func TestUnmarshalDataflowTuningConfig(t *testing.T) { } ctx := context.Background() for _, tc := range testCases { - got, err := UnmarshalDataflowTuningConfig(ctx, &tc.sam, "unused/path/due/to/mock") + got, err := UnmarshalDataflowTuningConfig(ctx, nil, &tc.sam, "unused/path/due/to/mock") assert.Equal(t, tc.expectError, err != nil) assert.Equal(t, tc.want, got) } diff --git a/conversion/conversion.go b/conversion/conversion.go index bf02976781..efbe86dec1 100644 --- a/conversion/conversion.go +++ b/conversion/conversion.go @@ -44,6 +44,7 @@ import ( datastream "cloud.google.com/go/datastream/apiv1" sp "cloud.google.com/go/spanner" database "cloud.google.com/go/spanner/admin/database/apiv1" + storageclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/storage" spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" @@ -356,9 +357,13 @@ func dataFromDatabase(ctx context.Context, sourceProfile profiles.SourceProfile, // Try to apply lifecycle rule to Datastream destination bucket. gcsConfig := streamingCfg.GcsCfg + sc, err := storageclient.NewStorageClientImpl(ctx) + if err != nil { + return nil, err + } sa := storageaccessor.StorageAccessorImpl{} if gcsConfig.TtlInDaysSet { - err = sa.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = sa.EnableBucketLifecycleDeleteRule(ctx, sc, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") @@ -478,9 +483,13 @@ func dataFromDatabaseForDataflowMigration(targetProfile profiles.TargetProfile, // Try to apply lifecycle rule to Datastream destination bucket. gcsConfig := streamingCfg.GcsCfg + sc, err := storageclient.NewStorageClientImpl(ctx) + if err != nil { + return common.TaskResult[*profiles.DataShard]{Result: p, Err: err} + } sa := storageaccessor.StorageAccessorImpl{} if gcsConfig.TtlInDaysSet { - err = sa.EnableBucketLifecycleDeleteRule(ctx, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = sa.EnableBucketLifecycleDeleteRule(ctx, sc, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") diff --git a/streaming/streaming.go b/streaming/streaming.go index f00af44198..550ffffc31 100644 --- a/streaming/streaming.go +++ b/streaming/streaming.go @@ -33,6 +33,7 @@ import ( resourcemanager "cloud.google.com/go/resourcemanager/apiv3" resourcemanagerpb "cloud.google.com/go/resourcemanager/apiv3/resourcemanagerpb" + storageclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/storage" dataflowaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/dataflow" storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" @@ -861,12 +862,16 @@ func StartDatastream(ctx context.Context, streamingCfg StreamingCfg, sourceProfi } func StartDataflow(ctx context.Context, targetProfile profiles.TargetProfile, streamingCfg StreamingCfg, conv *internal.Conv) (internal.DataflowOutput, error) { + sc, err := storageclient.NewStorageClientImpl(ctx) + if err != nil { + return internal.DataflowOutput{}, err + } sa := storageaccessor.StorageAccessorImpl{} convJSON, err := json.MarshalIndent(conv, "", " ") if err != nil { return internal.DataflowOutput{}, fmt.Errorf("can't encode session state to JSON: %v", err) } - err = sa.WriteDataToGCS(ctx, streamingCfg.TmpDir, "session.json", string(convJSON)) + err = sa.WriteDataToGCS(ctx, sc, streamingCfg.TmpDir, "session.json", string(convJSON)) if err != nil { return internal.DataflowOutput{}, fmt.Errorf("error while writing to GCS: %v", err) } @@ -877,7 +882,7 @@ func StartDataflow(ctx context.Context, targetProfile profiles.TargetProfile, st if err != nil { return internal.DataflowOutput{}, fmt.Errorf("failed to compute transformation context: %s", err.Error()) } - err = sa.WriteDataToGCS(ctx, streamingCfg.TmpDir, "transformationContext.json", string(transformationContext)) + err = sa.WriteDataToGCS(ctx, sc, streamingCfg.TmpDir, "transformationContext.json", string(transformationContext)) if err != nil { return internal.DataflowOutput{}, fmt.Errorf("error while writing to GCS: %v", err) } diff --git a/webv2/profile/profile.go b/webv2/profile/profile.go index 923bb5d84d..29b0958d43 100644 --- a/webv2/profile/profile.go +++ b/webv2/profile/profile.go @@ -11,6 +11,7 @@ import ( "strings" datastream "cloud.google.com/go/datastream/apiv1" + storageclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/storage" storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" @@ -154,6 +155,11 @@ func CreateConnectionProfile(w http.ResponseWriter, r *http.Request) { ValidateOnly: details.ValidateOnly, } var bucketName string + sc, err := storageclient.NewStorageClientImpl(ctx) + if err != nil { + http.Error(w, fmt.Sprintf("Error while StorageClientImpl: %v", err), http.StatusBadRequest) + return + } sa := storageaccessor.StorageAccessorImpl{} if !details.IsSource { @@ -162,7 +168,7 @@ func CreateConnectionProfile(w http.ResponseWriter, r *http.Request) { } else { bucketName = strings.ToLower(sessionState.Conv.Audit.MigrationRequestId) } - err = sa.CreateGCSBucket(ctx, bucketName, sessionState.GCPProjectID, sessionState.Region) + err = sa.CreateGCSBucket(ctx, sc, bucketName, sessionState.GCPProjectID, sessionState.Region) if err != nil { http.Error(w, fmt.Sprintf("Error while creating bucket: %v", err), http.StatusBadRequest) return diff --git a/webv2/web.go b/webv2/web.go index c45dab3514..971d42c701 100644 --- a/webv2/web.go +++ b/webv2/web.go @@ -36,6 +36,7 @@ import ( "time" instance "cloud.google.com/go/spanner/admin/instance/apiv1" + storageclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/storage" storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" "github.com/GoogleCloudPlatform/spanner-migration-tool/cmd" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" @@ -2486,8 +2487,12 @@ func createConfigFileForShardedBulkMigration(sessionState *session.SessionState, } func writeSessionFile(ctx context.Context, sessionState *session.SessionState) error { + sc, err := storageclient.NewStorageClientImpl(ctx) + if err != nil { + return err + } sa := storageaccessor.StorageAccessorImpl{} - err := sa.CreateGCSBucket(ctx, sessionState.Bucket, sessionState.GCPProjectID, sessionState.Region) + err = sa.CreateGCSBucket(ctx, sc, sessionState.Bucket, sessionState.GCPProjectID, sessionState.Region) if err != nil { return fmt.Errorf("error while creating bucket: %v", err) } @@ -2496,7 +2501,7 @@ func writeSessionFile(ctx context.Context, sessionState *session.SessionState) e if err != nil { return fmt.Errorf("can't encode session state to JSON: %v", err) } - err = sa.WriteDataToGCS(ctx, "gs://"+sessionState.Bucket+sessionState.RootPath, "session.json", string(convJSON)) + err = sa.WriteDataToGCS(ctx, sc, "gs://"+sessionState.Bucket+sessionState.RootPath, "session.json", string(convJSON)) if err != nil { return fmt.Errorf("error while writing to GCS: %v", err) } From f18a287a079f8c5cdb72dc56c09ddc99e1f53378 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Tue, 30 Jan 2024 17:14:57 +0530 Subject: [PATCH 17/25] Resolved storage comments --- .../dataflow/dataflow_helpers.go} | 6 +++- .../dataflow/dataflow_helpers_test.go} | 6 ++-- accessors/storage/mocks.go | 25 ++++++-------- accessors/storage/storage_accessor.go | 33 ++++++++----------- common/utils/storage_utils_test.go | 24 ++++++++++++-- conversion/conversion.go | 4 +-- webv2/profile/profile.go | 2 +- webv2/web.go | 2 +- 8 files changed, 58 insertions(+), 44 deletions(-) rename accessors/{utils/dataflow/dataflow_utils.go => helpers/dataflow/dataflow_helpers.go} (82%) rename accessors/{utils/dataflow/dataflow_utils_test.go => helpers/dataflow/dataflow_helpers_test.go} (97%) diff --git a/accessors/utils/dataflow/dataflow_utils.go b/accessors/helpers/dataflow/dataflow_helpers.go similarity index 82% rename from accessors/utils/dataflow/dataflow_utils.go rename to accessors/helpers/dataflow/dataflow_helpers.go index 476981a061..ee117ccced 100644 --- a/accessors/utils/dataflow/dataflow_utils.go +++ b/accessors/helpers/dataflow/dataflow_helpers.go @@ -14,7 +14,7 @@ // This is a package is kept with accessors because some functions import other accessors. // The common/utils package should not import any SMT dependency. -package dataflowutils +package dataflowhelpers import ( "context" @@ -25,6 +25,10 @@ import ( storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" ) +// This package contains common helper methods using the accessor package. This will be used by multiple flows. +// Do not move to util since util only expects methods that do no import any other internal dependency. + +// Reads any local or gcs file and unmarshals the data into a DataflowTuningConfig struct. func UnmarshalDataflowTuningConfig(ctx context.Context, sc storageclient.StorageClient, sa storageaccessor.StorageAccessor, filePath string) (dataflowaccessor.DataflowTuningConfig, error) { jsonStr, err := sa.ReadAnyFile(ctx, sc, filePath) if err != nil { diff --git a/accessors/utils/dataflow/dataflow_utils_test.go b/accessors/helpers/dataflow/dataflow_helpers_test.go similarity index 97% rename from accessors/utils/dataflow/dataflow_utils_test.go rename to accessors/helpers/dataflow/dataflow_helpers_test.go index 4b9edcc6d5..1418dea448 100644 --- a/accessors/utils/dataflow/dataflow_utils_test.go +++ b/accessors/helpers/dataflow/dataflow_helpers_test.go @@ -11,7 +11,7 @@ // 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 dataflowutils +package dataflowhelpers import ( "context" @@ -135,7 +135,7 @@ func TestUnmarshalDataflowTuningConfig(t *testing.T) { ctx := context.Background() for _, tc := range testCases { got, err := UnmarshalDataflowTuningConfig(ctx, nil, &tc.sam, "unused/path/due/to/mock") - assert.Equal(t, tc.expectError, err != nil) - assert.Equal(t, tc.want, got) + assert.Equal(t, tc.expectError, err != nil, tc.name) + assert.Equal(t, tc.want, got, tc.name) } } diff --git a/accessors/storage/mocks.go b/accessors/storage/mocks.go index 7c3f906f2a..0840a90737 100644 --- a/accessors/storage/mocks.go +++ b/accessors/storage/mocks.go @@ -20,25 +20,20 @@ import ( ) type StorageAccessorMock struct { - CreateGCSBucketMock func(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string) error - CreateGCSBucketWithLifecycleMock func(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error - EnableBucketLifecycleDeleteRuleMock func(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error - UploadLocalFileToGCSMock func(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error - WriteDataToGCSMock func(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, data string) error - ReadGcsFileMock func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) - ReadAnyFileMock func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) + CreateGCSBucketMock func(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, ttl int64, matchesPrefix []string) error + ApplyBucketLifecycleDeleteRuleMock func(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error + UploadLocalFileToGCSMock func(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error + WriteDataToGCSMock func(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, data string) error + ReadGcsFileMock func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) + ReadAnyFileMock func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) } -func (sam *StorageAccessorMock) CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string) error { - return sam.CreateGCSBucketMock(ctx, sc, bucketName, projectID, location) +func (sam *StorageAccessorMock) CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, ttl int64, matchesPrefix []string) error { + return sam.CreateGCSBucketMock(ctx, sc, bucketName, projectID, location, ttl, matchesPrefix) } -func (sam *StorageAccessorMock) CreateGCSBucketWithLifecycle(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { - return sam.CreateGCSBucketWithLifecycleMock(ctx, sc, bucketName, projectID, location, matchesPrefix, ttl) -} - -func (sam *StorageAccessorMock) EnableBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error { - return sam.EnableBucketLifecycleDeleteRuleMock(ctx, sc, bucketName, matchesPrefix, ttl) +func (sam *StorageAccessorMock) ApplyBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error { + return sam.ApplyBucketLifecycleDeleteRuleMock(ctx, sc, bucketName, matchesPrefix, ttl) } func (sam *StorageAccessorMock) UploadLocalFileToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error { diff --git a/accessors/storage/storage_accessor.go b/accessors/storage/storage_accessor.go index ee3202db22..ac391dd939 100644 --- a/accessors/storage/storage_accessor.go +++ b/accessors/storage/storage_accessor.go @@ -29,9 +29,8 @@ import ( ) type StorageAccessor interface { - CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string) error - CreateGCSBucketWithLifecycle(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error - EnableBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error + CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, ttl int64, matchesPrefix []string) error + ApplyBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error UploadLocalFileToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error WriteDataToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, data string) error ReadGcsFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) @@ -40,20 +39,9 @@ type StorageAccessor interface { type StorageAccessorImpl struct{} -func (sa *StorageAccessorImpl) CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string) error { - return sa.createGCSBucketUtil(ctx, sc, bucketName, projectID, location, nil, 0) -} - -func (sa *StorageAccessorImpl) CreateGCSBucketWithLifecycle(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { - return sa.createGCSBucketUtil(ctx, sc, bucketName, projectID, location, matchesPrefix, ttl) -} - -func (sa *StorageAccessorImpl) createGCSBucketUtil(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, matchesPrefix []string, ttl int64) error { - client, err := storageclient.GetOrCreateClient(ctx) - if err != nil { - return err - } - bucket := client.Bucket(bucketName) +// Create a GCS bucket with input parameters. Set @ttl to 0 to skip creating lifecycle rules. +func (sa *StorageAccessorImpl) CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, ttl int64, matchesPrefix []string) error { + bucket := sc.Bucket(bucketName) attrs := storage.BucketAttrs{ Location: location, } @@ -95,7 +83,7 @@ func (sa *StorageAccessorImpl) createGCSBucketUtil(ctx context.Context, sc stora // Applies the bucket lifecycle with delete rule. Only accepts the Age and // prefix rule conditions as it is only used for the Datastream destination // bucket currently. -func (sa *StorageAccessorImpl) EnableBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error { +func (sa *StorageAccessorImpl) ApplyBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error { for i, str := range matchesPrefix { matchesPrefix[i] = strings.TrimPrefix(str, "/") } @@ -135,6 +123,7 @@ func (sa *StorageAccessorImpl) UploadLocalFileToGCS(ctx context.Context, sc stor return sa.WriteDataToGCS(ctx, sc, filePath, fileName, string(data)) } +// Uploads a gcs object to gs://@filePath/@fileName with @data. func (sa *StorageAccessorImpl) WriteDataToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, data string) error { u, err := utils.ParseGCSFilePath(filePath) if err != nil { @@ -142,7 +131,11 @@ func (sa *StorageAccessorImpl) WriteDataToGCS(ctx context.Context, sc storagecli } bucketName := u.Host bucket := sc.Bucket(bucketName) - obj := bucket.Object(u.Path[1:] + fileName) + fullFilePath := u.Path + fileName + if strings.HasPrefix(fullFilePath, "/") { + fullFilePath = u.Path[1:] + fileName + } + obj := bucket.Object(fullFilePath) w := obj.NewWriter(ctx) logger.Log.Info(fmt.Sprintf("Writing data to %s", filePath)) @@ -160,6 +153,7 @@ func (sa *StorageAccessorImpl) WriteDataToGCS(ctx context.Context, sc storagecli return nil } +// Read a Gcs file path. func (sa *StorageAccessorImpl) ReadGcsFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { u, err := utils.ParseGCSFilePath(filePath) if err != nil { @@ -184,6 +178,7 @@ func (sa *StorageAccessorImpl) ReadGcsFile(ctx context.Context, sc storageclient return buf.String(), nil } +// Read local or gcs file path. Gcs files must start with a gs://. func (sa *StorageAccessorImpl) ReadAnyFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { if strings.HasPrefix(filePath, constants.GCS_FILE_PREFIX) { return sa.ReadGcsFile(ctx, sc, filePath) diff --git a/common/utils/storage_utils_test.go b/common/utils/storage_utils_test.go index c9ea7f75a5..fa00f0c302 100644 --- a/common/utils/storage_utils_test.go +++ b/common/utils/storage_utils_test.go @@ -59,6 +59,26 @@ func TestParseGCSFilePath(t *testing.T) { Path: "/path/to/folder/", }, }, + { + name: "Empty path", + filePath: "gs://test-bucket", + expectError: false, + want: &url.URL{ + Scheme: "gs", + Host: "test-bucket", + Path: "/", + }, + }, + { + name: "Empty path with leading slash", + filePath: "gs://test-bucket/", + expectError: false, + want: &url.URL{ + Scheme: "gs", + Host: "test-bucket", + Path: "/", + }, + }, { name: "Empty File path", filePath: "", @@ -81,7 +101,7 @@ func TestParseGCSFilePath(t *testing.T) { for _, tc := range testCases { got, err := ParseGCSFilePath(tc.filePath) - assert.Equal(t, tc.expectError, err != nil) - assert.Equal(t, tc.want, got) + assert.Equal(t, tc.expectError, err != nil, tc.name) + assert.Equal(t, tc.want, got, tc.name) } } diff --git a/conversion/conversion.go b/conversion/conversion.go index efbe86dec1..9c811d1c86 100644 --- a/conversion/conversion.go +++ b/conversion/conversion.go @@ -363,7 +363,7 @@ func dataFromDatabase(ctx context.Context, sourceProfile profiles.SourceProfile, } sa := storageaccessor.StorageAccessorImpl{} if gcsConfig.TtlInDaysSet { - err = sa.EnableBucketLifecycleDeleteRule(ctx, sc, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = sa.ApplyBucketLifecycleDeleteRule(ctx, sc, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") @@ -489,7 +489,7 @@ func dataFromDatabaseForDataflowMigration(targetProfile profiles.TargetProfile, } sa := storageaccessor.StorageAccessorImpl{} if gcsConfig.TtlInDaysSet { - err = sa.EnableBucketLifecycleDeleteRule(ctx, sc, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = sa.ApplyBucketLifecycleDeleteRule(ctx, sc, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") diff --git a/webv2/profile/profile.go b/webv2/profile/profile.go index 29b0958d43..1afb866241 100644 --- a/webv2/profile/profile.go +++ b/webv2/profile/profile.go @@ -168,7 +168,7 @@ func CreateConnectionProfile(w http.ResponseWriter, r *http.Request) { } else { bucketName = strings.ToLower(sessionState.Conv.Audit.MigrationRequestId) } - err = sa.CreateGCSBucket(ctx, sc, bucketName, sessionState.GCPProjectID, sessionState.Region) + err = sa.CreateGCSBucket(ctx, sc, bucketName, sessionState.GCPProjectID, sessionState.Region, 0, nil) if err != nil { http.Error(w, fmt.Sprintf("Error while creating bucket: %v", err), http.StatusBadRequest) return diff --git a/webv2/web.go b/webv2/web.go index 971d42c701..7ca1b03996 100644 --- a/webv2/web.go +++ b/webv2/web.go @@ -2492,7 +2492,7 @@ func writeSessionFile(ctx context.Context, sessionState *session.SessionState) e return err } sa := storageaccessor.StorageAccessorImpl{} - err = sa.CreateGCSBucket(ctx, sc, sessionState.Bucket, sessionState.GCPProjectID, sessionState.Region) + err = sa.CreateGCSBucket(ctx, sc, sessionState.Bucket, sessionState.GCPProjectID, sessionState.Region, 0, nil) if err != nil { return fmt.Errorf("error while creating bucket: %v", err) } From b017e1afb72e1675213e14064e94e178c7a3daa4 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 31 Jan 2024 01:14:47 +0530 Subject: [PATCH 18/25] Add spanner mocks and move out clients out of accessors --- .../clients/spanner/admin/admin_client.go | 64 +++++++++++++++++++ accessors/clients/spanner/admin/mocks.go | 14 ++++ .../clients/spanner/instanceadmin/mocks.go | 14 ++++ .../instanceadmin/spanner_instance_admin.go | 27 ++++++++ accessors/spanner/mocks.go | 59 +++++++++++++++++ accessors/spanner/spanner_accessor.go | 44 ++++--------- cmd/data.go | 7 +- conversion/conversion.go | 7 +- .../spanner/spanner_accessor_test.go | 7 +- webv2/helpers/helpers.go | 6 +- webv2/session/session_service.go | 7 +- 11 files changed, 215 insertions(+), 41 deletions(-) create mode 100644 accessors/clients/spanner/admin/mocks.go create mode 100644 accessors/clients/spanner/instanceadmin/mocks.go create mode 100644 accessors/spanner/mocks.go diff --git a/accessors/clients/spanner/admin/admin_client.go b/accessors/clients/spanner/admin/admin_client.go index 220621ee09..c5da8e6611 100644 --- a/accessors/clients/spanner/admin/admin_client.go +++ b/accessors/clients/spanner/admin/admin_client.go @@ -19,6 +19,8 @@ import ( "sync" database "cloud.google.com/go/spanner/admin/database/apiv1" + "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" + "github.com/googleapis/gax-go/v2" ) var once sync.Once @@ -41,3 +43,65 @@ func GetOrCreateClient(ctx context.Context) (*database.DatabaseAdminClient, erro } return spannerAdminClient, nil } + +type AdminClient interface { + GetDatabase(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) + CreateDatabase(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (CreateDatabaseOperation, error) + UpdateDatabaseDdl(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (UpdateDatabaseDdlOperation, error) +} + +type CreateDatabaseOperation interface { + Wait(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) +} + +type UpdateDatabaseDdlOperation interface { + Wait(ctx context.Context, opts ...gax.CallOption) error +} + +type AdminClientImpl struct { + adminClient *database.DatabaseAdminClient +} + +func NewAdminClientImpl(ctx context.Context) (*AdminClientImpl, error) { + c, err := GetOrCreateClient(ctx) + if err != nil { + return nil, err + } + return &AdminClientImpl{adminClient: c}, nil +} + +func (c *AdminClientImpl) GetDatabase(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) { + return c.adminClient.GetDatabase(ctx, req, opts...) +} + +func (c *AdminClientImpl) CreateDatabase(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (CreateDatabaseOperation, error) { + op, err := c.adminClient.CreateDatabase(ctx, req, opts...) + if err != nil { + return nil, err + } + return &CreateDatabaseOperationImpl{dbo: op}, nil +} + +func (c *AdminClientImpl) UpdateDatabaseDdl(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (UpdateDatabaseDdlOperation, error) { + op, err := c.adminClient.UpdateDatabaseDdl(ctx, req, opts...) + if err != nil { + return nil, err + } + return &UpdateDatabaseDdlImpl{dbo: op}, nil +} + +type CreateDatabaseOperationImpl struct { + dbo *database.CreateDatabaseOperation +} + +func (c *CreateDatabaseOperationImpl) Wait(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) { + return c.dbo.Wait(ctx, opts...) +} + +type UpdateDatabaseDdlImpl struct { + dbo *database.UpdateDatabaseDdlOperation +} + +func (c *UpdateDatabaseDdlImpl) Wait(ctx context.Context, opts ...gax.CallOption) error { + return c.dbo.Wait(ctx, opts...) +} diff --git a/accessors/clients/spanner/admin/mocks.go b/accessors/clients/spanner/admin/mocks.go new file mode 100644 index 0000000000..eafb726e1a --- /dev/null +++ b/accessors/clients/spanner/admin/mocks.go @@ -0,0 +1,14 @@ +// 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 diff --git a/accessors/clients/spanner/instanceadmin/mocks.go b/accessors/clients/spanner/instanceadmin/mocks.go new file mode 100644 index 0000000000..5bee8f5d97 --- /dev/null +++ b/accessors/clients/spanner/instanceadmin/mocks.go @@ -0,0 +1,14 @@ +// 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 spinstanceadmin diff --git a/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go b/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go index 8e6529bfc7..beae616a1a 100644 --- a/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go +++ b/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go @@ -19,6 +19,8 @@ import ( "sync" instance "cloud.google.com/go/spanner/admin/instance/apiv1" + "cloud.google.com/go/spanner/admin/instance/apiv1/instancepb" + "github.com/googleapis/gax-go/v2" ) var once sync.Once @@ -41,3 +43,28 @@ func GetOrCreateClient(ctx context.Context) (*instance.InstanceAdminClient, erro } return instanceAdminClient, nil } + +type InstanceAdminClient interface { + GetInstance(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) + GetInstanceConfig(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) +} + +type InstanceAdminClientImpl struct { + client *instance.InstanceAdminClient +} + +func NewInstanceAdminClientImpl(ctx context.Context) (*InstanceAdminClientImpl, error) { + c, err := GetOrCreateClient(ctx) + if err != nil { + return nil, err + } + return &InstanceAdminClientImpl{client: c}, nil +} + +func (c *InstanceAdminClientImpl) GetInstance(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) { + return c.client.GetInstance(ctx, req, opts...) +} + +func (c *InstanceAdminClientImpl) GetInstanceConfig(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) { + return c.client.GetInstanceConfig(ctx, req, opts...) +} diff --git a/accessors/spanner/mocks.go b/accessors/spanner/mocks.go new file mode 100644 index 0000000000..205bc84c4b --- /dev/null +++ b/accessors/spanner/mocks.go @@ -0,0 +1,59 @@ +// 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 spanneraccessor + +import ( + "context" + + spanneradmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/admin" + spinstanceadmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/instanceadmin" +) + +type SpannerAccessorMock struct { + GetDatabaseDialectMock func(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (string, error) + CheckExistingDbMock func(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (bool, error) + CreateEmptyDatabaseMock func(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) error + GetSpannerLeaderLocationMock func(ctx context.Context, instanceClient spinstanceadmin.InstanceAdminClient, instanceURI string) (string, error) + CheckIfChangeStreamExistsMock func(ctx context.Context, changeStreamName, dbURI string) (bool, error) + ValidateChangeStreamOptionsMock func(ctx context.Context, changeStreamName, dbURI string) error + CreateChangeStreamMock func(ctx context.Context, adminClient spanneradmin.AdminClient, changeStreamName, dbURI string) error +} + +func (sam *SpannerAccessorMock) GetDatabaseDialect(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (string, error) { + return sam.GetDatabaseDialectMock(ctx, adminClient, dbURI) +} + +func (sam *SpannerAccessorMock) CheckExistingDb(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (bool, error) { + return sam.CheckExistingDbMock(ctx, adminClient, dbURI) +} + +func (sam *SpannerAccessorMock) CreateEmptyDatabase(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) error { + return sam.CreateEmptyDatabaseMock(ctx, adminClient, dbURI) +} + +func (sam *SpannerAccessorMock) GetSpannerLeaderLocation(ctx context.Context, instanceClient spinstanceadmin.InstanceAdminClient, instanceURI string) (string, error) { + return sam.GetSpannerLeaderLocationMock(ctx, instanceClient, instanceURI) +} + +func (sam *SpannerAccessorMock) CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI string) (bool, error) { + return sam.CheckIfChangeStreamExistsMock(ctx, changeStreamName, dbURI) +} + +func (sam *SpannerAccessorMock) ValidateChangeStreamOptions(ctx context.Context, changeStreamName, dbURI string) error { + return sam.ValidateChangeStreamOptionsMock(ctx, changeStreamName, dbURI) +} + +func (sam *SpannerAccessorMock) CreateChangeStream(ctx context.Context, adminClient spanneradmin.AdminClient, changeStreamName, dbURI string) error { + return sam.CreateChangeStreamMock(ctx, adminClient, changeStreamName, dbURI) +} diff --git a/accessors/spanner/spanner_accessor.go b/accessors/spanner/spanner_accessor.go index a5972ae036..ebb1896427 100644 --- a/accessors/spanner/spanner_accessor.go +++ b/accessors/spanner/spanner_accessor.go @@ -30,28 +30,19 @@ import ( ) type SpannerAccessor interface { - GetDatabase(ctx context.Context, dbURI string) (*databasepb.Database, error) - GetDatabaseDialect(ctx context.Context, dbURI string) (string, error) - CheckExistingDb(ctx context.Context, dbURI string) (bool, error) - CreateEmptyDatabase(ctx context.Context, dbURI string) error - GetSpannerLeaderLocation(ctx context.Context, instanceURI string) (string, error) + GetDatabaseDialect(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (string, error) + CheckExistingDb(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (bool, error) + CreateEmptyDatabase(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) error + GetSpannerLeaderLocation(ctx context.Context, instanceClient spinstanceadmin.InstanceAdminClient, instanceURI string) (string, error) CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI string) (bool, error) ValidateChangeStreamOptions(ctx context.Context, changeStreamName, dbURI string) error - CreateChangeStream(ctx context.Context, changeStreamName, dbURI string) error + CreateChangeStream(ctx context.Context, adminClient spanneradmin.AdminClient, changeStreamName, dbURI string) error } type SpannerAccessorImpl struct{} -func (sp *SpannerAccessorImpl) GetDatabase(ctx context.Context, dbURI string) (*databasepb.Database, error) { - adminClient, err := spanneradmin.GetOrCreateClient(ctx) - if err != nil { - return nil, err - } - return adminClient.GetDatabase(ctx, &databasepb.GetDatabaseRequest{Name: dbURI}) -} - -func (sp *SpannerAccessorImpl) GetDatabaseDialect(ctx context.Context, dbURI string) (string, error) { - result, err := sp.GetDatabase(ctx, dbURI) +func (sp *SpannerAccessorImpl) GetDatabaseDialect(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (string, error) { + result, err := adminClient.GetDatabase(ctx, &databasepb.GetDatabaseRequest{Name: dbURI}) if err != nil { return "", fmt.Errorf("cannot connect to database: %v", err) } @@ -60,11 +51,11 @@ func (sp *SpannerAccessorImpl) GetDatabaseDialect(ctx context.Context, dbURI str // CheckExistingDb checks whether the database with dbURI exists or not. // If API call doesn't respond then user is informed after every 5 minutes on command line. -func (sp *SpannerAccessorImpl) CheckExistingDb(ctx context.Context, dbURI string) (bool, error) { +func (sp *SpannerAccessorImpl) CheckExistingDb(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (bool, error) { gotResponse := make(chan bool) var err error go func() { - _, err = sp.GetDatabase(ctx, dbURI) + _, err = adminClient.GetDatabase(ctx, &databasepb.GetDatabaseRequest{Name: dbURI}) gotResponse <- true }() for { @@ -83,11 +74,7 @@ func (sp *SpannerAccessorImpl) CheckExistingDb(ctx context.Context, dbURI string } } -func (sp *SpannerAccessorImpl) CreateEmptyDatabase(ctx context.Context, dbURI string) error { - adminClient, err := spanneradmin.GetOrCreateClient(ctx) - if err != nil { - return err - } +func (sp *SpannerAccessorImpl) CreateEmptyDatabase(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) error { project, instance, dbName := utils.ParseDbURI(dbURI) req := &databasepb.CreateDatabaseRequest{ Parent: fmt.Sprintf("projects/%s/instances/%s", project, instance), @@ -103,11 +90,7 @@ func (sp *SpannerAccessorImpl) CreateEmptyDatabase(ctx context.Context, dbURI st return nil } -func (sp *SpannerAccessorImpl) GetSpannerLeaderLocation(ctx context.Context, instanceURI string) (string, error) { - instanceClient, err := spinstanceadmin.GetOrCreateClient(ctx) - if err != nil { - return "", err - } +func (sp *SpannerAccessorImpl) GetSpannerLeaderLocation(ctx context.Context, instanceClient spinstanceadmin.InstanceAdminClient, instanceURI string) (string, error) { instanceInfo, err := instanceClient.GetInstance(ctx, &instancepb.GetInstanceRequest{Name: instanceURI}) if err != nil { return "", err @@ -192,9 +175,8 @@ func (sp *SpannerAccessorImpl) ValidateChangeStreamOptions(ctx context.Context, return nil } -func (sp *SpannerAccessorImpl) CreateChangeStream(ctx context.Context, changeStreamName, dbURI string) error { - spClient, _ := spanneradmin.GetOrCreateClient(ctx) - op, err := spClient.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{ +func (sp *SpannerAccessorImpl) CreateChangeStream(ctx context.Context, adminClient spanneradmin.AdminClient, changeStreamName, dbURI string) error { + op, err := adminClient.UpdateDatabaseDdl(ctx, &databasepb.UpdateDatabaseDdlRequest{ Database: dbURI, // TODO: create change stream for only the tables present in Spanner. Statements: []string{fmt.Sprintf("CREATE CHANGE STREAM %s FOR ALL OPTIONS (value_capture_type = 'NEW_ROW', retention_period = '7d')", changeStreamName)}, diff --git a/cmd/data.go b/cmd/data.go index b31f041269..9be772042e 100644 --- a/cmd/data.go +++ b/cmd/data.go @@ -26,6 +26,7 @@ import ( sp "cloud.google.com/go/spanner" database "cloud.google.com/go/spanner/admin/database/apiv1" + spanneradmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/admin" spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/utils" @@ -179,8 +180,12 @@ func (cmd *DataCmd) Execute(ctx context.Context, f *flag.FlagSet, _ ...interface // validateExistingDb validates that the existing spanner schema is in accordance with the one specified in the session file. func validateExistingDb(ctx context.Context, spDialect, dbURI string, adminClient *database.DatabaseAdminClient, client *sp.Client, conv *internal.Conv) error { + adminClientImpl, err := spanneradmin.NewAdminClientImpl(ctx) + if err != nil { + return err + } spA := spanneraccessor.SpannerAccessorImpl{} - dbExists, err := spA.CheckExistingDb(ctx, dbURI) + dbExists, err := spA.CheckExistingDb(ctx, adminClientImpl, dbURI) if err != nil { err = fmt.Errorf("can't verify target database: %v", err) return err diff --git a/conversion/conversion.go b/conversion/conversion.go index 9c811d1c86..06c6511809 100644 --- a/conversion/conversion.go +++ b/conversion/conversion.go @@ -44,6 +44,7 @@ import ( datastream "cloud.google.com/go/datastream/apiv1" sp "cloud.google.com/go/spanner" database "cloud.google.com/go/spanner/admin/database/apiv1" + spanneradmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/admin" storageclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/storage" spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" storageaccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/storage" @@ -811,8 +812,12 @@ func getSeekable(f *os.File) (*os.File, int64, error) { // VerifyDb checks whether the db exists and if it does, verifies if the schema is what we currently support. func VerifyDb(ctx context.Context, adminClient *database.DatabaseAdminClient, dbURI string) (dbExists bool, err error) { + adminClientImpl, err := spanneradmin.NewAdminClientImpl(ctx) + if err != nil { + return dbExists, err + } spA := spanneraccessor.SpannerAccessorImpl{} - dbExists, err = spA.CheckExistingDb(ctx, dbURI) + dbExists, err = spA.CheckExistingDb(ctx, adminClientImpl, dbURI) if err != nil { return dbExists, err } diff --git a/testing/accessors/spanner/spanner_accessor_test.go b/testing/accessors/spanner/spanner_accessor_test.go index 0dd0baa07f..311d2de0eb 100644 --- a/testing/accessors/spanner/spanner_accessor_test.go +++ b/testing/accessors/spanner/spanner_accessor_test.go @@ -28,6 +28,7 @@ import ( database "cloud.google.com/go/spanner/admin/database/apiv1" "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" + spanneradmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/admin" spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" "github.com/GoogleCloudPlatform/spanner-migration-tool/conversion" @@ -115,9 +116,13 @@ func TestCheckExistingDb(t *testing.T) { {"check-db-exists", true}, {"check-db-does-not-exist", false}, } + adminClientImpl, err := spanneradmin.NewAdminClientImpl(ctx) + if err != nil { + t.Fatal(err) + } spA := spanneraccessor.SpannerAccessorImpl{} for _, tc := range testCases { - dbExists, err := spA.CheckExistingDb(ctx, fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectID, instanceID, tc.dbName)) + dbExists, err := spA.CheckExistingDb(ctx, adminClientImpl, fmt.Sprintf("projects/%s/instances/%s/databases/%s", projectID, instanceID, tc.dbName)) assert.Nil(t, err) assert.Equal(t, tc.dbExists, dbExists) } diff --git a/webv2/helpers/helpers.go b/webv2/helpers/helpers.go index 0511070cb3..980bc7eac2 100644 --- a/webv2/helpers/helpers.go +++ b/webv2/helpers/helpers.go @@ -21,6 +21,7 @@ import ( "strings" database "cloud.google.com/go/spanner/admin/database/apiv1" + spanneradmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/admin" spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" "github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants" adminpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1" @@ -153,15 +154,14 @@ func CheckOrCreateMetadataDb(projectId string, instanceId string) bool { } ctx := context.Background() - adminClient, err := database.NewDatabaseAdminClient(ctx) + adminClientImpl, err := spanneradmin.NewAdminClientImpl(ctx) if err != nil { fmt.Println(err) return false } - defer adminClient.Close() spA := spanneraccessor.SpannerAccessorImpl{} - dbExists, err := spA.CheckExistingDb(ctx, uri) + dbExists, err := spA.CheckExistingDb(ctx, adminClientImpl, uri) if err != nil { fmt.Println(err) return false diff --git a/webv2/session/session_service.go b/webv2/session/session_service.go index 705e2284dc..7f2a089ba5 100644 --- a/webv2/session/session_service.go +++ b/webv2/session/session_service.go @@ -5,8 +5,8 @@ import ( "fmt" "cloud.google.com/go/spanner" - database "cloud.google.com/go/spanner/admin/database/apiv1" "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" + spanneradmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/admin" spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" helpers "github.com/GoogleCloudPlatform/spanner-migration-tool/webv2/helpers" ) @@ -80,16 +80,15 @@ func getOldMetadataDbUri(projectId string, instanceId string) string { func migrateMetadataDb(projectId, instanceId string) { ctx := context.Background() - adminClient, err := database.NewDatabaseAdminClient(ctx) + adminClientImpl, err := spanneradmin.NewAdminClientImpl(ctx) if err != nil { fmt.Println(err) return } - defer adminClient.Close() spA := spanneraccessor.SpannerAccessorImpl{} oldMetadataDbUri := getOldMetadataDbUri(projectId, instanceId) - oldMetadataDBExists, err := spA.CheckExistingDb(ctx, oldMetadataDbUri) + oldMetadataDBExists, err := spA.CheckExistingDb(ctx, adminClientImpl, oldMetadataDbUri) if err != nil { fmt.Printf("could not check if oldMetadataDB exists. error=%v\n", err) return From f0d8f15bf4c1bbf6ba840332846ffdd4d33dfce0 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 31 Jan 2024 14:23:56 +0530 Subject: [PATCH 19/25] add admin client mocks and few tests --- accessors/clients/spanner/admin/mocks.go | 41 +++++ accessors/spanner/spanner_accessor_test.go | 178 +++++++++++++++++++++ 2 files changed, 219 insertions(+) diff --git a/accessors/clients/spanner/admin/mocks.go b/accessors/clients/spanner/admin/mocks.go index eafb726e1a..e5792120d0 100644 --- a/accessors/clients/spanner/admin/mocks.go +++ b/accessors/clients/spanner/admin/mocks.go @@ -12,3 +12,44 @@ // See the License for the specific language governing permissions and // limitations under the License. package spanneradmin + +import ( + "context" + + "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" + "github.com/googleapis/gax-go/v2" +) + +type AdminClientMock struct { + GetDatabaseMock func(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) + CreateDatabaseMock func(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (CreateDatabaseOperation, error) + UpdateDatabaseDdlMock func(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (UpdateDatabaseDdlOperation, error) +} + +func (acm *AdminClientMock) GetDatabase(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) { + return acm.GetDatabaseMock(ctx, req, opts...) +} + +func (acm *AdminClientMock) CreateDatabase(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (CreateDatabaseOperation, error) { + return acm.CreateDatabaseMock(ctx, req, opts...) +} + +func (acm *AdminClientMock) UpdateDatabaseDdl(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (UpdateDatabaseDdlOperation, error) { + return acm.UpdateDatabaseDdlMock(ctx, req, opts...) +} + +type CreateDatabaseOperationMock struct { + WaitMock func(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) +} + +func (dbo *CreateDatabaseOperationMock) Wait(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) { + return dbo.WaitMock(ctx, opts...) +} + +type UpdateDatabaseDdlOperationMock struct { + WaitMock func(ctx context.Context, opts ...gax.CallOption) error +} + +func (dbo *UpdateDatabaseDdlOperationMock) Wait(ctx context.Context, opts ...gax.CallOption) error { + return dbo.WaitMock(ctx, opts...) +} diff --git a/accessors/spanner/spanner_accessor_test.go b/accessors/spanner/spanner_accessor_test.go index 5f0d2f60b9..7571bd3d26 100644 --- a/accessors/spanner/spanner_accessor_test.go +++ b/accessors/spanner/spanner_accessor_test.go @@ -12,3 +12,181 @@ // See the License for the specific language governing permissions and // limitations under the License. package spanneraccessor + +import ( + "context" + "fmt" + "os" + "testing" + + "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" + spanneradmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/admin" + "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" + "github.com/googleapis/gax-go/v2" + "github.com/stretchr/testify/assert" + "go.uber.org/zap" +) + +func init() { + logger.Log = zap.NewNop() +} + +func TestMain(m *testing.M) { + res := m.Run() + os.Exit(res) +} + +func TestSpannerAccessorImpl_GetDatabaseDialect(t *testing.T) { + testCases := []struct { + name string + acm spanneradmin.AdminClientMock + expectError bool + want string + }{ + { + name: "Basic", + acm: spanneradmin.AdminClientMock{ + GetDatabaseMock: func(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) { + return &databasepb.Database{DatabaseDialect: databasepb.DatabaseDialect_GOOGLE_STANDARD_SQL}, nil + }, + }, + expectError: false, + want: "google_standard_sql", + }, + { + name: "Pg Dialect", + acm: spanneradmin.AdminClientMock{ + GetDatabaseMock: func(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) { + return &databasepb.Database{DatabaseDialect: databasepb.DatabaseDialect_POSTGRESQL}, nil + }, + }, + expectError: false, + want: "postgresql", + }, + { + name: "Unspecified Dialect", + acm: spanneradmin.AdminClientMock{ + GetDatabaseMock: func(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) { + return &databasepb.Database{DatabaseDialect: databasepb.DatabaseDialect_DATABASE_DIALECT_UNSPECIFIED}, nil + }, + }, + expectError: false, + want: "database_dialect_unspecified", + }, + { + name: "Error case", + acm: spanneradmin.AdminClientMock{ + GetDatabaseMock: func(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) { + return nil, fmt.Errorf("test-error") + }, + }, + expectError: true, + want: "", + }, + } + ctx := context.Background() + spA := SpannerAccessorImpl{} + for _, tc := range testCases { + got, err := spA.GetDatabaseDialect(ctx, &tc.acm, "testUri") + assert.Equal(t, tc.expectError, err != nil, tc.name) + assert.Equal(t, tc.want, got, tc.name) + } +} + +func TestSpannerAccessorImpl_CreateEmptyDatabase(t *testing.T) { + testCases := []struct { + name string + acm spanneradmin.AdminClientMock + expectError bool + want string + }{ + { + name: "Basic", + acm: spanneradmin.AdminClientMock{ + CreateDatabaseMock: func(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (spanneradmin.CreateDatabaseOperation, error) { + return &spanneradmin.CreateDatabaseOperationMock{ + WaitMock: func(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) { return nil, nil }, + }, nil + }, + }, + expectError: false, + }, + { + name: "Create database returns error", + acm: spanneradmin.AdminClientMock{ + CreateDatabaseMock: func(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (spanneradmin.CreateDatabaseOperation, error) { + return nil, fmt.Errorf("test error") + }, + }, + expectError: true, + }, + { + name: "Wait returns error", + acm: spanneradmin.AdminClientMock{ + CreateDatabaseMock: func(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (spanneradmin.CreateDatabaseOperation, error) { + return &spanneradmin.CreateDatabaseOperationMock{ + WaitMock: func(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) { + return nil, fmt.Errorf("test error") + }, + }, nil + }, + }, + expectError: true, + }, + } + ctx := context.Background() + spA := SpannerAccessorImpl{} + for _, tc := range testCases { + err := spA.CreateEmptyDatabase(ctx, &tc.acm, "projects/test-project/instances/test-instance/databases/mydb") + assert.Equal(t, tc.expectError, err != nil, tc.name) + } +} + +func TestSpannerAccessorImpl_CreateChangeStream(t *testing.T) { + testCases := []struct { + name string + acm spanneradmin.AdminClientMock + expectError bool + want string + }{ + { + name: "Basic", + acm: spanneradmin.AdminClientMock{ + UpdateDatabaseDdlMock: func(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (spanneradmin.UpdateDatabaseDdlOperation, error) { + return &spanneradmin.UpdateDatabaseDdlOperationMock{ + WaitMock: func(ctx context.Context, opts ...gax.CallOption) error { return nil }, + }, nil + }, + }, + expectError: false, + }, + { + name: "Update database ddl returns error", + acm: spanneradmin.AdminClientMock{ + UpdateDatabaseDdlMock: func(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (spanneradmin.UpdateDatabaseDdlOperation, error) { + return nil, fmt.Errorf("test error") + }, + }, + expectError: true, + }, + { + name: "Wait returns error", + acm: spanneradmin.AdminClientMock{ + UpdateDatabaseDdlMock: func(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (spanneradmin.UpdateDatabaseDdlOperation, error) { + return &spanneradmin.UpdateDatabaseDdlOperationMock{ + WaitMock: func(ctx context.Context, opts ...gax.CallOption) error { + return fmt.Errorf("test error") + }, + }, nil + }, + }, + expectError: true, + }, + } + ctx := context.Background() + spA := SpannerAccessorImpl{} + for _, tc := range testCases { + err := spA.CreateChangeStream(ctx, &tc.acm, "my-changestream", "projects/test-project/instances/test-instance/databases/mydb") + assert.Equal(t, tc.expectError, err != nil, tc.name) + } +} From 9d55319b14de874c1354552df643258e9e967daf Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 31 Jan 2024 16:27:54 +0530 Subject: [PATCH 20/25] add instance accessor unit tests and mocks --- .../clients/spanner/instanceadmin/mocks.go | 20 +++ accessors/spanner/spanner_accessor_test.go | 144 ++++++++++++++++++ 2 files changed, 164 insertions(+) diff --git a/accessors/clients/spanner/instanceadmin/mocks.go b/accessors/clients/spanner/instanceadmin/mocks.go index 5bee8f5d97..b601c1a480 100644 --- a/accessors/clients/spanner/instanceadmin/mocks.go +++ b/accessors/clients/spanner/instanceadmin/mocks.go @@ -12,3 +12,23 @@ // See the License for the specific language governing permissions and // limitations under the License. package spinstanceadmin + +import ( + "context" + + "cloud.google.com/go/spanner/admin/instance/apiv1/instancepb" + "github.com/googleapis/gax-go/v2" +) + +type InstanceAdminClientMock struct { + GetInstanceMock func(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) + GetInstanceConfigMock func(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) +} + +func (iac *InstanceAdminClientMock) GetInstance(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) { + return iac.GetInstanceMock(ctx, req, opts...) +} + +func (iac *InstanceAdminClientMock) GetInstanceConfig(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) { + return iac.GetInstanceConfigMock(ctx, req, opts...) +} diff --git a/accessors/spanner/spanner_accessor_test.go b/accessors/spanner/spanner_accessor_test.go index 7571bd3d26..9f93050496 100644 --- a/accessors/spanner/spanner_accessor_test.go +++ b/accessors/spanner/spanner_accessor_test.go @@ -20,7 +20,9 @@ import ( "testing" "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" + "cloud.google.com/go/spanner/admin/instance/apiv1/instancepb" spanneradmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/admin" + spinstanceadmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/instanceadmin" "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" "github.com/googleapis/gax-go/v2" "github.com/stretchr/testify/assert" @@ -93,6 +95,53 @@ func TestSpannerAccessorImpl_GetDatabaseDialect(t *testing.T) { } } +func TestSpannerAccessorImpl_CheckExistingDb(t *testing.T) { + testCases := []struct { + name string + acm spanneradmin.AdminClientMock + expectError bool + want bool + }{ + { + name: "Basic", + acm: spanneradmin.AdminClientMock{ + GetDatabaseMock: func(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) { + return nil, nil + }, + }, + expectError: false, + want: true, + }, + { + name: "Database not found error", + acm: spanneradmin.AdminClientMock{ + GetDatabaseMock: func(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) { + return nil, fmt.Errorf("database not found") + }, + }, + expectError: false, + want: false, + }, + { + name: "Could not get db info", + acm: spanneradmin.AdminClientMock{ + GetDatabaseMock: func(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) { + return nil, fmt.Errorf("failed to connect") + }, + }, + expectError: true, + want: false, + }, + } + ctx := context.Background() + spA := SpannerAccessorImpl{} + for _, tc := range testCases { + got, err := spA.CheckExistingDb(ctx, &tc.acm, "testUri") + assert.Equal(t, tc.expectError, err != nil, tc.name) + assert.Equal(t, tc.want, got, tc.name) + } +} + func TestSpannerAccessorImpl_CreateEmptyDatabase(t *testing.T) { testCases := []struct { name string @@ -190,3 +239,98 @@ func TestSpannerAccessorImpl_CreateChangeStream(t *testing.T) { assert.Equal(t, tc.expectError, err != nil, tc.name) } } + +func TestSpannerAccessorImpl_GetSpannerLeaderLocation(t *testing.T) { + testCases := []struct { + name string + iac spinstanceadmin.InstanceAdminClientMock + expectError bool + want string + }{ + { + name: "Basic", + iac: spinstanceadmin.InstanceAdminClientMock{ + GetInstanceMock: func(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) { + return &instancepb.Instance{Config: "projects/test-project/instanceConfigs/test-config"}, nil + }, + GetInstanceConfigMock: func(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) { + return &instancepb.InstanceConfig{Replicas: []*instancepb.ReplicaInfo{ + &instancepb.ReplicaInfo{ + Location: "us-east1", + DefaultLeaderLocation: false, + }, + &instancepb.ReplicaInfo{ + Location: "india1", + DefaultLeaderLocation: true, + }, + &instancepb.ReplicaInfo{ + Location: "europe2", + DefaultLeaderLocation: false, + }, + }}, nil + }, + }, + expectError: false, + want: "india1", + }, + { + name: "GetInstanceMock returns error", + iac: spinstanceadmin.InstanceAdminClientMock{ + GetInstanceMock: func(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) { + return nil, fmt.Errorf("test-error") + }, + GetInstanceConfigMock: func(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) { + return nil, nil + }, + }, + expectError: true, + want: "", + }, + { + name: "GetInstanceConfigMock returns error", + iac: spinstanceadmin.InstanceAdminClientMock{ + GetInstanceMock: func(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) { + return &instancepb.Instance{Config: "projects/test-project/instanceConfigs/test-config"}, nil + }, + GetInstanceConfigMock: func(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) { + return nil, fmt.Errorf("test-error") + }, + }, + expectError: true, + want: "", + }, + { + name: "No leader found returns error", + iac: spinstanceadmin.InstanceAdminClientMock{ + GetInstanceMock: func(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) { + return &instancepb.Instance{Config: "projects/test-project/instanceConfigs/test-config"}, nil + }, + GetInstanceConfigMock: func(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) { + return &instancepb.InstanceConfig{Replicas: []*instancepb.ReplicaInfo{ + &instancepb.ReplicaInfo{ + Location: "us-east1", + DefaultLeaderLocation: false, + }, + &instancepb.ReplicaInfo{ + Location: "india1", + DefaultLeaderLocation: false, + }, + &instancepb.ReplicaInfo{ + Location: "europe2", + DefaultLeaderLocation: false, + }, + }}, nil + }, + }, + expectError: true, + want: "", + }, + } + ctx := context.Background() + spA := SpannerAccessorImpl{} + for _, tc := range testCases { + got, err := spA.GetSpannerLeaderLocation(ctx, &tc.iac, "projects/test-project/instances/test-instance") + assert.Equal(t, tc.expectError, err != nil, tc.name) + assert.Equal(t, tc.want, got, tc.name) + } +} From 4a52f6bae04d92c8f05c41fc86d9a063d7813ccb Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 31 Jan 2024 16:46:47 +0530 Subject: [PATCH 21/25] Rearrange files into mock.go and interface.go --- accessors/clients/dataflow/dataflow_client.go | 6 -- accessors/clients/dataflow/interface.go | 42 ++++++++++ accessors/clients/dataflow/mocks.go | 29 +++++++ .../clients/spanner/admin/admin_client.go | 64 -------------- accessors/clients/spanner/admin/interface.go | 84 +++++++++++++++++++ .../spanner/instanceadmin/interface.go | 47 +++++++++++ .../instanceadmin/spanner_instance_admin.go | 27 ------ accessors/clients/storage/interface.go | 80 ++++++++++++++++++ accessors/clients/storage/storage_client.go | 60 ------------- accessors/dataflow/dataflow_accessor_test.go | 17 ++-- 10 files changed, 287 insertions(+), 169 deletions(-) create mode 100644 accessors/clients/dataflow/interface.go create mode 100644 accessors/clients/dataflow/mocks.go create mode 100644 accessors/clients/spanner/admin/interface.go create mode 100644 accessors/clients/spanner/instanceadmin/interface.go create mode 100644 accessors/clients/storage/interface.go diff --git a/accessors/clients/dataflow/dataflow_client.go b/accessors/clients/dataflow/dataflow_client.go index cc76c34fda..94ac4584f7 100644 --- a/accessors/clients/dataflow/dataflow_client.go +++ b/accessors/clients/dataflow/dataflow_client.go @@ -19,14 +19,8 @@ import ( "sync" dataflow "cloud.google.com/go/dataflow/apiv1beta3" - "cloud.google.com/go/dataflow/apiv1beta3/dataflowpb" - "github.com/googleapis/gax-go/v2" ) -type DataflowClient interface { - LaunchFlexTemplate(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) -} - var once sync.Once var dfClient *dataflow.FlexTemplatesClient diff --git a/accessors/clients/dataflow/interface.go b/accessors/clients/dataflow/interface.go new file mode 100644 index 0000000000..dff61e482f --- /dev/null +++ b/accessors/clients/dataflow/interface.go @@ -0,0 +1,42 @@ +// 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" +) + +type DataflowClient interface { + LaunchFlexTemplate(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) +} + +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...) +} diff --git a/accessors/clients/dataflow/mocks.go b/accessors/clients/dataflow/mocks.go new file mode 100644 index 0000000000..1d71ca646c --- /dev/null +++ b/accessors/clients/dataflow/mocks.go @@ -0,0 +1,29 @@ +// 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" +) + +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...) +} diff --git a/accessors/clients/spanner/admin/admin_client.go b/accessors/clients/spanner/admin/admin_client.go index c5da8e6611..220621ee09 100644 --- a/accessors/clients/spanner/admin/admin_client.go +++ b/accessors/clients/spanner/admin/admin_client.go @@ -19,8 +19,6 @@ import ( "sync" database "cloud.google.com/go/spanner/admin/database/apiv1" - "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" - "github.com/googleapis/gax-go/v2" ) var once sync.Once @@ -43,65 +41,3 @@ func GetOrCreateClient(ctx context.Context) (*database.DatabaseAdminClient, erro } return spannerAdminClient, nil } - -type AdminClient interface { - GetDatabase(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) - CreateDatabase(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (CreateDatabaseOperation, error) - UpdateDatabaseDdl(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (UpdateDatabaseDdlOperation, error) -} - -type CreateDatabaseOperation interface { - Wait(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) -} - -type UpdateDatabaseDdlOperation interface { - Wait(ctx context.Context, opts ...gax.CallOption) error -} - -type AdminClientImpl struct { - adminClient *database.DatabaseAdminClient -} - -func NewAdminClientImpl(ctx context.Context) (*AdminClientImpl, error) { - c, err := GetOrCreateClient(ctx) - if err != nil { - return nil, err - } - return &AdminClientImpl{adminClient: c}, nil -} - -func (c *AdminClientImpl) GetDatabase(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) { - return c.adminClient.GetDatabase(ctx, req, opts...) -} - -func (c *AdminClientImpl) CreateDatabase(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (CreateDatabaseOperation, error) { - op, err := c.adminClient.CreateDatabase(ctx, req, opts...) - if err != nil { - return nil, err - } - return &CreateDatabaseOperationImpl{dbo: op}, nil -} - -func (c *AdminClientImpl) UpdateDatabaseDdl(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (UpdateDatabaseDdlOperation, error) { - op, err := c.adminClient.UpdateDatabaseDdl(ctx, req, opts...) - if err != nil { - return nil, err - } - return &UpdateDatabaseDdlImpl{dbo: op}, nil -} - -type CreateDatabaseOperationImpl struct { - dbo *database.CreateDatabaseOperation -} - -func (c *CreateDatabaseOperationImpl) Wait(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) { - return c.dbo.Wait(ctx, opts...) -} - -type UpdateDatabaseDdlImpl struct { - dbo *database.UpdateDatabaseDdlOperation -} - -func (c *UpdateDatabaseDdlImpl) Wait(ctx context.Context, opts ...gax.CallOption) error { - return c.dbo.Wait(ctx, opts...) -} diff --git a/accessors/clients/spanner/admin/interface.go b/accessors/clients/spanner/admin/interface.go new file mode 100644 index 0000000000..1ee5a3d450 --- /dev/null +++ b/accessors/clients/spanner/admin/interface.go @@ -0,0 +1,84 @@ +// 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" + + database "cloud.google.com/go/spanner/admin/database/apiv1" + "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" + "github.com/googleapis/gax-go/v2" +) + +type AdminClient interface { + GetDatabase(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) + CreateDatabase(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (CreateDatabaseOperation, error) + UpdateDatabaseDdl(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (UpdateDatabaseDdlOperation, error) +} + +type CreateDatabaseOperation interface { + Wait(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) +} + +type UpdateDatabaseDdlOperation interface { + Wait(ctx context.Context, opts ...gax.CallOption) error +} + +type AdminClientImpl struct { + adminClient *database.DatabaseAdminClient +} + +func NewAdminClientImpl(ctx context.Context) (*AdminClientImpl, error) { + c, err := GetOrCreateClient(ctx) + if err != nil { + return nil, err + } + return &AdminClientImpl{adminClient: c}, nil +} + +func (c *AdminClientImpl) GetDatabase(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) { + return c.adminClient.GetDatabase(ctx, req, opts...) +} + +func (c *AdminClientImpl) CreateDatabase(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (CreateDatabaseOperation, error) { + op, err := c.adminClient.CreateDatabase(ctx, req, opts...) + if err != nil { + return nil, err + } + return &CreateDatabaseOperationImpl{dbo: op}, nil +} + +func (c *AdminClientImpl) UpdateDatabaseDdl(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (UpdateDatabaseDdlOperation, error) { + op, err := c.adminClient.UpdateDatabaseDdl(ctx, req, opts...) + if err != nil { + return nil, err + } + return &UpdateDatabaseDdlImpl{dbo: op}, nil +} + +type CreateDatabaseOperationImpl struct { + dbo *database.CreateDatabaseOperation +} + +func (c *CreateDatabaseOperationImpl) Wait(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) { + return c.dbo.Wait(ctx, opts...) +} + +type UpdateDatabaseDdlImpl struct { + dbo *database.UpdateDatabaseDdlOperation +} + +func (c *UpdateDatabaseDdlImpl) Wait(ctx context.Context, opts ...gax.CallOption) error { + return c.dbo.Wait(ctx, opts...) +} diff --git a/accessors/clients/spanner/instanceadmin/interface.go b/accessors/clients/spanner/instanceadmin/interface.go new file mode 100644 index 0000000000..663d8ad114 --- /dev/null +++ b/accessors/clients/spanner/instanceadmin/interface.go @@ -0,0 +1,47 @@ +// 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 spinstanceadmin + +import ( + "context" + + instance "cloud.google.com/go/spanner/admin/instance/apiv1" + "cloud.google.com/go/spanner/admin/instance/apiv1/instancepb" + "github.com/googleapis/gax-go/v2" +) + +type InstanceAdminClient interface { + GetInstance(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) + GetInstanceConfig(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) +} + +type InstanceAdminClientImpl struct { + client *instance.InstanceAdminClient +} + +func NewInstanceAdminClientImpl(ctx context.Context) (*InstanceAdminClientImpl, error) { + c, err := GetOrCreateClient(ctx) + if err != nil { + return nil, err + } + return &InstanceAdminClientImpl{client: c}, nil +} + +func (c *InstanceAdminClientImpl) GetInstance(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) { + return c.client.GetInstance(ctx, req, opts...) +} + +func (c *InstanceAdminClientImpl) GetInstanceConfig(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) { + return c.client.GetInstanceConfig(ctx, req, opts...) +} diff --git a/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go b/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go index beae616a1a..8e6529bfc7 100644 --- a/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go +++ b/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go @@ -19,8 +19,6 @@ import ( "sync" instance "cloud.google.com/go/spanner/admin/instance/apiv1" - "cloud.google.com/go/spanner/admin/instance/apiv1/instancepb" - "github.com/googleapis/gax-go/v2" ) var once sync.Once @@ -43,28 +41,3 @@ func GetOrCreateClient(ctx context.Context) (*instance.InstanceAdminClient, erro } return instanceAdminClient, nil } - -type InstanceAdminClient interface { - GetInstance(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) - GetInstanceConfig(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) -} - -type InstanceAdminClientImpl struct { - client *instance.InstanceAdminClient -} - -func NewInstanceAdminClientImpl(ctx context.Context) (*InstanceAdminClientImpl, error) { - c, err := GetOrCreateClient(ctx) - if err != nil { - return nil, err - } - return &InstanceAdminClientImpl{client: c}, nil -} - -func (c *InstanceAdminClientImpl) GetInstance(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) { - return c.client.GetInstance(ctx, req, opts...) -} - -func (c *InstanceAdminClientImpl) GetInstanceConfig(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) { - return c.client.GetInstanceConfig(ctx, req, opts...) -} diff --git a/accessors/clients/storage/interface.go b/accessors/clients/storage/interface.go new file mode 100644 index 0000000000..6bc395ec97 --- /dev/null +++ b/accessors/clients/storage/interface.go @@ -0,0 +1,80 @@ +// 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 storageclient + +import ( + "context" + "io" + + "cloud.google.com/go/storage" +) + +type StorageClient interface { + Bucket(name string) BucketHandle +} + +type BucketHandle interface { + Create(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) + Update(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) + Object(name string) ObjectHandle +} + +type ObjectHandle interface { + NewWriter(ctx context.Context) io.WriteCloser + NewReader(ctx context.Context) (io.ReadCloser, error) +} + +type StorageClientImpl struct { + client *storage.Client +} + +func NewStorageClientImpl(ctx context.Context) (*StorageClientImpl, error) { + c, err := GetOrCreateClient(ctx) + if err != nil { + return nil, err + } + return &StorageClientImpl{client: c}, nil +} + +func (c *StorageClientImpl) Bucket(name string) BucketHandle { + return &BucketHandleImpl{bucketHandle: c.client.Bucket(name)} +} + +type BucketHandleImpl struct { + bucketHandle *storage.BucketHandle +} + +func (b *BucketHandleImpl) Create(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) { + return b.bucketHandle.Create(ctx, projectID, attrs) +} + +func (b *BucketHandleImpl) Update(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) { + return b.bucketHandle.Update(ctx, uattrs) +} + +func (b *BucketHandleImpl) Object(name string) ObjectHandle { + return &ObjectHandleImpl{objectHandle: b.bucketHandle.Object(name)} +} + +type ObjectHandleImpl struct { + objectHandle *storage.ObjectHandle +} + +func (o *ObjectHandleImpl) NewWriter(ctx context.Context) io.WriteCloser { + return o.objectHandle.NewWriter(ctx) +} + +func (o *ObjectHandleImpl) NewReader(ctx context.Context) (io.ReadCloser, error) { + return o.objectHandle.NewReader(ctx) +} diff --git a/accessors/clients/storage/storage_client.go b/accessors/clients/storage/storage_client.go index 4ce4026b83..28c62b868f 100644 --- a/accessors/clients/storage/storage_client.go +++ b/accessors/clients/storage/storage_client.go @@ -16,7 +16,6 @@ package storageclient import ( "context" "fmt" - "io" "sync" "cloud.google.com/go/storage" @@ -42,62 +41,3 @@ func GetOrCreateClient(ctx context.Context) (*storage.Client, error) { } return gcsClient, nil } - -type StorageClient interface { - Bucket(name string) BucketHandle -} - -type BucketHandle interface { - Create(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) - Update(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) - Object(name string) ObjectHandle -} - -type ObjectHandle interface { - NewWriter(ctx context.Context) io.WriteCloser - NewReader(ctx context.Context) (io.ReadCloser, error) -} - -type StorageClientImpl struct { - client *storage.Client -} - -func NewStorageClientImpl(ctx context.Context) (*StorageClientImpl, error) { - c, err := GetOrCreateClient(ctx) - if err != nil { - return nil, err - } - return &StorageClientImpl{client: c}, nil -} - -func (c *StorageClientImpl) Bucket(name string) BucketHandle { - return &BucketHandleImpl{bucketHandle: c.client.Bucket(name)} -} - -type BucketHandleImpl struct { - bucketHandle *storage.BucketHandle -} - -func (b *BucketHandleImpl) Create(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) { - return b.bucketHandle.Create(ctx, projectID, attrs) -} - -func (b *BucketHandleImpl) Update(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) { - return b.bucketHandle.Update(ctx, uattrs) -} - -func (b *BucketHandleImpl) Object(name string) ObjectHandle { - return &ObjectHandleImpl{objectHandle: b.bucketHandle.Object(name)} -} - -type ObjectHandleImpl struct { - objectHandle *storage.ObjectHandle -} - -func (o *ObjectHandleImpl) NewWriter(ctx context.Context) io.WriteCloser { - return o.objectHandle.NewWriter(ctx) -} - -func (o *ObjectHandleImpl) NewReader(ctx context.Context) (io.ReadCloser, error) { - return o.objectHandle.NewReader(ctx) -} diff --git a/accessors/dataflow/dataflow_accessor_test.go b/accessors/dataflow/dataflow_accessor_test.go index ffb8b6537b..76cec60007 100644 --- a/accessors/dataflow/dataflow_accessor_test.go +++ b/accessors/dataflow/dataflow_accessor_test.go @@ -20,6 +20,7 @@ import ( "testing" "cloud.google.com/go/dataflow/apiv1beta3/dataflowpb" + dataflowclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/dataflow" "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" "github.com/google/go-cmp/cmp" "github.com/googleapis/gax-go/v2" @@ -157,14 +158,6 @@ func getExpectedGcloudCmd2() string { "transformationContextFilePath=gs://transformationContext.json" } -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...) -} - func TestLaunchDataflowTemplate(t *testing.T) { ctx := context.Background() da := DataflowAccessorImpl{} @@ -172,7 +165,7 @@ func TestLaunchDataflowTemplate(t *testing.T) { name string params map[string]string cfg DataflowTuningConfig - dcm DataflowClientMock + dcm dataflowclient.DataflowClientMock expectError bool expectedJobId string expectedGcloudCmd string @@ -181,7 +174,7 @@ func TestLaunchDataflowTemplate(t *testing.T) { name: "Basic Correct", params: getParameters(), cfg: getTuningConfig(), - dcm: DataflowClientMock{ + dcm: dataflowclient.DataflowClientMock{ LaunchFlexTemplateMock: func(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) { return &dataflowpb.LaunchFlexTemplateResponse{Job: &dataflowpb.Job{Id: "1234"}}, nil }, @@ -194,7 +187,7 @@ func TestLaunchDataflowTemplate(t *testing.T) { name: "Request builder error", params: getParameters(), cfg: DataflowTuningConfig{Subnetwork: "test"}, - dcm: DataflowClientMock{ + dcm: dataflowclient.DataflowClientMock{ LaunchFlexTemplateMock: func(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) { return &dataflowpb.LaunchFlexTemplateResponse{Job: &dataflowpb.Job{Id: "1234"}}, nil }, @@ -207,7 +200,7 @@ func TestLaunchDataflowTemplate(t *testing.T) { name: "Launch flex template throws error", params: getParameters(), cfg: getTuningConfig(), - dcm: DataflowClientMock{ + dcm: dataflowclient.DataflowClientMock{ LaunchFlexTemplateMock: func(ctx context.Context, req *dataflowpb.LaunchFlexTemplateRequest, opts ...gax.CallOption) (*dataflowpb.LaunchFlexTemplateResponse, error) { return nil, fmt.Errorf("test error") }, From 7a89822d387294e663d18d9df5b11a870d12ffcf Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Wed, 31 Jan 2024 17:37:11 +0530 Subject: [PATCH 22/25] add unit tests and mock for storage --- accessors/clients/storage/mocks.go | 72 ++++ accessors/storage/storage_accessor.go | 7 +- accessors/storage/storage_accessor_test.go | 389 ++++++++++++++++++++- webv2/session/session_service.go | 8 +- 4 files changed, 462 insertions(+), 14 deletions(-) diff --git a/accessors/clients/storage/mocks.go b/accessors/clients/storage/mocks.go index 02bff2f3ac..b5f7b993f2 100644 --- a/accessors/clients/storage/mocks.go +++ b/accessors/clients/storage/mocks.go @@ -12,3 +12,75 @@ // See the License for the specific language governing permissions and // limitations under the License. package storageclient + +import ( + "context" + "io" + + "cloud.google.com/go/storage" +) + +type StorageClientMock struct { + BucketMock func(name string) BucketHandle +} + +func (scm *StorageClientMock) Bucket(name string) BucketHandle { + return scm.BucketMock(name) +} + +type BucketHandleMock struct { + CreateMock func(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) + UpdateMock func(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) + ObjectMock func(name string) ObjectHandle +} + +func (b *BucketHandleMock) Create(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) { + return b.CreateMock(ctx, projectID, attrs) +} + +func (b *BucketHandleMock) Update(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) { + return b.UpdateMock(ctx, uattrs) +} + +func (b *BucketHandleMock) Object(name string) ObjectHandle { + return b.ObjectMock(name) +} + +type ObjectHandleMock struct { + NewWriterMock func(ctx context.Context) io.WriteCloser + NewReaderMock func(ctx context.Context) (io.ReadCloser, error) +} + +func (o *ObjectHandleMock) NewWriter(ctx context.Context) io.WriteCloser { + return o.NewWriterMock(ctx) +} + +func (o *ObjectHandleMock) NewReader(ctx context.Context) (io.ReadCloser, error) { + return o.NewReaderMock(ctx) +} + +type WriterMock struct { + WriteMock func(p []byte) (n int, err error) + CloseMock func() error +} + +func (w *WriterMock) Write(p []byte) (n int, err error) { + return w.WriteMock(p) +} + +func (w *WriterMock) Close() error { + return w.CloseMock() +} + +type ReaderMock struct { + ReadMock func(p []byte) (n int, err error) + CloseMock func() error +} + +func (r *ReaderMock) Read(p []byte) (n int, err error) { + return r.ReadMock(p) +} + +func (r *ReaderMock) Close() error { + return r.CloseMock() +} diff --git a/accessors/storage/storage_accessor.go b/accessors/storage/storage_accessor.go index ac391dd939..9760d2a453 100644 --- a/accessors/storage/storage_accessor.go +++ b/accessors/storage/storage_accessor.go @@ -114,7 +114,7 @@ func (sa *StorageAccessorImpl) ApplyBucketLifecycleDeleteRule(ctx context.Contex return nil } -// UploadLocalFileToGCS uploads an object. +// UploadLocalFileToGCS uploads a local file at @localFilePath to a gcs file path @filePath with name @fileName. func (sa *StorageAccessorImpl) UploadLocalFileToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error { data, err := os.ReadFile(localFilePath) if err != nil { @@ -141,13 +141,13 @@ func (sa *StorageAccessorImpl) WriteDataToGCS(ctx context.Context, sc storagecli logger.Log.Info(fmt.Sprintf("Writing data to %s", filePath)) n, err := fmt.Fprint(w, data) if err != nil { - fmt.Printf("Failed to write to Cloud Storage: %s", filePath) + fmt.Printf("Failed to write to Cloud Storage: %s\n", filePath) return err } logger.Log.Info(fmt.Sprintf("Wrote %d bytes to GCS", n)) if err := w.Close(); err != nil { - fmt.Printf("Failed to close GCS file: %s", filePath) + fmt.Printf("Failed to close GCS file: %s\n", filePath) return err } return nil @@ -162,7 +162,6 @@ func (sa *StorageAccessorImpl) ReadGcsFile(ctx context.Context, sc storageclient bucketName := u.Host bucket := sc.Bucket(bucketName) obj := bucket.Object(u.Path[1:]) - rc, err := obj.NewReader(ctx) if err != nil { return "", err diff --git a/accessors/storage/storage_accessor_test.go b/accessors/storage/storage_accessor_test.go index 1519db317b..409d148a77 100644 --- a/accessors/storage/storage_accessor_test.go +++ b/accessors/storage/storage_accessor_test.go @@ -2,13 +2,384 @@ // // 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. +// You may obtain a package storageaccessor + +import ( + "context" + "fmt" + "io" + "os" + "strings" + "testing" + + "cloud.google.com/go/storage" + storageclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/storage" + "github.com/GoogleCloudPlatform/spanner-migration-tool/logger" + "github.com/stretchr/testify/assert" + "go.uber.org/zap" + "google.golang.org/api/googleapi" +) + +func init() { + logger.Log = zap.NewNop() +} + +func TestMain(m *testing.M) { + res := m.Run() + os.Exit(res) +} + +func TestStorageAccessorImpl_CreateGCSBucket(t *testing.T) { + testCases := []struct { + name string + scm storageclient.StorageClientMock + expectError bool + }{ + { + name: "Basic", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + CreateMock: func(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) { + return nil + }, + } + }, + }, + expectError: false, + }, + { + name: "random error", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + CreateMock: func(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) { + return fmt.Errorf("random error") + }, + } + }, + }, + expectError: true, + }, + { + name: "Bucket already exists", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + CreateMock: func(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) { + return &googleapi.Error{Code: 409} + }, + } + }, + }, + expectError: false, + }, + { + name: "Other google api error", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + CreateMock: func(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) { + return &googleapi.Error{Code: 100} + }, + } + }, + }, + expectError: true, + }, + } + ctx := context.Background() + sa := StorageAccessorImpl{} + for _, tc := range testCases { + err := sa.CreateGCSBucket(ctx, &tc.scm, "test-bucket", "test-project", "india2", 1, nil) + assert.Equal(t, tc.expectError, err != nil, tc.name) + } +} + +func TestStorageAccessorImpl_ApplyBucketLifecycleDeleteRule(t *testing.T) { + testCases := []struct { + name string + scm storageclient.StorageClientMock + ttl int64 + matchesPrefix []string + expectError bool + }{ + { + name: "Basic", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + UpdateMock: func(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) { + if strings.HasPrefix(uattrs.Lifecycle.Rules[0].Condition.MatchesPrefix[0], "/") { + return nil, fmt.Errorf("test error") + } + return &storage.BucketAttrs{Lifecycle: storage.Lifecycle{Rules: []storage.LifecycleRule{ + { + Action: storage.LifecycleAction{Type: "Delete"}, + Condition: storage.LifecycleCondition{ + AgeInDays: 5, + MatchesPrefix: []string{}, + }, + }, + }}}, nil + }, + } + }, + }, + ttl: 5, + matchesPrefix: []string{"test"}, + expectError: false, + }, + { + name: "Update error", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + UpdateMock: func(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) { + return nil, fmt.Errorf("test error") + }, + } + }, + }, + ttl: 5, + matchesPrefix: []string{"test"}, + expectError: true, + }, + { + name: "Prefix '/' gets removed", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + UpdateMock: func(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) { + if strings.HasPrefix(uattrs.Lifecycle.Rules[0].Condition.MatchesPrefix[0], "/") { + return nil, fmt.Errorf("test error") + } + return &storage.BucketAttrs{Lifecycle: storage.Lifecycle{Rules: []storage.LifecycleRule{ + { + Action: storage.LifecycleAction{Type: "Delete"}, + Condition: storage.LifecycleCondition{ + AgeInDays: 5, + MatchesPrefix: []string{}, + }, + }, + }}}, nil + }, + } + }, + }, + ttl: 5, + matchesPrefix: []string{"/test"}, + expectError: false, + }, + } + ctx := context.Background() + sa := StorageAccessorImpl{} + for _, tc := range testCases { + err := sa.ApplyBucketLifecycleDeleteRule(ctx, &tc.scm, "test-bucket", tc.matchesPrefix, tc.ttl) + assert.Equal(t, tc.expectError, err != nil, tc.name) + } +} + +func TestStorageAccessorImpl_WriteDataToGCS(t *testing.T) { + testCases := []struct { + name string + scm storageclient.StorageClientMock + filePath string + fileName string + data string + expectError bool + }{ + { + name: "Basic", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + ObjectMock: func(name string) storageclient.ObjectHandle { + return &storageclient.ObjectHandleMock{ + NewWriterMock: func(ctx context.Context) io.WriteCloser { + return &storageclient.WriterMock{ + WriteMock: func(p []byte) (n int, err error) { return len(p), nil }, + CloseMock: func() error { return nil }, + } + }, + } + }, + } + }, + }, + filePath: "gs://bucket/path", + fileName: "test-file", + data: "abcd", + expectError: false, + }, + { + name: "File parsing error", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + ObjectMock: func(name string) storageclient.ObjectHandle { + return &storageclient.ObjectHandleMock{ + NewWriterMock: func(ctx context.Context) io.WriteCloser { + return &storageclient.WriterMock{ + WriteMock: func(p []byte) (n int, err error) { return len(p), nil }, + CloseMock: func() error { return nil }, + } + }, + } + }, + } + }, + }, + filePath: "://bucket/path", + fileName: "test-file", + data: "abcd", + expectError: true, + }, + { + name: "Write error", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + ObjectMock: func(name string) storageclient.ObjectHandle { + return &storageclient.ObjectHandleMock{ + NewWriterMock: func(ctx context.Context) io.WriteCloser { + return &storageclient.WriterMock{ + WriteMock: func(p []byte) (n int, err error) { return 0, fmt.Errorf("test-error") }, + CloseMock: func() error { return nil }, + } + }, + } + }, + } + }, + }, + filePath: "gs://bucket/path", + fileName: "test-file", + data: "abcd", + expectError: true, + }, + { + name: "Close error", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + ObjectMock: func(name string) storageclient.ObjectHandle { + return &storageclient.ObjectHandleMock{ + NewWriterMock: func(ctx context.Context) io.WriteCloser { + return &storageclient.WriterMock{ + WriteMock: func(p []byte) (n int, err error) { return len(p), nil }, + CloseMock: func() error { return fmt.Errorf("test error") }, + } + }, + } + }, + } + }, + }, + filePath: "gs://bucket/path", + fileName: "test-file", + data: "abcd", + expectError: true, + }, + } + ctx := context.Background() + sa := StorageAccessorImpl{} + for _, tc := range testCases { + err := sa.WriteDataToGCS(ctx, &tc.scm, tc.filePath, tc.fileName, tc.data) + assert.Equal(t, tc.expectError, err != nil, tc.name) + } +} + +func TestStorageAccessorImpl_ReadGcsFile(t *testing.T) { + testCases := []struct { + name string + scm storageclient.StorageClientMock + filePath string + expectError bool + want string + }{ + { + name: "Basic", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + ObjectMock: func(name string) storageclient.ObjectHandle { + return &storageclient.ObjectHandleMock{ + NewReaderMock: func(ctx context.Context) (io.ReadCloser, error) { + return &storageclient.ReaderMock{ + ReadMock: func(p []byte) (n int, err error) { + copy(p, "hello") + return 5, io.EOF + }, + CloseMock: func() error { return nil }, + }, nil + }, + } + }, + } + }, + }, + filePath: "gs://bucket/path", + expectError: false, + want: "hello", + }, + { + name: "Parse error", + scm: storageclient.StorageClientMock{}, + filePath: "://bucket/path", + expectError: true, + want: "", + }, + { + name: "New reader error", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + ObjectMock: func(name string) storageclient.ObjectHandle { + return &storageclient.ObjectHandleMock{ + NewReaderMock: func(ctx context.Context) (io.ReadCloser, error) { + return nil, fmt.Errorf("test error") + }, + } + }, + } + }, + }, + filePath: "gs://bucket/path", + expectError: true, + want: "", + }, + { + name: "Read error", + scm: storageclient.StorageClientMock{ + BucketMock: func(name string) storageclient.BucketHandle { + return &storageclient.BucketHandleMock{ + ObjectMock: func(name string) storageclient.ObjectHandle { + return &storageclient.ObjectHandleMock{ + NewReaderMock: func(ctx context.Context) (io.ReadCloser, error) { + return &storageclient.ReaderMock{ + ReadMock: func(p []byte) (n int, err error) { + return 0, fmt.Errorf("test error") + }, + CloseMock: func() error { return nil }, + }, nil + }, + } + }, + } + }, + }, + filePath: "gs://bucket/path", + expectError: true, + want: "", + }, + } + ctx := context.Background() + sa := StorageAccessorImpl{} + for _, tc := range testCases { + got, err := sa.ReadGcsFile(ctx, &tc.scm, tc.filePath) + assert.Equal(t, tc.expectError, err != nil, tc.name) + assert.Equal(t, tc.want, got, tc.name) + } +} diff --git a/webv2/session/session_service.go b/webv2/session/session_service.go index 7f2a089ba5..c30e9d8a62 100644 --- a/webv2/session/session_service.go +++ b/webv2/session/session_service.go @@ -5,6 +5,7 @@ import ( "fmt" "cloud.google.com/go/spanner" + database "cloud.google.com/go/spanner/admin/database/apiv1" "cloud.google.com/go/spanner/admin/database/apiv1/databasepb" spanneradmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/admin" spanneraccessor "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/spanner" @@ -157,7 +158,12 @@ func migrateMetadataDb(projectId, instanceId string) { } fmt.Println("Successfully wrote data to new metadata DB.") - + adminClient, err := database.NewDatabaseAdminClient(ctx) + if err != nil { + fmt.Println(err) + return + } + defer adminClient.Close() err = adminClient.DropDatabase(ctx, &databasepb.DropDatabaseRequest{ Database: oldMetadataDbUri, }) From 1756228d6fbff6cb1c0601c538ce8e854a565f32 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Thu, 1 Feb 2024 14:54:49 +0530 Subject: [PATCH 23/25] add dataflow accessor mock --- accessors/dataflow/mocks.go | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100644 accessors/dataflow/mocks.go diff --git a/accessors/dataflow/mocks.go b/accessors/dataflow/mocks.go new file mode 100644 index 0000000000..c210809e31 --- /dev/null +++ b/accessors/dataflow/mocks.go @@ -0,0 +1,28 @@ +// 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 dataflowaccessor + +import ( + "context" + + dataflowclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/dataflow" +) + +type DataflowAccessorMock struct { + LaunchFlexTemplateMock func(ctx context.Context, c dataflowclient.DataflowClient, parameters map[string]string, cfg DataflowTuningConfig) (string, string, error) +} + +func (d *DataflowAccessorMock) LaunchDataflowTemplate(ctx context.Context, c dataflowclient.DataflowClient, parameters map[string]string, cfg DataflowTuningConfig) (string, string, error) { + return d.LaunchFlexTemplateMock(ctx, c, parameters, cfg) +} From aee36007988f5529fb0ea9334a9e2c8919705489 Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Thu, 1 Feb 2024 16:05:22 +0530 Subject: [PATCH 24/25] Add comments --- accessors/clients/dataflow/dataflow_client.go | 2 +- accessors/clients/dataflow/interface.go | 2 ++ accessors/clients/dataflow/mocks.go | 2 ++ .../clients/spanner/admin/admin_client.go | 2 +- accessors/clients/spanner/admin/interface.go | 6 ++++++ accessors/clients/spanner/admin/mocks.go | 6 ++++++ .../spanner/instanceadmin/interface.go | 2 ++ .../clients/spanner/instanceadmin/mocks.go | 2 ++ .../instanceadmin/spanner_instance_admin.go | 2 +- accessors/clients/storage/interface.go | 6 ++++++ accessors/clients/storage/mocks.go | 10 ++++++++++ accessors/clients/storage/storage_client.go | 2 +- accessors/dataflow/dataflow_accessor.go | 4 +++- accessors/dataflow/mocks.go | 2 ++ accessors/spanner/mocks.go | 2 ++ accessors/spanner/spanner_accessor.go | 14 ++++++++++++-- accessors/storage/mocks.go | 2 ++ accessors/storage/storage_accessor.go | 19 +++++++++++-------- 18 files changed, 72 insertions(+), 15 deletions(-) diff --git a/accessors/clients/dataflow/dataflow_client.go b/accessors/clients/dataflow/dataflow_client.go index 94ac4584f7..5f70698c8d 100644 --- a/accessors/clients/dataflow/dataflow_client.go +++ b/accessors/clients/dataflow/dataflow_client.go @@ -25,7 +25,7 @@ var once sync.Once var dfClient *dataflow.FlexTemplatesClient // This function is declared as a global variable to make it testable. The unit -// tests edit this function, acting like a double. +// tests update this function, acting like a double. var newFlexTemplatesClient = dataflow.NewFlexTemplatesClient func GetOrCreateClient(ctx context.Context) (*dataflow.FlexTemplatesClient, error) { diff --git a/accessors/clients/dataflow/interface.go b/accessors/clients/dataflow/interface.go index dff61e482f..9918b48a02 100644 --- a/accessors/clients/dataflow/interface.go +++ b/accessors/clients/dataflow/interface.go @@ -21,10 +21,12 @@ import ( "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 } diff --git a/accessors/clients/dataflow/mocks.go b/accessors/clients/dataflow/mocks.go index 1d71ca646c..ac899a7345 100644 --- a/accessors/clients/dataflow/mocks.go +++ b/accessors/clients/dataflow/mocks.go @@ -20,6 +20,8 @@ import ( "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) } diff --git a/accessors/clients/spanner/admin/admin_client.go b/accessors/clients/spanner/admin/admin_client.go index 220621ee09..6d60d8f194 100644 --- a/accessors/clients/spanner/admin/admin_client.go +++ b/accessors/clients/spanner/admin/admin_client.go @@ -25,7 +25,7 @@ var once sync.Once var spannerAdminClient *database.DatabaseAdminClient // This function is declared as a global variable to make it testable. The unit -// tests edit this function, acting like a double. +// tests update this function, acting like a double. var newDatabaseAdminClient = database.NewDatabaseAdminClient func GetOrCreateClient(ctx context.Context) (*database.DatabaseAdminClient, error) { diff --git a/accessors/clients/spanner/admin/interface.go b/accessors/clients/spanner/admin/interface.go index 1ee5a3d450..d0a2ab41b4 100644 --- a/accessors/clients/spanner/admin/interface.go +++ b/accessors/clients/spanner/admin/interface.go @@ -21,20 +21,24 @@ import ( "github.com/googleapis/gax-go/v2" ) +// Use this interface instead of database.DatabaseAdminClient to support mocking. type AdminClient interface { GetDatabase(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) CreateDatabase(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (CreateDatabaseOperation, error) UpdateDatabaseDdl(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest, opts ...gax.CallOption) (UpdateDatabaseDdlOperation, error) } +// Use this interface instead of database.CreateDatabaseOperation to support mocking. type CreateDatabaseOperation interface { Wait(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) } +// Use this interface instead of database.UpdateDatabaseDdlOperation to support mocking. type UpdateDatabaseDdlOperation interface { Wait(ctx context.Context, opts ...gax.CallOption) error } +// This implements the AdminClient interface. This is the primary implementation that should be used in all places other than tests. type AdminClientImpl struct { adminClient *database.DatabaseAdminClient } @@ -67,6 +71,7 @@ func (c *AdminClientImpl) UpdateDatabaseDdl(ctx context.Context, req *databasepb return &UpdateDatabaseDdlImpl{dbo: op}, nil } +// This implements the CreateDatabaseOperation interface. This is the primary implementation that should be used in all places other than tests. type CreateDatabaseOperationImpl struct { dbo *database.CreateDatabaseOperation } @@ -75,6 +80,7 @@ func (c *CreateDatabaseOperationImpl) Wait(ctx context.Context, opts ...gax.Call return c.dbo.Wait(ctx, opts...) } +// This implements the UpdateDatabaseDdl interface. This is the primary implementation that should be used in all places other than tests. type UpdateDatabaseDdlImpl struct { dbo *database.UpdateDatabaseDdlOperation } diff --git a/accessors/clients/spanner/admin/mocks.go b/accessors/clients/spanner/admin/mocks.go index e5792120d0..5d84ddeadd 100644 --- a/accessors/clients/spanner/admin/mocks.go +++ b/accessors/clients/spanner/admin/mocks.go @@ -20,6 +20,8 @@ import ( "github.com/googleapis/gax-go/v2" ) +// Mock that implements the AdminClient interface. +// Pass in unit tests where AdminClient is an input parameter. type AdminClientMock struct { GetDatabaseMock func(ctx context.Context, req *databasepb.GetDatabaseRequest, opts ...gax.CallOption) (*databasepb.Database, error) CreateDatabaseMock func(ctx context.Context, req *databasepb.CreateDatabaseRequest, opts ...gax.CallOption) (CreateDatabaseOperation, error) @@ -38,6 +40,8 @@ func (acm *AdminClientMock) UpdateDatabaseDdl(ctx context.Context, req *database return acm.UpdateDatabaseDdlMock(ctx, req, opts...) } +// Mock that implements the CreateDatabaseOperation interface. +// Pass in unit tests where CreateDatabaseOperation is an input parameter. type CreateDatabaseOperationMock struct { WaitMock func(ctx context.Context, opts ...gax.CallOption) (*databasepb.Database, error) } @@ -46,6 +50,8 @@ func (dbo *CreateDatabaseOperationMock) Wait(ctx context.Context, opts ...gax.Ca return dbo.WaitMock(ctx, opts...) } +// Mock that implements the UpdateDatabaseDdlOperation interface. +// Pass in unit tests where UpdateDatabaseDdlOperation is an input parameter. type UpdateDatabaseDdlOperationMock struct { WaitMock func(ctx context.Context, opts ...gax.CallOption) error } diff --git a/accessors/clients/spanner/instanceadmin/interface.go b/accessors/clients/spanner/instanceadmin/interface.go index 663d8ad114..5e195ebf1c 100644 --- a/accessors/clients/spanner/instanceadmin/interface.go +++ b/accessors/clients/spanner/instanceadmin/interface.go @@ -21,11 +21,13 @@ import ( "github.com/googleapis/gax-go/v2" ) +// Use this interface instead of instance.InstanceAdminClient to support mocking. type InstanceAdminClient interface { GetInstance(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) GetInstanceConfig(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) } +// This implements the InstanceAdminClient interface. This is the primary implementation that should be used in all places other than tests. type InstanceAdminClientImpl struct { client *instance.InstanceAdminClient } diff --git a/accessors/clients/spanner/instanceadmin/mocks.go b/accessors/clients/spanner/instanceadmin/mocks.go index b601c1a480..bb88cca00a 100644 --- a/accessors/clients/spanner/instanceadmin/mocks.go +++ b/accessors/clients/spanner/instanceadmin/mocks.go @@ -20,6 +20,8 @@ import ( "github.com/googleapis/gax-go/v2" ) +// Mock that implements the InstanceAdminClient interface. +// Pass in unit tests where InstanceAdminClient is an input parameter. type InstanceAdminClientMock struct { GetInstanceMock func(ctx context.Context, req *instancepb.GetInstanceRequest, opts ...gax.CallOption) (*instancepb.Instance, error) GetInstanceConfigMock func(ctx context.Context, req *instancepb.GetInstanceConfigRequest, opts ...gax.CallOption) (*instancepb.InstanceConfig, error) diff --git a/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go b/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go index 8e6529bfc7..a3d597e173 100644 --- a/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go +++ b/accessors/clients/spanner/instanceadmin/spanner_instance_admin.go @@ -25,7 +25,7 @@ var once sync.Once var instanceAdminClient *instance.InstanceAdminClient // This function is declared as a global variable to make it testable. The unit -// tests edit this function, acting like a double. +// tests update this function, acting like a double. var newInstanceAdminClient = instance.NewInstanceAdminClient func GetOrCreateClient(ctx context.Context) (*instance.InstanceAdminClient, error) { diff --git a/accessors/clients/storage/interface.go b/accessors/clients/storage/interface.go index 6bc395ec97..978acd841c 100644 --- a/accessors/clients/storage/interface.go +++ b/accessors/clients/storage/interface.go @@ -20,21 +20,25 @@ import ( "cloud.google.com/go/storage" ) +// Use this interface instead of storage.Client to support mocking. type StorageClient interface { Bucket(name string) BucketHandle } +// Use this interface instead of storage.BucketHandle to support mocking. type BucketHandle interface { Create(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) Update(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) Object(name string) ObjectHandle } +// Use this interface instead of storage.ObjectHandle to support mocking. type ObjectHandle interface { NewWriter(ctx context.Context) io.WriteCloser NewReader(ctx context.Context) (io.ReadCloser, error) } +// This implements the StorageClient interface. This is the primary implementation that should be used in all places other than tests. type StorageClientImpl struct { client *storage.Client } @@ -51,6 +55,7 @@ func (c *StorageClientImpl) Bucket(name string) BucketHandle { return &BucketHandleImpl{bucketHandle: c.client.Bucket(name)} } +// This implements the BucketHandle interface. This is the primary implementation that should be used in all places other than tests. type BucketHandleImpl struct { bucketHandle *storage.BucketHandle } @@ -67,6 +72,7 @@ func (b *BucketHandleImpl) Object(name string) ObjectHandle { return &ObjectHandleImpl{objectHandle: b.bucketHandle.Object(name)} } +// This implements the ObjectHandle interface. This is the primary implementation that should be used in all places other than tests. type ObjectHandleImpl struct { objectHandle *storage.ObjectHandle } diff --git a/accessors/clients/storage/mocks.go b/accessors/clients/storage/mocks.go index b5f7b993f2..f6effe2414 100644 --- a/accessors/clients/storage/mocks.go +++ b/accessors/clients/storage/mocks.go @@ -20,6 +20,8 @@ import ( "cloud.google.com/go/storage" ) +// Mock that implements the StorageClient interface. +// Pass in unit tests where StorageClient is an input parameter. type StorageClientMock struct { BucketMock func(name string) BucketHandle } @@ -28,6 +30,8 @@ func (scm *StorageClientMock) Bucket(name string) BucketHandle { return scm.BucketMock(name) } +// Mock that implements the BucketHandle interface. +// Pass in unit tests where BucketHandle is an input parameter. type BucketHandleMock struct { CreateMock func(ctx context.Context, projectID string, attrs *storage.BucketAttrs) (err error) UpdateMock func(ctx context.Context, uattrs storage.BucketAttrsToUpdate) (attrs *storage.BucketAttrs, err error) @@ -46,6 +50,8 @@ func (b *BucketHandleMock) Object(name string) ObjectHandle { return b.ObjectMock(name) } +// Mock that implements the ObjectHandle interface. +// Pass in unit tests where ObjectHandle is an input parameter. type ObjectHandleMock struct { NewWriterMock func(ctx context.Context) io.WriteCloser NewReaderMock func(ctx context.Context) (io.ReadCloser, error) @@ -59,6 +65,8 @@ func (o *ObjectHandleMock) NewReader(ctx context.Context) (io.ReadCloser, error) return o.NewReaderMock(ctx) } +// Mock that implements the io.WriteCloser interface. +// Pass in unit tests where io.WriteCloser is an input parameter. type WriterMock struct { WriteMock func(p []byte) (n int, err error) CloseMock func() error @@ -72,6 +80,8 @@ func (w *WriterMock) Close() error { return w.CloseMock() } +// Mock that implements the io.ReadCloser interface. +// Pass in unit tests where io.ReadCloser is an input parameter. type ReaderMock struct { ReadMock func(p []byte) (n int, err error) CloseMock func() error diff --git a/accessors/clients/storage/storage_client.go b/accessors/clients/storage/storage_client.go index 28c62b868f..09a66c401a 100644 --- a/accessors/clients/storage/storage_client.go +++ b/accessors/clients/storage/storage_client.go @@ -25,7 +25,7 @@ var once sync.Once var gcsClient *storage.Client // This function is declared as a global variable to make it testable. The unit -// tests edit this function, acting like a double. +// tests update this function, acting like a double. var newClient = storage.NewClient func GetOrCreateClient(ctx context.Context) (*storage.Client, error) { diff --git a/accessors/dataflow/dataflow_accessor.go b/accessors/dataflow/dataflow_accessor.go index aa7f1bea49..3fed9b3cc9 100644 --- a/accessors/dataflow/dataflow_accessor.go +++ b/accessors/dataflow/dataflow_accessor.go @@ -25,12 +25,14 @@ import ( "golang.org/x/exp/maps" ) +// The DataflowAccessor provides methods that internally use the dataflow client. Methods should only contain generic logic here that can be used by multiple workflows. type DataflowAccessor interface { // This function takes the template parameters (@parameters) and runtime environment config (@cfg) as input, and returns // the generated jobId, equivalentGcloudCommand and error if any. - LaunchFlexTemplate(ctx context.Context, c dataflowclient.DataflowClient, parameters map[string]string, cfg DataflowTuningConfig) (string, string, error) + LaunchDataflowTemplate(ctx context.Context, c dataflowclient.DataflowClient, parameters map[string]string, cfg DataflowTuningConfig) (string, string, error) } +// This implements the DataflowAccessor interface. This is the primary implementation that should be used in all places other than tests. type DataflowAccessorImpl struct{} func (dfA *DataflowAccessorImpl) LaunchDataflowTemplate(ctx context.Context, c dataflowclient.DataflowClient, parameters map[string]string, cfg DataflowTuningConfig) (string, string, error) { diff --git a/accessors/dataflow/mocks.go b/accessors/dataflow/mocks.go index c210809e31..9a65d9d512 100644 --- a/accessors/dataflow/mocks.go +++ b/accessors/dataflow/mocks.go @@ -19,6 +19,8 @@ import ( dataflowclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/dataflow" ) +// Mock that implements the DataflowAccessor interface. +// Pass in unit tests where DataflowAccessor is an input parameter. type DataflowAccessorMock struct { LaunchFlexTemplateMock func(ctx context.Context, c dataflowclient.DataflowClient, parameters map[string]string, cfg DataflowTuningConfig) (string, string, error) } diff --git a/accessors/spanner/mocks.go b/accessors/spanner/mocks.go index 205bc84c4b..828d2e17ce 100644 --- a/accessors/spanner/mocks.go +++ b/accessors/spanner/mocks.go @@ -20,6 +20,8 @@ import ( spinstanceadmin "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/spanner/instanceadmin" ) +// Mock that implements the SpannerAccessor interface. +// Pass in unit tests where SpannerAccessor is an input parameter. type SpannerAccessorMock struct { GetDatabaseDialectMock func(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (string, error) CheckExistingDbMock func(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (bool, error) diff --git a/accessors/spanner/spanner_accessor.go b/accessors/spanner/spanner_accessor.go index ebb1896427..fd5c2cae22 100644 --- a/accessors/spanner/spanner_accessor.go +++ b/accessors/spanner/spanner_accessor.go @@ -29,16 +29,27 @@ import ( "google.golang.org/api/iterator" ) +// The SpannerAccessor provides methods that internally use a spanner client (can be adminClient/databaseclient/instanceclient etc). +// Methods should only contain generic logic here that can be used by multiple workflows. type SpannerAccessor interface { + // Fetch the dialect of the spanner database. GetDatabaseDialect(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (string, error) + // CheckExistingDb checks whether the database with dbURI exists or not. + // If API call doesn't respond then user is informed after every 5 minutes on command line. CheckExistingDb(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (bool, error) + // Create a database with no schema. CreateEmptyDatabase(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) error + // Fetch the leader of the Spanner instance. GetSpannerLeaderLocation(ctx context.Context, instanceClient spinstanceadmin.InstanceAdminClient, instanceURI string) (string, error) + // Check if a change stream already exists. CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI string) (bool, error) + // Validate that change stream option 'VALUE_CAPTURE_TYPE' is 'NEW_ROW'. ValidateChangeStreamOptions(ctx context.Context, changeStreamName, dbURI string) error + // Create a change stream with default options. CreateChangeStream(ctx context.Context, adminClient spanneradmin.AdminClient, changeStreamName, dbURI string) error } +// This implements the SpannerAccessor interface. This is the primary implementation that should be used in all places other than tests. type SpannerAccessorImpl struct{} func (sp *SpannerAccessorImpl) GetDatabaseDialect(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (string, error) { @@ -49,8 +60,6 @@ func (sp *SpannerAccessorImpl) GetDatabaseDialect(ctx context.Context, adminClie return strings.ToLower(result.DatabaseDialect.String()), nil } -// CheckExistingDb checks whether the database with dbURI exists or not. -// If API call doesn't respond then user is informed after every 5 minutes on command line. func (sp *SpannerAccessorImpl) CheckExistingDb(ctx context.Context, adminClient spanneradmin.AdminClient, dbURI string) (bool, error) { gotResponse := make(chan bool) var err error @@ -108,6 +117,7 @@ func (sp *SpannerAccessorImpl) GetSpannerLeaderLocation(ctx context.Context, ins return "", fmt.Errorf("no leader found for spanner instance %s while trying fetch location", instanceURI) } +// Consider using a CreateChangestream operation and check for alreadyExists error. That uses adminClient which can be unit tested. func (sp *SpannerAccessorImpl) CheckIfChangeStreamExists(ctx context.Context, changeStreamName, dbURI string) (bool, error) { spClient, err := spannerclient.GetOrCreateClient(ctx, dbURI) if err != nil { diff --git a/accessors/storage/mocks.go b/accessors/storage/mocks.go index 0840a90737..2b790b0a5f 100644 --- a/accessors/storage/mocks.go +++ b/accessors/storage/mocks.go @@ -19,6 +19,8 @@ import ( storageclient "github.com/GoogleCloudPlatform/spanner-migration-tool/accessors/clients/storage" ) +// Mock that implements the StorageAccessor interface. +// Pass in unit tests where StorageAccessor is an input parameter. type StorageAccessorMock struct { CreateGCSBucketMock func(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, ttl int64, matchesPrefix []string) error ApplyBucketLifecycleDeleteRuleMock func(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error diff --git a/accessors/storage/storage_accessor.go b/accessors/storage/storage_accessor.go index 9760d2a453..e95a486128 100644 --- a/accessors/storage/storage_accessor.go +++ b/accessors/storage/storage_accessor.go @@ -28,18 +28,28 @@ import ( "google.golang.org/api/googleapi" ) +// The StorageAccessor provides methods that internally use a storage client. +// Methods should only contain generic logic here that can be used by multiple workflows. type StorageAccessor interface { + // Create a GCS bucket with the given name in the input projectId and location. If ttl is > 0, + // also apply a delete lifecycle rule with the input ttl and prefixes. Set @ttl to 0 to skip creating lifecycle rules. CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, ttl int64, matchesPrefix []string) error + // Applies the bucket lifecycle with delete rule. Only accepts the Age and prefix rule conditions as it is only used for the Datastream destination + // bucket currently. ApplyBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error + // UploadLocalFileToGCS uploads a local file at @localFilePath to a gcs file path @filePath with name @fileName. UploadLocalFileToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error + // Uploads a gcs object to gs://@filePath/@fileName with @data as content. WriteDataToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, data string) error + // Read a Gcs file path and returns the contents as a string. ReadGcsFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) + // Read a local or gcs file path. Files starting with a 'gs://' are treated as GCS files. ReadAnyFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) } +// This implements the StorageAccessor interface. This is the primary implementation that should be used in all places other than tests. type StorageAccessorImpl struct{} -// Create a GCS bucket with input parameters. Set @ttl to 0 to skip creating lifecycle rules. func (sa *StorageAccessorImpl) CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, ttl int64, matchesPrefix []string) error { bucket := sc.Bucket(bucketName) attrs := storage.BucketAttrs{ @@ -80,9 +90,6 @@ func (sa *StorageAccessorImpl) CreateGCSBucket(ctx context.Context, sc storagecl return nil } -// Applies the bucket lifecycle with delete rule. Only accepts the Age and -// prefix rule conditions as it is only used for the Datastream destination -// bucket currently. func (sa *StorageAccessorImpl) ApplyBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error { for i, str := range matchesPrefix { matchesPrefix[i] = strings.TrimPrefix(str, "/") @@ -114,7 +121,6 @@ func (sa *StorageAccessorImpl) ApplyBucketLifecycleDeleteRule(ctx context.Contex return nil } -// UploadLocalFileToGCS uploads a local file at @localFilePath to a gcs file path @filePath with name @fileName. func (sa *StorageAccessorImpl) UploadLocalFileToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error { data, err := os.ReadFile(localFilePath) if err != nil { @@ -123,7 +129,6 @@ func (sa *StorageAccessorImpl) UploadLocalFileToGCS(ctx context.Context, sc stor return sa.WriteDataToGCS(ctx, sc, filePath, fileName, string(data)) } -// Uploads a gcs object to gs://@filePath/@fileName with @data. func (sa *StorageAccessorImpl) WriteDataToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, data string) error { u, err := utils.ParseGCSFilePath(filePath) if err != nil { @@ -153,7 +158,6 @@ func (sa *StorageAccessorImpl) WriteDataToGCS(ctx context.Context, sc storagecli return nil } -// Read a Gcs file path. func (sa *StorageAccessorImpl) ReadGcsFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { u, err := utils.ParseGCSFilePath(filePath) if err != nil { @@ -177,7 +181,6 @@ func (sa *StorageAccessorImpl) ReadGcsFile(ctx context.Context, sc storageclient return buf.String(), nil } -// Read local or gcs file path. Gcs files must start with a gs://. func (sa *StorageAccessorImpl) ReadAnyFile(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) { if strings.HasPrefix(filePath, constants.GCS_FILE_PREFIX) { return sa.ReadGcsFile(ctx, sc, filePath) From 3585d91d9eab87c642fd67c73954e4e098b1603b Mon Sep 17 00:00:00 2001 From: Deep1998 Date: Thu, 1 Feb 2024 20:46:59 +0530 Subject: [PATCH 25/25] Move storage parameters into a struct --- accessors/storage/mocks.go | 12 ++++---- accessors/storage/storage_accessor.go | 36 +++++++++++----------- accessors/storage/storage_accessor_test.go | 14 +++++++-- accessors/storage/types.go | 24 +++++++++++++++ conversion/conversion.go | 12 ++++++-- webv2/profile/profile.go | 8 ++++- webv2/web.go | 8 ++++- 7 files changed, 84 insertions(+), 30 deletions(-) create mode 100644 accessors/storage/types.go diff --git a/accessors/storage/mocks.go b/accessors/storage/mocks.go index 2b790b0a5f..21673be905 100644 --- a/accessors/storage/mocks.go +++ b/accessors/storage/mocks.go @@ -22,20 +22,20 @@ import ( // Mock that implements the StorageAccessor interface. // Pass in unit tests where StorageAccessor is an input parameter. type StorageAccessorMock struct { - CreateGCSBucketMock func(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, ttl int64, matchesPrefix []string) error - ApplyBucketLifecycleDeleteRuleMock func(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error + CreateGCSBucketMock func(ctx context.Context, sc storageclient.StorageClient, req StorageBucketMetadata) error + ApplyBucketLifecycleDeleteRuleMock func(ctx context.Context, sc storageclient.StorageClient, req StorageBucketMetadata) error UploadLocalFileToGCSMock func(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error WriteDataToGCSMock func(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, data string) error ReadGcsFileMock func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) ReadAnyFileMock func(ctx context.Context, sc storageclient.StorageClient, filePath string) (string, error) } -func (sam *StorageAccessorMock) CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, ttl int64, matchesPrefix []string) error { - return sam.CreateGCSBucketMock(ctx, sc, bucketName, projectID, location, ttl, matchesPrefix) +func (sam *StorageAccessorMock) CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, req StorageBucketMetadata) error { + return sam.CreateGCSBucketMock(ctx, sc, req) } -func (sam *StorageAccessorMock) ApplyBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error { - return sam.ApplyBucketLifecycleDeleteRuleMock(ctx, sc, bucketName, matchesPrefix, ttl) +func (sam *StorageAccessorMock) ApplyBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, req StorageBucketMetadata) error { + return sam.ApplyBucketLifecycleDeleteRuleMock(ctx, sc, req) } func (sam *StorageAccessorMock) UploadLocalFileToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error { diff --git a/accessors/storage/storage_accessor.go b/accessors/storage/storage_accessor.go index e95a486128..42399648ce 100644 --- a/accessors/storage/storage_accessor.go +++ b/accessors/storage/storage_accessor.go @@ -33,10 +33,10 @@ import ( type StorageAccessor interface { // Create a GCS bucket with the given name in the input projectId and location. If ttl is > 0, // also apply a delete lifecycle rule with the input ttl and prefixes. Set @ttl to 0 to skip creating lifecycle rules. - CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, ttl int64, matchesPrefix []string) error + CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, req StorageBucketMetadata) error // Applies the bucket lifecycle with delete rule. Only accepts the Age and prefix rule conditions as it is only used for the Datastream destination // bucket currently. - ApplyBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error + ApplyBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, req StorageBucketMetadata) error // UploadLocalFileToGCS uploads a local file at @localFilePath to a gcs file path @filePath with name @fileName. UploadLocalFileToGCS(ctx context.Context, sc storageclient.StorageClient, filePath, fileName, localFilePath string) error // Uploads a gcs object to gs://@filePath/@fileName with @data as content. @@ -50,62 +50,62 @@ type StorageAccessor interface { // This implements the StorageAccessor interface. This is the primary implementation that should be used in all places other than tests. type StorageAccessorImpl struct{} -func (sa *StorageAccessorImpl) CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, bucketName, projectID, location string, ttl int64, matchesPrefix []string) error { - bucket := sc.Bucket(bucketName) +func (sa *StorageAccessorImpl) CreateGCSBucket(ctx context.Context, sc storageclient.StorageClient, req StorageBucketMetadata) error { + bucket := sc.Bucket(req.BucketName) attrs := storage.BucketAttrs{ - Location: location, + Location: req.Location, } - if ttl > 0 { + if req.Ttl > 0 { attrs.Lifecycle = storage.Lifecycle{ Rules: []storage.LifecycleRule{ { Action: storage.LifecycleAction{Type: "Delete"}, Condition: storage.LifecycleCondition{ - AgeInDays: ttl, + AgeInDays: req.Ttl, // The prefixes should not contain the bucket names and starting slash. // For object gs://my_bucket/pictures/paris_2022.jpg, // you would use a condition such as "matchesPrefix":["pictures/paris_"]. - MatchesPrefix: matchesPrefix, + MatchesPrefix: req.MatchesPrefix, }, }, }, } } - if err := bucket.Create(ctx, projectID, &attrs); err != nil { + if err := bucket.Create(ctx, req.ProjectID, &attrs); err != nil { if e, ok := err.(*googleapi.Error); ok { // Ignoring the bucket already exists error. if e.Code != 409 { return fmt.Errorf("failed to create bucket: %v", err) } else { - fmt.Printf("Using the existing bucket: %v \n", bucketName) + fmt.Printf("Using the existing bucket: %v \n", req.BucketName) } } else { return fmt.Errorf("failed to create bucket: %v", err) } } else { - logger.Log.Info(fmt.Sprintf("Created new GCS bucket: %v\n", bucketName)) + logger.Log.Info(fmt.Sprintf("Created new GCS bucket: %v\n", req.BucketName)) } return nil } -func (sa *StorageAccessorImpl) ApplyBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, bucketName string, matchesPrefix []string, ttl int64) error { - for i, str := range matchesPrefix { - matchesPrefix[i] = strings.TrimPrefix(str, "/") +func (sa *StorageAccessorImpl) ApplyBucketLifecycleDeleteRule(ctx context.Context, sc storageclient.StorageClient, req StorageBucketMetadata) error { + for i, str := range req.MatchesPrefix { + req.MatchesPrefix[i] = strings.TrimPrefix(str, "/") } - bucket := sc.Bucket(bucketName) + bucket := sc.Bucket(req.BucketName) bucketAttrsToUpdate := storage.BucketAttrsToUpdate{ Lifecycle: &storage.Lifecycle{ Rules: []storage.LifecycleRule{ { Action: storage.LifecycleAction{Type: "Delete"}, Condition: storage.LifecycleCondition{ - AgeInDays: ttl, + AgeInDays: req.Ttl, // The prefixes should not contain the bucket names and starting slash. // For object gs://my_bucket/pictures/paris_2022.jpg, // you would use a condition such as "matchesPrefix":["pictures/paris_"]. - MatchesPrefix: matchesPrefix, + MatchesPrefix: req.MatchesPrefix, }, }, }, @@ -117,7 +117,7 @@ func (sa *StorageAccessorImpl) ApplyBucketLifecycleDeleteRule(ctx context.Contex return fmt.Errorf("could not bucket with lifecycle: %w", err) } logger.Log.Info(fmt.Sprintf("Added lifecycle rule to bucket %v\n. Rule Action: %v\t Rule Condition: %v\n", - bucketName, attrs.Lifecycle.Rules[0].Action, attrs.Lifecycle.Rules[0].Condition)) + req.BucketName, attrs.Lifecycle.Rules[0].Action, attrs.Lifecycle.Rules[0].Condition)) return nil } diff --git a/accessors/storage/storage_accessor_test.go b/accessors/storage/storage_accessor_test.go index 409d148a77..4380a6a8f1 100644 --- a/accessors/storage/storage_accessor_test.go +++ b/accessors/storage/storage_accessor_test.go @@ -92,7 +92,13 @@ func TestStorageAccessorImpl_CreateGCSBucket(t *testing.T) { ctx := context.Background() sa := StorageAccessorImpl{} for _, tc := range testCases { - err := sa.CreateGCSBucket(ctx, &tc.scm, "test-bucket", "test-project", "india2", 1, nil) + err := sa.CreateGCSBucket(ctx, &tc.scm, StorageBucketMetadata{ + BucketName: "test-bucket", + ProjectID: "test-project", + Location: "india2", + Ttl: 1, + MatchesPrefix: nil, + }) assert.Equal(t, tc.expectError, err != nil, tc.name) } } @@ -176,7 +182,11 @@ func TestStorageAccessorImpl_ApplyBucketLifecycleDeleteRule(t *testing.T) { ctx := context.Background() sa := StorageAccessorImpl{} for _, tc := range testCases { - err := sa.ApplyBucketLifecycleDeleteRule(ctx, &tc.scm, "test-bucket", tc.matchesPrefix, tc.ttl) + err := sa.ApplyBucketLifecycleDeleteRule(ctx, &tc.scm, StorageBucketMetadata{ + BucketName: "test-bucket", + Ttl: tc.ttl, + MatchesPrefix: tc.matchesPrefix, + }) assert.Equal(t, tc.expectError, err != nil, tc.name) } } diff --git a/accessors/storage/types.go b/accessors/storage/types.go new file mode 100644 index 0000000000..5a20e1a002 --- /dev/null +++ b/accessors/storage/types.go @@ -0,0 +1,24 @@ +// 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 storageaccessor + +type StorageBucketMetadata struct { + BucketName string + // Not required for Updates. + ProjectID string + // Not required for Updates. + Location string + Ttl int64 + MatchesPrefix []string +} diff --git a/conversion/conversion.go b/conversion/conversion.go index 06c6511809..f1cf091dc4 100644 --- a/conversion/conversion.go +++ b/conversion/conversion.go @@ -364,7 +364,11 @@ func dataFromDatabase(ctx context.Context, sourceProfile profiles.SourceProfile, } sa := storageaccessor.StorageAccessorImpl{} if gcsConfig.TtlInDaysSet { - err = sa.ApplyBucketLifecycleDeleteRule(ctx, sc, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = sa.ApplyBucketLifecycleDeleteRule(ctx, sc, storageaccessor.StorageBucketMetadata{ + BucketName: gcsBucket, + Ttl: gcsConfig.TtlInDays, + MatchesPrefix: []string{gcsDestPrefix}, + }) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") @@ -490,7 +494,11 @@ func dataFromDatabaseForDataflowMigration(targetProfile profiles.TargetProfile, } sa := storageaccessor.StorageAccessorImpl{} if gcsConfig.TtlInDaysSet { - err = sa.ApplyBucketLifecycleDeleteRule(ctx, sc, gcsBucket, []string{gcsDestPrefix}, gcsConfig.TtlInDays) + err = sa.ApplyBucketLifecycleDeleteRule(ctx, sc, storageaccessor.StorageBucketMetadata{ + BucketName: gcsBucket, + Ttl: gcsConfig.TtlInDays, + MatchesPrefix: []string{gcsDestPrefix}, + }) if err != nil { logger.Log.Warn(fmt.Sprintf("\nWARNING: could not update Datastream destination GCS bucket with lifecycle rule, error: %v\n", err)) logger.Log.Warn("Please apply the lifecycle rule manually. Continuing...\n") diff --git a/webv2/profile/profile.go b/webv2/profile/profile.go index 1afb866241..f3935f5a29 100644 --- a/webv2/profile/profile.go +++ b/webv2/profile/profile.go @@ -168,7 +168,13 @@ func CreateConnectionProfile(w http.ResponseWriter, r *http.Request) { } else { bucketName = strings.ToLower(sessionState.Conv.Audit.MigrationRequestId) } - err = sa.CreateGCSBucket(ctx, sc, bucketName, sessionState.GCPProjectID, sessionState.Region, 0, nil) + err = sa.CreateGCSBucket(ctx, sc, storageaccessor.StorageBucketMetadata{ + BucketName: bucketName, + ProjectID: sessionState.GCPProjectID, + Location: sessionState.Region, + Ttl: 0, + MatchesPrefix: nil, + }) if err != nil { http.Error(w, fmt.Sprintf("Error while creating bucket: %v", err), http.StatusBadRequest) return diff --git a/webv2/web.go b/webv2/web.go index 7ca1b03996..22ed1a40f1 100644 --- a/webv2/web.go +++ b/webv2/web.go @@ -2492,7 +2492,13 @@ func writeSessionFile(ctx context.Context, sessionState *session.SessionState) e return err } sa := storageaccessor.StorageAccessorImpl{} - err = sa.CreateGCSBucket(ctx, sc, sessionState.Bucket, sessionState.GCPProjectID, sessionState.Region, 0, nil) + err = sa.CreateGCSBucket(ctx, sc, storageaccessor.StorageBucketMetadata{ + BucketName: sessionState.Bucket, + ProjectID: sessionState.GCPProjectID, + Location: sessionState.Region, + Ttl: 0, + MatchesPrefix: nil, + }) if err != nil { return fmt.Errorf("error while creating bucket: %v", err) }