diff --git a/ext/store/bigquery/bigquery.go b/ext/store/bigquery/bigquery.go index 33fad6d981..5c77f6255e 100644 --- a/ext/store/bigquery/bigquery.go +++ b/ext/store/bigquery/bigquery.go @@ -44,6 +44,7 @@ type Client interface { ExternalTableHandleFrom(dataset Dataset, name string) ResourceHandle ViewHandleFrom(dataset Dataset, name string) ResourceHandle RoutineHandleFrom(ds Dataset, name string) ResourceHandle + ModelHandleFrom(ds Dataset, name string) ResourceHandle BulkGetDDLView(ctx context.Context, dataset ProjectDataset, names []string) (map[ResourceURN]string, error) Close() error } @@ -291,6 +292,7 @@ func (s Store) Exist(ctx context.Context, tnnt tenant.Tenant, urn resource.URN) KindExternalTable: client.ExternalTableHandleFrom, KindView: client.ViewHandleFrom, KindRoutine: client.RoutineHandleFrom, + KindModel: client.ModelHandleFrom, } for kind, resourceHandleFn := range kindToHandleFn { diff --git a/ext/store/bigquery/bigquery_test.go b/ext/store/bigquery/bigquery_test.go index d4c6aa3051..db04b92a40 100644 --- a/ext/store/bigquery/bigquery_test.go +++ b/ext/store/bigquery/bigquery_test.go @@ -935,6 +935,7 @@ func TestBigqueryStore(t *testing.T) { client.On("ExternalTableHandleFrom", mock.Anything, mock.Anything).Return(handle).Maybe() client.On("ViewHandleFrom", mock.Anything, mock.Anything).Return(handle).Maybe() client.On("RoutineHandleFrom", mock.Anything, mock.Anything).Return(handle).Maybe() + client.On("ModelHandleFrom", mock.Anything, mock.Anything).Return(handle).Maybe() actualExist, actualError := bqStore.Exist(ctx, tnnt, urn) @@ -957,12 +958,14 @@ func TestBigqueryStore(t *testing.T) { externalTableHandle := new(mockTableResourceHandle) viewHandle := new(mockTableResourceHandle) routineHandle := new(mockTableResourceHandle) + modelHandle := new(mockTableResourceHandle) defer func() { dataSetHandle.AssertExpectations(t) tableHandle.AssertExpectations(t) externalTableHandle.AssertExpectations(t) viewHandle.AssertExpectations(t) routineHandle.AssertExpectations(t) + modelHandle.AssertExpectations(t) }() bqStore := bigquery.NewBigqueryDataStore(secretProvider, clientProvider) @@ -986,9 +989,10 @@ func TestBigqueryStore(t *testing.T) { externalTableHandle.On("Exists", mock.Anything).Return(false).Maybe() client.On("ViewHandleFrom", mock.Anything, mock.Anything).Return(viewHandle).Maybe() viewHandle.On("Exists", mock.Anything).Return(false).Maybe() - client.On("RoutineHandleFrom", mock.Anything, mock.Anything).Return(routineHandle) routineHandle.On("Exists", mock.Anything).Return(true) + client.On("ModelHandleFrom", mock.Anything, mock.Anything).Return(modelHandle).Maybe() + modelHandle.On("Exists", mock.Anything).Return(false).Maybe() actualExist, actualError := bqStore.Exist(ctx, tnnt, urn) @@ -996,7 +1000,7 @@ func TestBigqueryStore(t *testing.T) { assert.NoError(t, actualError) }) - t.Run("returns false and nil error if dataset exists and underlying routine does not exist", func(t *testing.T) { + t.Run("returns true and nil if dataset exists and underlying model does exist", func(t *testing.T) { secretProvider := new(mockSecretProvider) defer secretProvider.AssertExpectations(t) @@ -1011,12 +1015,14 @@ func TestBigqueryStore(t *testing.T) { externalTableHandle := new(mockTableResourceHandle) viewHandle := new(mockTableResourceHandle) routineHandle := new(mockTableResourceHandle) + modelHandle := new(mockTableResourceHandle) defer func() { dataSetHandle.AssertExpectations(t) tableHandle.AssertExpectations(t) externalTableHandle.AssertExpectations(t) viewHandle.AssertExpectations(t) routineHandle.AssertExpectations(t) + modelHandle.AssertExpectations(t) }() bqStore := bigquery.NewBigqueryDataStore(secretProvider, clientProvider) @@ -1040,9 +1046,67 @@ func TestBigqueryStore(t *testing.T) { externalTableHandle.On("Exists", mock.Anything).Return(false).Maybe() client.On("ViewHandleFrom", mock.Anything, mock.Anything).Return(viewHandle).Maybe() viewHandle.On("Exists", mock.Anything).Return(false).Maybe() + client.On("RoutineHandleFrom", mock.Anything, mock.Anything).Return(routineHandle).Maybe() + routineHandle.On("Exists", mock.Anything).Return(false).Maybe() + client.On("ModelHandleFrom", mock.Anything, mock.Anything).Return(modelHandle) + modelHandle.On("Exists", mock.Anything).Return(true) + + actualExist, actualError := bqStore.Exist(ctx, tnnt, urn) + + assert.True(t, actualExist) + assert.NoError(t, actualError) + }) + + t.Run("returns false and nil error if dataset exists and underlying routine does not exist", func(t *testing.T) { + secretProvider := new(mockSecretProvider) + defer secretProvider.AssertExpectations(t) + + client := new(mockClient) + defer client.AssertExpectations(t) + + clientProvider := new(mockClientProvider) + defer clientProvider.AssertExpectations(t) + + dataSetHandle := new(mockTableResourceHandle) + tableHandle := new(mockTableResourceHandle) + externalTableHandle := new(mockTableResourceHandle) + viewHandle := new(mockTableResourceHandle) + routineHandle := new(mockTableResourceHandle) + modelHandle := new(mockTableResourceHandle) + defer func() { + dataSetHandle.AssertExpectations(t) + tableHandle.AssertExpectations(t) + externalTableHandle.AssertExpectations(t) + viewHandle.AssertExpectations(t) + routineHandle.AssertExpectations(t) + modelHandle.AssertExpectations(t) + }() + + bqStore := bigquery.NewBigqueryDataStore(secretProvider, clientProvider) + + urn, err := resource.NewURN("bigquery", "project.dataset.routine") + assert.NoError(t, err) + + pts, _ := tenant.NewPlainTextSecret("secret_name", "secret_value") + secretProvider.On("GetSecret", mock.Anything, tnnt, "DATASTORE_BIGQUERY").Return(pts, nil) + + clientProvider.On("Get", mock.Anything, pts.Value()).Return(client, nil) + + client.On("Close").Return(nil) + + client.On("DatasetHandleFrom", mock.Anything, mock.Anything).Return(dataSetHandle) + dataSetHandle.On("Exists", mock.Anything).Return(true) + client.On("TableHandleFrom", mock.Anything, mock.Anything).Return(tableHandle).Maybe() + tableHandle.On("Exists", mock.Anything).Return(false).Maybe() + client.On("ExternalTableHandleFrom", mock.Anything, mock.Anything).Return(externalTableHandle).Maybe() + externalTableHandle.On("Exists", mock.Anything).Return(false).Maybe() + client.On("ViewHandleFrom", mock.Anything, mock.Anything).Return(viewHandle).Maybe() + viewHandle.On("Exists", mock.Anything).Return(false).Maybe() client.On("RoutineHandleFrom", mock.Anything, mock.Anything).Return(routineHandle) routineHandle.On("Exists", mock.Anything).Return(false) + client.On("ModelHandleFrom", mock.Anything, mock.Anything).Return(modelHandle).Maybe() + modelHandle.On("Exists", mock.Anything).Return(false).Maybe() actualExist, actualError := bqStore.Exist(ctx, tnnt, urn) @@ -1079,13 +1143,14 @@ func TestBigqueryStore(t *testing.T) { client.On("Close").Return(nil) handle.On("Exists", mock.Anything).Return(true).Once() - handle.On("Exists", mock.Anything).Return(false).Times(4) + handle.On("Exists", mock.Anything).Return(false).Times(5) client.On("DatasetHandleFrom", mock.Anything, mock.Anything).Return(handle) client.On("TableHandleFrom", mock.Anything, mock.Anything).Return(handle) client.On("ExternalTableHandleFrom", mock.Anything, mock.Anything).Return(handle) client.On("ViewHandleFrom", mock.Anything, mock.Anything).Return(handle) client.On("RoutineHandleFrom", mock.Anything, mock.Anything).Return(handle) + client.On("ModelHandleFrom", mock.Anything, mock.Anything).Return(handle) actualExist, actualError := bqStore.Exist(ctx, tnnt, urn) @@ -1144,6 +1209,11 @@ func (m *mockClient) RoutineHandleFrom(ds bigquery.Dataset, name string) bigquer return args.Get(0).(bigquery.ResourceHandle) } +func (m *mockClient) ModelHandleFrom(ds bigquery.Dataset, name string) bigquery.ResourceHandle { + args := m.Called(ds, name) + return args.Get(0).(bigquery.ResourceHandle) +} + func (m *mockClient) ViewHandleFrom(ds bigquery.Dataset, name string) bigquery.ResourceHandle { args := m.Called(ds, name) return args.Get(0).(bigquery.ResourceHandle) diff --git a/ext/store/bigquery/client.go b/ext/store/bigquery/client.go index d827aa5790..d430c70c0b 100644 --- a/ext/store/bigquery/client.go +++ b/ext/store/bigquery/client.go @@ -61,6 +61,11 @@ func (c *BqClient) ExternalTableHandleFrom(ds Dataset, name string) ResourceHand return NewExternalTableHandle(t) } +func (c *BqClient) ModelHandleFrom(ds Dataset, name string) ResourceHandle { + t := c.DatasetInProject(ds.Project, ds.DatasetName).Model(name) + return NewModelHandle(t) +} + func (c *BqClient) BulkGetDDLView(ctx context.Context, pd ProjectDataset, names []string) (map[ResourceURN]string, error) { me := errors.NewMultiError("bulk get ddl view errors") urnToDDL := make(map[ResourceURN]string, len(names)) diff --git a/ext/store/bigquery/model.go b/ext/store/bigquery/model.go new file mode 100644 index 0000000000..81c1fd654a --- /dev/null +++ b/ext/store/bigquery/model.go @@ -0,0 +1,41 @@ +package bigquery + +import ( + "context" + + "cloud.google.com/go/bigquery" + + "github.com/goto/optimus/core/resource" + "github.com/goto/optimus/internal/errors" +) + +const EntityModel = "resource_model" + +// BqModel is BigQuery Model +type BqModel interface { + Metadata(ctx context.Context) (mm *bigquery.ModelMetadata, err error) +} + +type ModelHandle struct { + bqModel BqModel +} + +func (ModelHandle) Create(_ context.Context, _ *resource.Resource) error { + return errors.FailedPrecondition(EntityModel, "create is not supported") +} + +func (ModelHandle) Update(_ context.Context, _ *resource.Resource) error { + return errors.FailedPrecondition(EntityModel, "update is not supported") +} + +func (r ModelHandle) Exists(ctx context.Context) bool { + if r.bqModel == nil { + return false + } + _, err := r.bqModel.Metadata(ctx) + return err == nil +} + +func NewModelHandle(bq BqModel) *ModelHandle { + return &ModelHandle{bqModel: bq} +} diff --git a/ext/store/bigquery/model_test.go b/ext/store/bigquery/model_test.go new file mode 100644 index 0000000000..a5862e2790 --- /dev/null +++ b/ext/store/bigquery/model_test.go @@ -0,0 +1,141 @@ +package bigquery_test + +import ( + "context" + "errors" + "testing" + + "cloud.google.com/go/bigquery" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + + "github.com/goto/optimus/core/resource" + "github.com/goto/optimus/core/tenant" + storebigquery "github.com/goto/optimus/ext/store/bigquery" +) + +func TestModelHandle(t *testing.T) { + ctx := context.Background() + bqStore := resource.Bigquery + tnnt, _ := tenant.NewTenant("proj", "ns") + metadata := resource.Metadata{ + Version: 1, + Description: "resource description", + Labels: map[string]string{"owner": "optimus"}, + } + spec := map[string]any{"description": []string{"a", "b"}} + res, err := resource.NewResource("proj.dataset.view1", storebigquery.KindView, bqStore, tnnt, &metadata, spec) + assert.Nil(t, err) + + t.Run("Create", func(t *testing.T) { + t.Run("return error, not supported", func(t *testing.T) { + v := NewMockBigQueryModel(t) + defer v.AssertExpectations(t) + handle := storebigquery.NewModelHandle(v) + + err := handle.Create(ctx, res) + assert.EqualError(t, err, "failed precondition for entity resource_model: create is not supported") + }) + }) + + t.Run("Update", func(t *testing.T) { + t.Run("return error, not supported", func(t *testing.T) { + v := NewMockBigQueryModel(t) + defer v.AssertExpectations(t) + handle := storebigquery.NewModelHandle(v) + + err := handle.Update(ctx, res) + assert.EqualError(t, err, "failed precondition for entity resource_model: update is not supported") + }) + }) + + t.Run("Exists", func(t *testing.T) { + t.Run("return true, model exists", func(t *testing.T) { + v := NewMockBigQueryModel(t) + defer v.AssertExpectations(t) + handle := storebigquery.NewModelHandle(v) + + v.On("Metadata", ctx).Return(nil, nil) + + actual := handle.Exists(ctx) + assert.True(t, actual) + }) + + t.Run("return false, model not exists", func(t *testing.T) { + v := NewMockBigQueryModel(t) + defer v.AssertExpectations(t) + handle := storebigquery.NewModelHandle(v) + + v.On("Metadata", ctx).Return(nil, errors.New("some error")) + + actual := handle.Exists(ctx) + assert.False(t, actual) + }) + + t.Run("return false, connection error", func(t *testing.T) { + v := NewMockBigQueryModel(t) + defer v.AssertExpectations(t) + handle := storebigquery.NewModelHandle(v) + + v.On("Metadata", ctx).Return(nil, context.DeadlineExceeded) + + actual := handle.Exists(ctx) + assert.False(t, actual) + }) + + t.Run("return false, BqModel is nil", func(t *testing.T) { + v := NewMockBigQueryModel(t) + defer v.AssertExpectations(t) + handle := storebigquery.NewModelHandle(nil) + + actual := handle.Exists(ctx) + assert.False(t, actual) + }) + }) +} + +type mockBigQueryModel struct { + mock.Mock +} + +func (_m *mockBigQueryModel) Metadata(ctx context.Context) (*bigquery.ModelMetadata, error) { + ret := _m.Called(ctx) + + if len(ret) == 0 { + panic("no return value specified for Metadata") + } + + var r0 *bigquery.ModelMetadata + var r1 error + if rf, ok := ret.Get(0).(func(context.Context) (*bigquery.ModelMetadata, error)); ok { + return rf(ctx) + } + if rf, ok := ret.Get(0).(func(context.Context) *bigquery.ModelMetadata); ok { + r0 = rf(ctx) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*bigquery.ModelMetadata) + } + } + + if rf, ok := ret.Get(1).(func(context.Context) error); ok { + r1 = rf(ctx) + } else { + r1 = ret.Error(1) + } + + return r0, r1 +} + +func NewMockBigQueryModel(t interface { + mock.TestingT + Cleanup(func()) +}, +) *mockBigQueryModel { + mock := &mockBigQueryModel{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} diff --git a/ext/store/bigquery/schema.go b/ext/store/bigquery/schema.go index 5bddd200d4..9d4af1414b 100644 --- a/ext/store/bigquery/schema.go +++ b/ext/store/bigquery/schema.go @@ -20,6 +20,7 @@ const ( KindView string = "view" KindExternalTable string = "external_table" KindRoutine string = "routine" + KindModel string = "model" ) type Schema []Field