Skip to content

Commit 751b3a0

Browse files
sudo-suhassudo-suhas
authored
feat: accept optional field ID in create asset probe (#35)
- Accept an optional ID field for create asset probe. ID, if specified, needs to be a valid UUID. The probe is created with the given ID if specified in the request and auto-generated otherwise. If the ID already exists, an appropriate error is returned. - Bump proton commit for generating protos.
1 parent 97f95b0 commit 751b3a0

File tree

9 files changed

+143
-16
lines changed

9 files changed

+143
-16
lines changed

‎Makefile

+1-2
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
NAME="github.com/goto/compass"
22
VERSION=$(shell git describe --always --tags 2>/dev/null)
33
COVERFILE="/tmp/compass.coverprofile"
4-
PROTON_COMMIT := "374259a5277ce724ecb0cda5091123db5d55c118"
4+
PROTON_COMMIT := "a6b2821e8ddd1127a63d3b376f860990d58931da"
55
.PHONY: all build test clean install proto
66

77
all: build
@@ -32,7 +32,6 @@ lint: ## Lint checker
3232

3333
proto: ## Generate the protobuf files
3434
@echo " > generating protobuf from goto/proton"
35-
@echo " > [info] make sure correct version of dependencies are installed using 'make install'"
3635
@buf generate https://github.com/goto/proton/archive/${PROTON_COMMIT}.zip#strip_components=1 --template buf.gen.yaml --path gotocompany/compass -v
3736
@echo " > protobuf compilation finished"
3837

‎core/asset/errors.go

+1
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88

99
var (
1010
ErrEmptyID = errors.New("asset does not have ID")
11+
ErrProbeExists = errors.New("asset probe already exists")
1112
ErrEmptyURN = errors.New("asset does not have URN")
1213
ErrUnknownType = errors.New("unknown type")
1314
ErrNilAsset = errors.New("nil asset")

‎internal/server/v1beta1/asset.go

+16-1
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"fmt"
88
"time"
99

10+
"github.com/google/uuid"
1011
"github.com/goto/compass/core/asset"
1112
"github.com/goto/compass/core/star"
1213
"github.com/goto/compass/core/user"
@@ -318,6 +319,9 @@ func (server *APIServer) CreateAssetProbe(ctx context.Context, req *compassv1bet
318319
return nil, err
319320
}
320321

322+
if req.Probe.Id != "" && !isValidUUID(req.Probe.Id) {
323+
return nil, status.Error(codes.InvalidArgument, "id should be a valid UUID")
324+
}
321325
if req.Probe.Status == "" {
322326
return nil, status.Error(codes.InvalidArgument, "Status is required")
323327
}
@@ -326,15 +330,21 @@ func (server *APIServer) CreateAssetProbe(ctx context.Context, req *compassv1bet
326330
}
327331

328332
probe := asset.Probe{
333+
ID: req.Probe.Id,
329334
Status: req.Probe.Status,
330335
StatusReason: req.Probe.StatusReason,
331336
Metadata: req.Probe.Metadata.AsMap(),
332337
Timestamp: req.Probe.Timestamp.AsTime(),
333338
}
334339
if err := server.assetService.AddProbe(ctx, req.AssetUrn, &probe); err != nil {
335-
if errors.As(err, &asset.NotFoundError{}) {
340+
switch {
341+
case errors.As(err, &asset.NotFoundError{}):
336342
return nil, status.Error(codes.NotFound, err.Error())
343+
344+
case errors.Is(err, asset.ErrProbeExists):
345+
return nil, status.Error(codes.AlreadyExists, err.Error())
337346
}
347+
338348
return nil, status.Error(codes.Internal, err.Error())
339349
}
340350

@@ -762,3 +772,8 @@ func diffChangeFromProto(pb *compassv1beta1.Change) diff.Change {
762772
To: toItf,
763773
}
764774
}
775+
776+
func isValidUUID(u string) bool {
777+
_, err := uuid.Parse(u)
778+
return err == nil
779+
}

‎internal/server/v1beta1/asset_test.go

+32
Original file line numberDiff line numberDiff line change
@@ -1233,6 +1233,18 @@ func TestCreateAssetProbe(t *testing.T) {
12331233
}
12341234

12351235
var testCases = []testCase{
1236+
{
1237+
Description: `should return error if id is not a valid UUID`,
1238+
ExpectStatus: codes.InvalidArgument,
1239+
Request: &compassv1beta1.CreateAssetProbeRequest{
1240+
AssetUrn: assetURN,
1241+
Probe: &compassv1beta1.CreateAssetProbeRequest_Probe{
1242+
Id: "invaliduuid",
1243+
Status: "RUNNING",
1244+
Timestamp: timestamppb.New(now),
1245+
},
1246+
},
1247+
},
12361248
{
12371249
Description: `should return error if status is missing`,
12381250
ExpectStatus: codes.InvalidArgument,
@@ -1269,6 +1281,26 @@ func TestCreateAssetProbe(t *testing.T) {
12691281
Return(asset.NotFoundError{URN: assetURN})
12701282
},
12711283
},
1284+
{
1285+
Description: `should return already exists if probe already exists`,
1286+
ExpectStatus: codes.AlreadyExists,
1287+
Request: &compassv1beta1.CreateAssetProbeRequest{
1288+
AssetUrn: assetURN,
1289+
Probe: &compassv1beta1.CreateAssetProbeRequest_Probe{
1290+
Id: probeID,
1291+
Status: "RUNNING",
1292+
Timestamp: timestamppb.New(now),
1293+
},
1294+
},
1295+
Setup: func(ctx context.Context, as *mocks.AssetService) {
1296+
as.EXPECT().AddProbe(ctx, assetURN, &asset.Probe{
1297+
ID: probeID,
1298+
Status: "RUNNING",
1299+
Metadata: map[string]interface{}{},
1300+
Timestamp: now,
1301+
}).Return(asset.ErrProbeExists)
1302+
},
1303+
},
12721304
{
12731305
Description: `should return internal server error if adding probe fails`,
12741306
ExpectStatus: codes.Internal,

‎internal/store/postgres/asset_repository.go

+21-10
Original file line numberDiff line numberDiff line change
@@ -346,21 +346,32 @@ func (r *AssetRepository) AddProbe(ctx context.Context, assetURN string, probe *
346346
probe.Timestamp = probe.Timestamp.UTC()
347347
}
348348

349-
query, args, err := sq.Insert("asset_probes").
350-
Columns("asset_urn", "status", "status_reason", "metadata", "timestamp", "created_at").
351-
Values(assetURN, probe.Status, probe.StatusReason, probe.Metadata, probe.Timestamp, probe.CreatedAt).
352-
Suffix("RETURNING \"id\"").
349+
insert := sq.Insert("asset_probes")
350+
if probe.ID != "" {
351+
insert = insert.Columns("id", "asset_urn", "status", "status_reason", "metadata", "timestamp", "created_at").
352+
Values(probe.ID, assetURN, probe.Status, probe.StatusReason, probe.Metadata, probe.Timestamp, probe.CreatedAt)
353+
} else {
354+
insert = insert.Columns("asset_urn", "status", "status_reason", "metadata", "timestamp", "created_at").
355+
Values(assetURN, probe.Status, probe.StatusReason, probe.Metadata, probe.Timestamp, probe.CreatedAt)
356+
}
357+
358+
query, args, err := insert.Suffix("RETURNING \"id\"").
353359
PlaceholderFormat(sq.Dollar).
354360
ToSql()
355361
if err != nil {
356-
return fmt.Errorf("error building insert asset probe query: %w", err)
362+
return fmt.Errorf("build insert asset probe query: %w", err)
357363
}
358364

359-
err = r.client.db.QueryRowContext(ctx, query, args...).Scan(&probe.ID)
360-
if errors.Is(checkPostgresError(err), errForeignKeyViolation) {
361-
return asset.NotFoundError{URN: assetURN}
362-
} else if err != nil {
363-
return fmt.Errorf("error running insert asset probe query: %w", err)
365+
if err = r.client.db.QueryRowContext(ctx, query, args...).Scan(&probe.ID); err != nil {
366+
switch e := checkPostgresError(err); {
367+
case errors.Is(e, errForeignKeyViolation):
368+
return asset.NotFoundError{URN: assetURN}
369+
370+
case errors.Is(e, errDuplicateKey):
371+
return asset.ErrProbeExists
372+
}
373+
374+
return fmt.Errorf("run insert asset probe query: %w", err)
364375
}
365376

366377
return nil

‎internal/store/postgres/asset_repository_test.go

+56
Original file line numberDiff line numberDiff line change
@@ -1290,7 +1290,36 @@ func (r *AssetRepositoryTestSuite) TestAddProbe() {
12901290
r.ErrorAs(err, &asset.NotFoundError{URN: urn})
12911291
})
12921292

1293+
r.Run("should return error if probe already exists", func() {
1294+
ast := asset.Asset{
1295+
URN: "urn-add-probe-1",
1296+
Type: asset.TypeJob,
1297+
Service: "airflow",
1298+
UpdatedBy: user.User{ID: defaultAssetUpdaterUserID},
1299+
}
1300+
probeID := uuid.NewString()
1301+
probe := asset.Probe{
1302+
ID: probeID,
1303+
Status: "COMPLETED",
1304+
StatusReason: "Sample Reason",
1305+
Timestamp: time.Now().Add(2 * time.Minute),
1306+
Metadata: map[string]interface{}{
1307+
"foo": "bar",
1308+
},
1309+
}
1310+
1311+
_, err := r.repository.Upsert(r.ctx, &ast)
1312+
r.Require().NoError(err)
1313+
1314+
err = r.repository.AddProbe(r.ctx, ast.URN, &probe)
1315+
r.NoError(err)
1316+
1317+
err = r.repository.AddProbe(r.ctx, ast.URN, &probe)
1318+
r.ErrorIs(err, asset.ErrProbeExists)
1319+
})
1320+
12931321
r.Run("should populate CreatedAt and persist probe", func() {
1322+
r.BeforeTest("", "")
12941323
ast := asset.Asset{
12951324
URN: "urn-add-probe-1",
12961325
Type: asset.TypeJob,
@@ -1337,6 +1366,33 @@ func (r *AssetRepositoryTestSuite) TestAddProbe() {
13371366
r.Require().NoError(err)
13381367
})
13391368

1369+
r.Run("should insert ID if specified", func() {
1370+
ast := asset.Asset{
1371+
URN: "urn-add-probe-1",
1372+
Type: asset.TypeJob,
1373+
Service: "airflow",
1374+
UpdatedBy: user.User{ID: defaultAssetUpdaterUserID},
1375+
}
1376+
probeID := uuid.NewString()
1377+
probe := asset.Probe{
1378+
ID: probeID,
1379+
Status: "COMPLETED",
1380+
StatusReason: "Sample Reason",
1381+
Timestamp: time.Now().Add(2 * time.Minute),
1382+
Metadata: map[string]interface{}{
1383+
"foo": "bar",
1384+
},
1385+
}
1386+
1387+
_, err := r.repository.Upsert(r.ctx, &ast)
1388+
r.Require().NoError(err)
1389+
1390+
err = r.repository.AddProbe(r.ctx, ast.URN, &probe)
1391+
r.NoError(err)
1392+
1393+
r.Equal(probeID, probe.ID)
1394+
})
1395+
13401396
r.Run("should populate Timestamp if empty", func() {
13411397
ast := asset.Asset{
13421398
URN: "urn-add-probe-2",

‎proto/compass.swagger.yaml

+2
Original file line numberDiff line numberDiff line change
@@ -1877,6 +1877,8 @@ definitions:
18771877
CreateAssetProbeRequest.Probe:
18781878
type: object
18791879
properties:
1880+
id:
1881+
type: string
18801882
metadata:
18811883
type: object
18821884
status:

‎proto/gotocompany/compass/v1beta1/service.pb.go

+12-3
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎proto/gotocompany/compass/v1beta1/service.pb.validate.go

+2
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)