diff --git a/api/docInitRunner/runner_test.go b/api/docInitRunner/runner_test.go index 556321a2..df9771cf 100644 --- a/api/docInitRunner/runner_test.go +++ b/api/docInitRunner/runner_test.go @@ -10,6 +10,7 @@ import ( "queryorchestration/internal/document" documentinit "queryorchestration/internal/document/init" "queryorchestration/internal/server/runner" + "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/serviceconfig/queue/documentsync" queuemock "queryorchestration/mocks/queue" @@ -47,7 +48,11 @@ func TestDocInitRunner(t *testing.T) { t.Run("valid", func(t *testing.T) { clientId := "clientid" bucketName := "bucketName" - location := fmt.Sprintf("%s/import/aaa", clientId) + location := objectstore.BucketKey{ + Location: objectstore.Import, + ClientID: clientId, + Filename: "aaa", + } docinfo := document.DocumentSummary{ ID: uuid.New(), ClientID: "hello", @@ -55,7 +60,7 @@ func TestDocInitRunner(t *testing.T) { } doc := docinitrunner.Body{ Bucket: bucketName, - Key: location, + Key: location.String(), Hash: docinfo.Hash, ClientID: clientId, } @@ -69,7 +74,7 @@ func TestDocInitRunner(t *testing.T) { pgxmock.NewRows([]string{"id"}). AddRow(docinfo.ID), ) - pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(docinfo.ID, bucketName, location). + pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(docinfo.ID, bucketName, location.String()). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(docinfo.ID). diff --git a/api/docTextProcessRunner/runner.go b/api/docTextProcessRunner/runner.go index b64a454f..ad99c96f 100644 --- a/api/docTextProcessRunner/runner.go +++ b/api/docTextProcessRunner/runner.go @@ -5,6 +5,7 @@ import ( "log/slog" documenttext "queryorchestration/internal/document/text" + "queryorchestration/internal/serviceconfig/objectstore" "github.com/go-playground/validator/v10" ) @@ -35,9 +36,15 @@ type Body struct { } func (s Runner) Process(ctx context.Context, body Body) bool { - err := s.svc.Text.Process(ctx, &documenttext.ProcessParams{ + key, err := objectstore.ParseBucketKey(body.Key) + if err != nil { + slog.Error("unable to parse key", "key", body.Key) + return true + } + + err = s.svc.Text.Process(ctx, &documenttext.ProcessParams{ Bucket: body.Bucket, - Key: body.Key, + Key: key, Hash: body.Hash, ClientID: body.ClientID, }) diff --git a/api/docTextProcessRunner/runner_test.go b/api/docTextProcessRunner/runner_test.go index dcc888ef..80769649 100644 --- a/api/docTextProcessRunner/runner_test.go +++ b/api/docTextProcessRunner/runner_test.go @@ -2,10 +2,12 @@ package doctextprocessrunner import ( "context" + "errors" "fmt" "io" "strings" "testing" + "time" "queryorchestration/internal/database/repository" "queryorchestration/internal/document" @@ -55,7 +57,7 @@ func TestDocTextProcessRunner(t *testing.T) { }) assert.NotNil(t, runner) - t.Run("valid", func(t *testing.T) { + t.Run("pass through", func(t *testing.T) { triggerId := uuid.New() docId := uuid.New() textractId := "heeloo" @@ -67,10 +69,15 @@ func TestDocTextProcessRunner(t *testing.T) { ClientID: "hello", Hash: `"5d41402abc4b2a76b9719d911017c592"`, } - location := fmt.Sprintf("%s/text/textract/%s", docinfo.ClientID, triggerId) + key := objectstore.BucketKey{ + ClientID: docinfo.ClientID, + Location: objectstore.TextTextract, + Filename: triggerId.String(), + CreatedAt: time.Now().UTC(), + } doc := Body{ Bucket: bucketName, - Key: location, + Key: key.String(), Hash: docinfo.Hash, ClientID: docinfo.ClientID, } @@ -113,4 +120,38 @@ func TestDocTextProcessRunner(t *testing.T) { assert.True(t, runner.Process(ctx, doc)) }) + t.Run("invalid key", func(t *testing.T) { + key := objectstore.BucketKey{} + doc := Body{ + Key: key.String(), + } + + assert.True(t, runner.Process(ctx, doc)) + }) + t.Run("unable to process", func(t *testing.T) { + triggerId := uuid.New() + bucketName := "bucketName" + docinfo := document.DocumentSummary{ + ID: uuid.New(), + ClientID: "hello", + Hash: `"5d41402abc4b2a76b9719d911017c592"`, + } + key := objectstore.BucketKey{ + ClientID: docinfo.ClientID, + Location: objectstore.TextTextract, + Filename: triggerId.String(), + CreatedAt: time.Now().UTC(), + } + doc := Body{ + Bucket: bucketName, + Key: key.String(), + Hash: docinfo.Hash, + ClientID: docinfo.ClientID, + } + + pool.ExpectQuery("name: GetDocumentTextTrigger :one").WithArgs(triggerId). + WillReturnError(errors.New("db fail")) + + assert.False(t, runner.Process(ctx, doc)) + }) } diff --git a/api/storeEventRunner/runner_test.go b/api/storeEventRunner/runner_test.go index 24017109..0a4e8b19 100644 --- a/api/storeEventRunner/runner_test.go +++ b/api/storeEventRunner/runner_test.go @@ -2,13 +2,14 @@ package storeeventrunner import ( "context" - "fmt" "testing" + "time" "queryorchestration/internal/database/repository" "queryorchestration/internal/document" documentstore "queryorchestration/internal/document/store" "queryorchestration/internal/server/runner" + "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/serviceconfig/queue/documentinit" "queryorchestration/internal/serviceconfig/queue/documenttextprocess" queuemock "queryorchestration/mocks/queue" @@ -50,7 +51,12 @@ func TestDocInitRunner(t *testing.T) { t.Run("valid", func(t *testing.T) { clientId := "hi" bucketName := "bucketName" - location := fmt.Sprintf("%s/import/aaa", clientId) + location := objectstore.BucketKey{ + Location: objectstore.Import, + ClientID: clientId, + Filename: "aaa", + CreatedAt: time.Now().UTC(), + } docinfo := document.DocumentSummary{ ID: uuid.New(), ClientID: clientId, @@ -65,7 +71,7 @@ func TestDocInitRunner(t *testing.T) { Name: bucketName, }, Object: S3Object{ - Key: location, + Key: location.String(), ETag: docinfo.Hash, }, }, @@ -82,7 +88,7 @@ func TestDocInitRunner(t *testing.T) { pgxmock.NewRows([]string{"id"}). AddRow(docinfo.ID), ) - pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(docinfo.ID, bucketName, location). + pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(docinfo.ID, bucketName, location.String()). WillReturnResult(pgxmock.NewResult("", 1)) pool.ExpectCommit() pool.ExpectQuery("name: GetDocumentSummary :one").WithArgs(docinfo.ID). diff --git a/internal/document/store/process.go b/internal/document/store/process.go index e95370f9..825adf76 100644 --- a/internal/document/store/process.go +++ b/internal/document/store/process.go @@ -2,13 +2,11 @@ package documentstore import ( "context" - "fmt" "log/slog" - "regexp" docinitrunner "queryorchestration/api/docInitRunner" doctextprocessrunner "queryorchestration/api/docTextProcessRunner" - "queryorchestration/internal/client" + "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/serviceconfig/queue" ) @@ -31,25 +29,29 @@ func (s *Service) Process(ctx context.Context, params Params) error { return nil } - task, clientId := getKeyMetadata(params.Key) + key, err := objectstore.ParseBucketKey(params.Key) + if err != nil { + slog.Error("unable to parse key", "key", params.Key) + return nil + } queueParams := &queue.SendParams{} - switch task { - case DocInit: + switch key.Location { + case objectstore.Import: queueParams.QueueURL = s.cfg.GetDocInitURL() queueParams.Body = docinitrunner.Body{ Bucket: params.Bucket, Key: params.Key, Hash: params.Hash, - ClientID: clientId, + ClientID: key.ClientID, } - case DocTextProcess: + case objectstore.TextTextract: queueParams.QueueURL = s.cfg.GetDocumentTextProcessURL() queueParams.Body = doctextprocessrunner.Body{ Bucket: params.Bucket, Key: params.Key, Hash: params.Hash, - ClientID: clientId, + ClientID: key.ClientID, } default: slog.Info("unsupported key", "key", params.Key) @@ -59,32 +61,6 @@ func (s *Service) Process(ctx context.Context, params Params) error { return s.cfg.SendToQueue(ctx, queueParams) } -type Task string - -const ( - DocInit Task = "import" - DocTextProcess Task = "text/textract" - Invalid Task = "" -) - -func getKeyMetadata(key string) (Task, string) { - regexStr := fmt.Sprintf(`^(%s)/(import|text/textract)/.+$`, client.CLIENT_ID_REGEX) - re := regexp.MustCompile(regexStr) - match := re.FindStringSubmatch(key) - if len(match) < 3 { - return Invalid, "" - } - - task := Task(match[2]) - if task == Invalid { - return Invalid, "" - } - - clientId := match[1] - - return task, clientId -} - func (s *Service) isSupportedEvent(name EventS3) bool { return name == EventS3ObjectCreatedPut } diff --git a/internal/document/store/process_test.go b/internal/document/store/process_test.go index 5fd2c16f..cb17386e 100644 --- a/internal/document/store/process_test.go +++ b/internal/document/store/process_test.go @@ -2,10 +2,11 @@ package documentstore import ( "context" - "fmt" "testing" + "time" "queryorchestration/internal/serviceconfig" + "queryorchestration/internal/serviceconfig/objectstore" "queryorchestration/internal/serviceconfig/queue/documentinit" "queryorchestration/internal/serviceconfig/queue/documenttextprocess" queuemock "queryorchestration/mocks/queue" @@ -56,9 +57,15 @@ func TestProcess(t *testing.T) { ). Return(&sqs.SendMessageOutput{}, nil) + location := objectstore.BucketKey{ + ClientID: "7db16095-9155-47d4-8004-b3b3ead93c83", + Location: objectstore.Import, + Filename: "aaa", + CreatedAt: time.Now().UTC(), + } err := svc.Process(ctx, Params{ Event: EventS3ObjectCreatedPut, - Key: "7db16095-9155-47d4-8004-b3b3ead93c83/import/a", + Key: location.String(), }) assert.NoError(t, err) }) @@ -73,9 +80,15 @@ func TestProcess(t *testing.T) { ). Return(&sqs.SendMessageOutput{}, nil) + location := objectstore.BucketKey{ + ClientID: "7db16095-9155-47d4-8004-b3b3ead93c83", + Location: objectstore.TextTextract, + Filename: "aaa", + CreatedAt: time.Now().UTC(), + } err := svc.Process(ctx, Params{ Event: EventS3ObjectCreatedPut, - Key: "7db16095-9155-47d4-8004-b3b3ead93c83/text/textract/a", + Key: location.String(), }) assert.NoError(t, err) }) @@ -91,27 +104,3 @@ func TestIsSupportedEvent(t *testing.T) { assert.False(t, svc.isSupportedEvent("invalid")) }) } - -func TestGetKeyMetadata(t *testing.T) { - baseClientId := "hello" - - task, clientId := getKeyMetadata(fmt.Sprintf("%s/import/a", baseClientId)) - assert.Equal(t, DocInit, task) - assert.Equal(t, baseClientId, clientId) - - task, clientId = getKeyMetadata(fmt.Sprintf("%s/text/textract/a", baseClientId)) - assert.Equal(t, DocTextProcess, task) - assert.Equal(t, baseClientId, clientId) - - task, clientId = getKeyMetadata("a") - assert.Equal(t, Invalid, task) - assert.Equal(t, "", clientId) - - task, clientId = getKeyMetadata("fdee2e22-ad0f-4e30-91c8-d63c3501162f/something/a") - assert.Equal(t, Invalid, task) - assert.Equal(t, "", clientId) - - task, clientId = getKeyMetadata("a/import/a") - assert.Equal(t, DocInit, task) - assert.Equal(t, "a", clientId) -} diff --git a/internal/document/text/processExtract.go b/internal/document/text/processExtract.go index a35441b2..2fe67fc2 100644 --- a/internal/document/text/processExtract.go +++ b/internal/document/text/processExtract.go @@ -5,12 +5,13 @@ import ( "context" "database/sql" "errors" - "fmt" "io" "strings" + "time" "queryorchestration/internal/database/repository" "queryorchestration/internal/serviceconfig/build" + "queryorchestration/internal/serviceconfig/objectstore" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/google/uuid" @@ -18,7 +19,7 @@ import ( type ProcessParams struct { Bucket string - Key string + Key objectstore.BucketKey Hash string ClientID string } @@ -44,9 +45,10 @@ func (s *Service) mergePages(pages []*Page) (io.Reader, error) { } func (s *Service) processTrigger(ctx context.Context, trigger *repository.GetDocumentTextTriggerRow, params *ProcessParams) error { + keyStr := params.Key.String() response, err := s.cfg.GetStoreClient().GetObject(ctx, &s3.GetObjectInput{ Bucket: ¶ms.Bucket, - Key: ¶ms.Key, + Key: &keyStr, IfMatch: ¶ms.Hash, }) if err != nil { @@ -71,10 +73,6 @@ func (s *Service) processTrigger(ctx context.Context, trigger *repository.GetDoc return nil } -func (s *Service) createTextKey(clientId string, textId uuid.UUID) string { - return fmt.Sprintf("%s/text/out/%s", clientId, textId) -} - func (s *Service) storeProcess(ctx context.Context, text io.Reader, trigger *repository.GetDocumentTextTriggerRow, params *ProcessParams) error { version := build.GetVersionUnixTimestamp() @@ -99,11 +97,17 @@ func (s *Service) storeProcess(ctx context.Context, text io.Reader, trigger *rep var textId uuid.UUID if existingId == uuid.Nil || errors.Is(err, sql.ErrNoRows) { - key := s.createTextKey(params.ClientID, trigger.ID) + key := objectstore.BucketKey{ + ClientID: params.ClientID, + Location: objectstore.TextOut, + CreatedAt: time.Now().UTC(), + Filename: trigger.ID.String(), + } + keyStr := key.String() textId, err = q.AddDocumentText(ctx, &repository.AddDocumentTextParams{ Bucket: params.Bucket, - Key: key, + Key: keyStr, Hash: hash, }) if err != nil { @@ -112,7 +116,7 @@ func (s *Service) storeProcess(ctx context.Context, text io.Reader, trigger *rep _, err = s.cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{ Bucket: ¶ms.Bucket, - Key: &key, + Key: &keyStr, Body: storeReader, }) if err != nil { diff --git a/internal/document/text/processExtract_test.go b/internal/document/text/processExtract_test.go index d05a052c..3e778a3e 100644 --- a/internal/document/text/processExtract_test.go +++ b/internal/document/text/processExtract_test.go @@ -7,8 +7,10 @@ import ( "io" "strings" "testing" + "time" "queryorchestration/internal/database/repository" + "queryorchestration/internal/serviceconfig/objectstore" objectstoremock "queryorchestration/mocks/objectstore" queuemock "queryorchestration/mocks/queue" @@ -82,9 +84,14 @@ func TestProcess(t *testing.T) { ). Return(&sqs.SendMessageOutput{}, nil) + key := objectstore.BucketKey{ + ClientID: "my_client", + Location: objectstore.TextTextract, + Filename: triggerId.String(), + } err = svc.Process(ctx, &ProcessParams{ Bucket: "bucket", - Key: "6ece4a77-702c-4cc9-8654-43c557e21658/text/textract/6ac21b2b-f36c-4ae7-a815-307d4a9bfe4d", + Key: key, Hash: "hash", }) assert.NoError(t, err) @@ -131,16 +138,22 @@ func TestStoreTrigger(t *testing.T) { Cleanid: cleanId, }, &ProcessParams{ Bucket: bucket, + Key: objectstore.BucketKey{}, }) require.NoError(t, err) }) t.Run("not exists", func(t *testing.T) { - key := fmt.Sprintf("%s/text/out/%s", clientId, triggerId) + key := objectstore.BucketKey{ + ClientID: clientId, + Location: objectstore.TextOut, + CreatedAt: time.Now().UTC(), + Filename: triggerId.String(), + } hash := `"09644372e99020106946045c6fd2d70b"` pool.ExpectBegin() pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash). WillReturnError(sql.ErrNoRows) - pool.ExpectQuery("name: AddDocumentText :one").WithArgs(bucket, key, hash). + pool.ExpectQuery("name: AddDocumentText :one").WithArgs(bucket, key.String(), hash). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(textId), @@ -152,7 +165,7 @@ func TestStoreTrigger(t *testing.T) { mock.MatchedBy(func(in *s3.PutObjectInput) bool { bod, err := io.ReadAll(in.Body) require.NoError(t, err) - return string(bod) == "byebye" && *in.Key == key && *in.Bucket == bucket + return string(bod) == "byebye" && *in.Key == key.String() && *in.Bucket == bucket }), mock.Anything, ). @@ -167,6 +180,7 @@ func TestStoreTrigger(t *testing.T) { ID: triggerId, Cleanid: cleanId, }, &ProcessParams{ + Key: key, Bucket: bucket, ClientID: clientId, }) @@ -209,13 +223,18 @@ func TestProcessTrigger(t *testing.T) { triggerId := uuid.New() clientId := "hello" bucket := "bucket" - key := fmt.Sprintf("%s/text/out/%s", clientId, triggerId) + key := objectstore.BucketKey{ + ClientID: clientId, + Location: objectstore.TextOut, + Filename: objectstore.BucketFilename(triggerId.String()), + CreatedAt: time.Now().UTC(), + } hash := `"5d41402abc4b2a76b9719d911017c592"` pool.ExpectBegin() pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash). WillReturnError(sql.ErrNoRows) - pool.ExpectQuery("name: AddDocumentText :one").WithArgs(bucket, key, hash). + pool.ExpectQuery("name: AddDocumentText :one").WithArgs(bucket, key.String(), hash). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(textId), @@ -226,7 +245,7 @@ func TestProcessTrigger(t *testing.T) { GetObject( mock.Anything, mock.MatchedBy(func(in *s3.GetObjectInput) bool { - return *in.Key == key && *in.Bucket == bucket && *in.IfMatch == hash + return *in.Key == key.String() && *in.Bucket == bucket && *in.IfMatch == hash }), mock.Anything, ). @@ -240,7 +259,7 @@ func TestProcessTrigger(t *testing.T) { mock.MatchedBy(func(in *s3.PutObjectInput) bool { bod, err := io.ReadAll(in.Body) require.NoError(t, err) - return string(bod) == "hello" && *in.Key == key && *in.Bucket == bucket + return string(bod) == "hello" && *in.Key == key.String() && *in.Bucket == bucket }), mock.Anything, ). diff --git a/internal/document/text/triggerExtract.go b/internal/document/text/triggerExtract.go index 9c405ff3..a557f6af 100644 --- a/internal/document/text/triggerExtract.go +++ b/internal/document/text/triggerExtract.go @@ -3,12 +3,11 @@ package documenttext import ( "context" "errors" - "fmt" - "regexp" + "time" - "queryorchestration/internal/client" "queryorchestration/internal/database/repository" "queryorchestration/internal/serviceconfig/build" + "queryorchestration/internal/serviceconfig/objectstore" "github.com/aws/aws-sdk-go-v2/service/textract" "github.com/aws/aws-sdk-go-v2/service/textract/types" @@ -29,11 +28,19 @@ func (s *Service) triggerExtract(ctx context.Context, clean *repository.Currentc jobTag := triggerId.String() - key, err := s.getOutputLocation(clean.Clientid, triggerId) + filename, err := s.getOutputFilename(triggerId) if err != nil { return err } + key := objectstore.BucketKey{ + CreatedAt: time.Now().UTC(), + Location: objectstore.TextTextract, + ClientID: clean.Clientid, + Filename: filename, + } + keyStr := key.String() + res, err := s.cfg.GetTextractClient().StartDocumentTextDetection(ctx, &textract.StartDocumentTextDetectionInput{ DocumentLocation: &types.DocumentLocation{ S3Object: &types.S3Object{ @@ -44,7 +51,7 @@ func (s *Service) triggerExtract(ctx context.Context, clean *repository.Currentc JobTag: &jobTag, OutputConfig: &types.OutputConfig{ S3Bucket: clean.Bucket, - S3Prefix: &key, + S3Prefix: &keyStr, }, }) if err != nil { @@ -65,25 +72,16 @@ func (s *Service) triggerExtract(ctx context.Context, clean *repository.Currentc }) } -func (s *Service) getOutputLocation(clientId string, triggerId uuid.UUID) (string, error) { - if clientId == "" { - return "", errors.New("client id required") - } else if triggerId == uuid.Nil { +func (s *Service) getOutputFilename(triggerId uuid.UUID) (string, error) { + if triggerId == uuid.Nil { return "", errors.New("trigger id required") } - return fmt.Sprintf("%s/text/textract/%s", clientId, triggerId), nil + return triggerId.String(), nil } -func (s *Service) getTriggerIdFromKey(key string) (uuid.UUID, error) { - regexPattern := fmt.Sprintf(`^%s/text/textract/([0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12})$`, client.CLIENT_ID_REGEX) - re := regexp.MustCompile(regexPattern) - match := re.FindStringSubmatch(key) - if len(match) < 2 { - return uuid.Nil, errors.New("Trigger Id not found") - } - - triggerId, err := uuid.Parse(match[1]) +func (s *Service) getTriggerIdFromKey(key objectstore.BucketKey) (uuid.UUID, error) { + triggerId, err := uuid.Parse(key.Filename) if err != nil { return uuid.Nil, err } diff --git a/internal/document/text/triggerExtract_test.go b/internal/document/text/triggerExtract_test.go index fdfe0b40..8026297c 100644 --- a/internal/document/text/triggerExtract_test.go +++ b/internal/document/text/triggerExtract_test.go @@ -2,11 +2,11 @@ package documenttext import ( "context" - "fmt" "testing" "queryorchestration/internal/database/repository" "queryorchestration/internal/document" + "queryorchestration/internal/serviceconfig/objectstore" objectstoremock "queryorchestration/mocks/objectstore" textractmock "queryorchestration/mocks/textract" @@ -82,40 +82,32 @@ func TestTriggerExtract(t *testing.T) { }) } -func TestGetOutputLocation(t *testing.T) { +func TestGetOutputFilename(t *testing.T) { svc := Service{} - _, err := svc.getOutputLocation("", uuid.Nil) - assert.EqualError(t, err, "client id required") - - _, err = svc.getOutputLocation("a", uuid.Nil) + _, err := svc.getOutputFilename(uuid.Nil) assert.EqualError(t, err, "trigger id required") triggerId := uuid.New() - key, err := svc.getOutputLocation("a", triggerId) + key, err := svc.getOutputFilename(triggerId) assert.NoError(t, err) - assert.Equal(t, fmt.Sprintf("a/text/textract/%s", triggerId), key) + assert.Equal(t, triggerId.String(), key) } func TestGetTriggerIdFromKey(t *testing.T) { svc := Service{} - _, err := svc.getTriggerIdFromKey("") - assert.EqualError(t, err, "Trigger Id not found") + _, err := svc.getTriggerIdFromKey(objectstore.BucketKey{}) + assert.EqualError(t, err, "invalid UUID length: 0") - _, err = svc.getTriggerIdFromKey("aa") - assert.EqualError(t, err, "Trigger Id not found") + _, err = svc.getTriggerIdFromKey(objectstore.BucketKey{ + Filename: "aa", + }) + assert.EqualError(t, err, "invalid UUID length: 2") - _, err = svc.getTriggerIdFromKey("6ece4a77-702c-4cc9-8654-43c557e21658") - assert.EqualError(t, err, "Trigger Id not found") - - _, err = svc.getTriggerIdFromKey("6ece4a77-702c-4cc9-8654-43c557e21658/text/textract/") - assert.EqualError(t, err, "Trigger Id not found") - - _, err = svc.getTriggerIdFromKey("6ece4a77-702c-4cc9-8654-43c557e21658/text/textract/a") - assert.EqualError(t, err, "Trigger Id not found") - - triggerId, err := svc.getTriggerIdFromKey("6ece4a77-702c-4cc9-8654-43c557e21658/text/textract/6ac21b2b-f36c-4ae7-a815-307d4a9bfe4d") + triggerId, err := svc.getTriggerIdFromKey(objectstore.BucketKey{ + Filename: "6ac21b2b-f36c-4ae7-a815-307d4a9bfe4d", + }) assert.NoError(t, err) assert.Equal(t, "6ac21b2b-f36c-4ae7-a815-307d4a9bfe4d", triggerId.String()) } diff --git a/internal/serviceconfig/objectstore/config.go b/internal/serviceconfig/objectstore/config.go index 6bb2c633..1b7ac01c 100644 --- a/internal/serviceconfig/objectstore/config.go +++ b/internal/serviceconfig/objectstore/config.go @@ -4,11 +4,15 @@ import ( "context" "crypto/md5" "encoding/hex" + "errors" "fmt" "io" "log/slog" + "regexp" + "strconv" "time" + "queryorchestration/internal/client" awsc "queryorchestration/internal/serviceconfig/aws" "github.com/aws/aws-sdk-go-v2/service/s3" @@ -140,3 +144,75 @@ func (c *ObjectStoreConfig) CalculateETag(ctx context.Context, doc io.Reader) (s return fmt.Sprintf(`"%s-%d"`, hex.EncodeToString(finalHash.Sum(nil)), len(partHashes)), nil } + +type BucketLocation string +type BucketFilename = string +type BucketKey struct { + ClientID string + Location BucketLocation + CreatedAt time.Time + Part *uint8 + Filename BucketFilename +} + +const ( + Import BucketLocation = "import" + TextTextract BucketLocation = "text/textract" + TextOut BucketLocation = "text/out" + Export BucketLocation = "export" +) + +func (c *BucketKey) String() string { + dateStr := c.CreatedAt.Format("20060102") + partStr := "" + if c.Part != nil { + partStr = fmt.Sprintf("/%d", *c.Part) + } + + return fmt.Sprintf("%s/%s/%s%s/%s", c.ClientID, c.Location, dateStr, partStr, c.Filename) +} + +func ParseBucketKey(key string) (BucketKey, error) { + locations := fmt.Sprintf("%s|%s|%s|%s", Import, TextTextract, TextOut, Export) + pattern := fmt.Sprintf(`^(%s)/(%s)/([0-9]{8})/(.+)$`, client.CLIENT_ID_REGEX, locations) + reg := regexp.MustCompile(pattern) + matches := reg.FindStringSubmatch(key) + if len(matches) != 5 { + return BucketKey{}, errors.New("does not match expectation") + } + + formattedKey := BucketKey{} + + formattedKey.ClientID = matches[1] + formattedKey.Location = BucketLocation(matches[2]) + + parsedTime, err := time.Parse("20060102", matches[3]) + if err != nil { + return BucketKey{}, fmt.Errorf("failed to parse date: %w", err) + } + formattedKey.CreatedAt = parsedTime + + filename := `[a-zA-Z0-9\_\.\-]+` + reg = regexp.MustCompile(fmt.Sprintf(`^([0-9]+)/(%s)$`, filename)) + subMatches := reg.FindStringSubmatch(matches[4]) + if len(subMatches) == 3 { + part, err := strconv.ParseUint(subMatches[1], 10, 8) + if err != nil { + return BucketKey{}, err + } + safePart := uint8(part) + formattedKey.Part = &safePart + + formattedKey.Filename = subMatches[2] + } else { + reg = regexp.MustCompile(fmt.Sprintf(`^(%s)$`, filename)) + subMatches := reg.FindStringSubmatch(matches[4]) + if len(subMatches) != 2 { + return BucketKey{}, errors.New("does not match expectation") + } + + formattedKey.Filename = subMatches[1] + } + + return formattedKey, nil +} diff --git a/internal/serviceconfig/objectstore/config_test.go b/internal/serviceconfig/objectstore/config_test.go index 5385da39..42f44cf1 100644 --- a/internal/serviceconfig/objectstore/config_test.go +++ b/internal/serviceconfig/objectstore/config_test.go @@ -115,3 +115,160 @@ func TestCalculateETag(t *testing.T) { assert.Equal(t, *res.ETag, hash) } + +func TestGetBucketKey(t *testing.T) { + key := objectstore.BucketKey{} + assert.Equal(t, "//00010101/", key.String()) + key.CreatedAt = time.Date(2025, time.March, 31, 0, 0, 0, 0, time.UTC) + assert.Equal(t, "//20250331/", key.String()) + key.ClientID = "AAA" + assert.Equal(t, "AAA//20250331/", key.String()) + key.Filename = "coolfile" + assert.Equal(t, "AAA//20250331/coolfile", key.String()) + key.Location = objectstore.Import + assert.Equal(t, "AAA/import/20250331/coolfile", key.String()) + part := uint8(3) + key.Part = &part + assert.Equal(t, "AAA/import/20250331/3/coolfile", key.String()) + key.Location = objectstore.TextTextract + assert.Equal(t, "AAA/text/textract/20250331/3/coolfile", key.String()) + key.Location = objectstore.TextOut + assert.Equal(t, "AAA/text/out/20250331/3/coolfile", key.String()) + key.Location = objectstore.Export + assert.Equal(t, "AAA/export/20250331/3/coolfile", key.String()) +} +func TestParseBucketKey(t *testing.T) { + _, err := objectstore.ParseBucketKey("") + assert.Error(t, err) + + _, err = objectstore.ParseBucketKey("AAA") + assert.Error(t, err) + + _, err = objectstore.ParseBucketKey("//00010101/") + assert.Error(t, err) + + _, err = objectstore.ParseBucketKey("//20250331/") + assert.Error(t, err) + + _, err = objectstore.ParseBucketKey("AAA//20250331/") + assert.Error(t, err) + + _, err = objectstore.ParseBucketKey("AAA//20250331/coolfile") + assert.Error(t, err) + + key, err := objectstore.ParseBucketKey("AAA/import/20250331/coolfile") + assert.NoError(t, err) + assert.EqualExportedValues(t, objectstore.BucketKey{ + ClientID: "AAA", + Location: objectstore.Import, + Filename: "coolfile", + CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC), + }, key) + + key, err = objectstore.ParseBucketKey("AAA/text/textract/20250331/coolfile") + assert.NoError(t, err) + assert.EqualExportedValues(t, objectstore.BucketKey{ + ClientID: "AAA", + Location: objectstore.TextTextract, + Filename: "coolfile", + CreatedAt: time.Date(2025, 2, 31, 0, 0, 0, 0, time.UTC), + }, key) + + key, err = objectstore.ParseBucketKey("AAA/text/out/20250331/coolfile") + assert.NoError(t, err) + assert.EqualExportedValues(t, objectstore.BucketKey{ + ClientID: "AAA", + Location: objectstore.TextOut, + Filename: "coolfile", + CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC), + }, key) + + key, err = objectstore.ParseBucketKey("AAA/export/20250331/coolfile") + assert.NoError(t, err) + assert.EqualExportedValues(t, objectstore.BucketKey{ + ClientID: "AAA", + Location: objectstore.Export, + Filename: "coolfile", + CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC), + }, key) + + part := uint8(3) + key, err = objectstore.ParseBucketKey("AAA/import/20250331/3/coolfile") + assert.NoError(t, err) + assert.EqualExportedValues(t, objectstore.BucketKey{ + ClientID: "AAA", + Location: objectstore.Import, + Filename: "coolfile", + CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC), + Part: &part, + }, key) + + key, err = objectstore.ParseBucketKey("AAA/text/textract/20250331/3/coolfile") + assert.NoError(t, err) + assert.EqualExportedValues(t, objectstore.BucketKey{ + ClientID: "AAA", + Location: objectstore.TextTextract, + Filename: "coolfile", + CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC), + Part: &part, + }, key) + + key, err = objectstore.ParseBucketKey("AAA/text/out/20250331/3/coolfile") + assert.NoError(t, err) + assert.EqualExportedValues(t, objectstore.BucketKey{ + ClientID: "AAA", + Location: objectstore.TextOut, + Filename: "coolfile", + CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC), + Part: &part, + }, key) + + key, err = objectstore.ParseBucketKey("AAA/export/20250331/3/coolfile") + assert.NoError(t, err) + assert.EqualExportedValues(t, objectstore.BucketKey{ + ClientID: "AAA", + Location: objectstore.Export, + Filename: "coolfile", + CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC), + Part: &part, + }, key) + + key, err = objectstore.ParseBucketKey("AAA/import/20250331/cool-file") + assert.NoError(t, err) + assert.EqualExportedValues(t, objectstore.BucketKey{ + ClientID: "AAA", + Location: objectstore.Import, + Filename: "cool-file", + CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC), + }, key) + + key, err = objectstore.ParseBucketKey("AAA/import/20250331/cool_file") + assert.NoError(t, err) + assert.EqualExportedValues(t, objectstore.BucketKey{ + ClientID: "AAA", + Location: objectstore.Import, + Filename: "cool_file", + CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC), + }, key) + + key, err = objectstore.ParseBucketKey("AAA/import/20250331/coolfile365") + assert.NoError(t, err) + assert.EqualExportedValues(t, objectstore.BucketKey{ + ClientID: "AAA", + Location: objectstore.Import, + Filename: "coolfile365", + CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC), + }, key) + + _, err = objectstore.ParseBucketKey("AAA/import/99999999/coolfile") + assert.Error(t, err) + + _, err = objectstore.ParseBucketKey("AAA/import/20250331/99999999/coolfile") + assert.Error(t, err) + + _, err = objectstore.ParseBucketKey("AAA/import/20250331/999/coolfile") + assert.Error(t, err) + + _, err = objectstore.ParseBucketKey("AAA/import/20250331/1/cool**file") + assert.Error(t, err) +} diff --git a/internal/test/ecosystem.go b/internal/test/ecosystem.go index 5b9b6198..1c26ff5e 100644 --- a/internal/test/ecosystem.go +++ b/internal/test/ecosystem.go @@ -10,7 +10,6 @@ import ( "io" "log/slog" "net/http" - "regexp" "testing" "time" @@ -279,10 +278,9 @@ func CreateMockExpectation(t testing.TB, server MockServer, expectation MockExpe } type StartDocTextDetectionExpectationParams struct { - Bucket string - ClientID string - Filename string - JobID string + Bucket string + JobID string + Key objectstore.BucketKey } func CreateStartDocTextDetectionExpectation(t testing.TB, mockServer MockServer, params StartDocTextDetectionExpectationParams) MockExpectation { @@ -297,7 +295,7 @@ func CreateStartDocTextDetectionExpectation(t testing.TB, mockServer MockServer, "DocumentLocation": { "S3Object": map[string]any{ "Bucket": params.Bucket, - "Name": fmt.Sprintf("%s/import/%s", params.ClientID, params.Filename), + "Name": params.Key.String(), }, }, "OutputConfig": { @@ -434,17 +432,16 @@ func WaitForMockTextractOutput(t testing.TB, ctx context.Context, cfg objectstor bucket, ok := outputConfig["S3Bucket"].(string) require.True(t, ok) - key, ok := outputConfig["S3Prefix"].(string) + keyStr, ok := outputConfig["S3Prefix"].(string) require.True(t, ok) - matches := regexp.MustCompile(`^(.+)/text/textract/(.+)$`). - FindStringSubmatch(key) - require.Len(t, matches, 3) + key, err := objectstore.ParseBucketKey(keyStr) + require.NoError(t, err) + require.Equal(t, objectstore.TextTextract, key.Location) - PutTextractObject(t, ctx, cfg, PutObjectParams{ - ClientId: matches[1], - Bucket: bucket, - Filepath: matches[2], - File: params.File, + PutObject(t, ctx, cfg, PutObjectParams{ + Key: key, + Bucket: bucket, + File: params.File, }) } diff --git a/internal/test/ecosystem_test.go b/internal/test/ecosystem_test.go index 4d41108f..03cf5844 100644 --- a/internal/test/ecosystem_test.go +++ b/internal/test/ecosystem_test.go @@ -8,6 +8,7 @@ import ( "net/http/httptest" "strings" "testing" + "time" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/serviceconfig/aws" @@ -174,14 +175,27 @@ func TestWaitForTextractMockEndpoint(t *testing.T) { deps, cleanup := CreateFullDependencies(t, ctx, cfg) defer cleanup() + clientId := "CLIENTID" + importKey := objectstore.BucketKey{ + ClientID: clientId, + CreatedAt: time.Now().UTC(), + Location: objectstore.Import, + Filename: "input", + } + outputKey := objectstore.BucketKey{ + ClientID: clientId, + CreatedAt: time.Now().UTC(), + Location: objectstore.TextTextract, + Filename: "output", + } + e := CreateStartDocTextDetectionExpectation(t, deps.MockServer, StartDocTextDetectionExpectationParams{ - Bucket: deps.BucketName, - ClientID: "hi", - Filename: "hi", - JobID: "hi", + Bucket: deps.BucketName, + Key: importKey, + JobID: "hi", }) - body := strings.NewReader(fmt.Sprintf(`{"DocumentLocation":{"S3Object":{"Bucket":"%s","Name":"hi/import/hi"}},"OutputConfig":{"S3Bucket":"%s","S3Prefix":"hi/text/textract/hi"}}`, deps.BucketName, deps.BucketName)) + body := strings.NewReader(fmt.Sprintf(`{"DocumentLocation":{"S3Object":{"Bucket":"%s","Name":"%s"}},"OutputConfig":{"S3Bucket":"%s","S3Prefix":"%s"}}`, deps.BucketName, importKey.String(), deps.BucketName, outputKey.String())) req, err := http.NewRequest("POST", string(deps.MockServer.External), body) require.NoError(t, err) req.Header.Add("X-Amz-Target", "Textract.StartDocumentTextDetection") diff --git a/internal/test/objectstore.go b/internal/test/objectstore.go index c95d2bac..994bb410 100644 --- a/internal/test/objectstore.go +++ b/internal/test/objectstore.go @@ -2,7 +2,6 @@ package test import ( "context" - "fmt" "io" "log/slog" "testing" @@ -55,29 +54,18 @@ func SetStoreClient(t testing.TB, ctx context.Context, cfg objectstore.ConfigPro } type PutObjectParams struct { - ClientId string - Bucket string - Filepath string - File io.Reader + Key objectstore.BucketKey + Bucket string + File io.Reader } func PutObject(t testing.TB, ctx context.Context, cfg objectstore.ConfigProvider, params PutObjectParams) { - location := fmt.Sprintf("%s/%s", params.ClientId, params.Filepath) + key := params.Key.String() _, err := cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{ Bucket: ¶ms.Bucket, - Key: &location, + Key: &key, Body: params.File, }) require.NoError(t, err) - slog.Info("put object", "bucket", params.Bucket, "key", location) -} - -func PutImportObject(t testing.TB, ctx context.Context, cfg objectstore.ConfigProvider, params PutObjectParams) { - params.Filepath = fmt.Sprintf("import/%s", params.Filepath) - PutObject(t, ctx, cfg, params) -} - -func PutTextractObject(t testing.TB, ctx context.Context, cfg objectstore.ConfigProvider, params PutObjectParams) { - params.Filepath = fmt.Sprintf("text/textract/%s", params.Filepath) - PutObject(t, ctx, cfg, params) + slog.Info("put object", "bucket", params.Bucket, "key", key) } diff --git a/internal/test/objectstore_test.go b/internal/test/objectstore_test.go index 2291bc1e..b10cb182 100644 --- a/internal/test/objectstore_test.go +++ b/internal/test/objectstore_test.go @@ -2,7 +2,6 @@ package test import ( "context" - "fmt" "strings" "testing" @@ -63,90 +62,26 @@ func TestPutObject(t *testing.T) { mockS3 := objectstoremock.NewMockS3Client(t) cfg.StoreClient = mockS3 - clientId := "byebye" bucketName := "bucket" file := strings.NewReader("hello") - filename := "hello" - expectedLocation := fmt.Sprintf("%s/%s", clientId, filename) + key := objectstore.BucketKey{ + ClientID: "byebye", + Filename: "hello", + } mockS3.EXPECT(). PutObject( mock.Anything, mock.MatchedBy(func(in *s3.PutObjectInput) bool { - return *in.Bucket == bucketName && *in.Key == expectedLocation + return *in.Bucket == bucketName && *in.Key == key.String() }), mock.Anything, ). Return(&s3.PutObjectOutput{}, nil) PutObject(t, ctx, cfg, PutObjectParams{ - ClientId: clientId, - Bucket: bucketName, - Filepath: filename, - File: file, - }) -} - -func TestPutImportObject(t *testing.T) { - ctx := context.Background() - - cfg := &StoreConfig{} - SetCfgProvider(t, cfg) - mockS3 := objectstoremock.NewMockS3Client(t) - cfg.StoreClient = mockS3 - - clientId := "byebye" - bucketName := "bucket" - file := strings.NewReader("hello") - filename := "hello" - expectedLocation := fmt.Sprintf("%s/import/%s", clientId, filename) - - mockS3.EXPECT(). - PutObject( - mock.Anything, - mock.MatchedBy(func(in *s3.PutObjectInput) bool { - return *in.Bucket == bucketName && *in.Key == expectedLocation - }), - mock.Anything, - ). - Return(&s3.PutObjectOutput{}, nil) - - PutImportObject(t, ctx, cfg, PutObjectParams{ - ClientId: clientId, - Bucket: bucketName, - Filepath: filename, - File: file, - }) -} - -func TestPutTextractObject(t *testing.T) { - ctx := context.Background() - - cfg := &StoreConfig{} - SetCfgProvider(t, cfg) - mockS3 := objectstoremock.NewMockS3Client(t) - cfg.StoreClient = mockS3 - - clientId := "byebye" - bucketName := "bucket" - file := strings.NewReader("hello") - filename := "hello" - expectedLocation := fmt.Sprintf("%s/text/textract/%s", clientId, filename) - - mockS3.EXPECT(). - PutObject( - mock.Anything, - mock.MatchedBy(func(in *s3.PutObjectInput) bool { - return *in.Bucket == bucketName && *in.Key == expectedLocation - }), - mock.Anything, - ). - Return(&s3.PutObjectOutput{}, nil) - - PutTextractObject(t, ctx, cfg, PutObjectParams{ - ClientId: clientId, - Bucket: bucketName, - Filepath: filename, - File: file, + Bucket: bucketName, + File: file, + Key: key, }) } diff --git a/test/process_test.go b/test/process_test.go index 2ab86191..d255915a 100644 --- a/test/process_test.go +++ b/test/process_test.go @@ -5,6 +5,7 @@ import ( "net/http" "strings" "testing" + "time" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/serviceconfig/aws" @@ -65,18 +66,26 @@ func TestProcess(t *testing.T) { filename := "objectname" expectation := test.CreateStartDocTextDetectionExpectation(t, net.Dependencies.MockServer, test.StartDocTextDetectionExpectationParams{ - Bucket: net.Dependencies.BucketName, - ClientID: client.Id, - Filename: filename, - JobID: "textractId", + Bucket: net.Dependencies.BucketName, + JobID: "textractId", + Key: objectstore.BucketKey{ + ClientID: client.Id, + Filename: filename, + Location: objectstore.Import, + CreatedAt: time.Now().UTC(), + }, }) body := strings.NewReader(pdfHelloWorld) - test.PutImportObject(t, ctx, cfg, test.PutObjectParams{ - ClientId: client.Id, - Bucket: net.Dependencies.BucketName, - Filepath: filename, - File: body, + test.PutObject(t, ctx, cfg, test.PutObjectParams{ + Bucket: net.Dependencies.BucketName, + File: body, + Key: objectstore.BucketKey{ + ClientID: client.Id, + Filename: filename, + Location: objectstore.Import, + CreatedAt: time.Now().UTC(), + }, }) queryapitest.WaitForClientStatus(t, ctx, net.Client, client.Id, queryapi.NOTSYNCED)