jacob-mathison/query-orchestration
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Merged in feature/clean (pull request #76)

Browse Source 
Clean Set Up

* fail

* starteddb

* constraint

* cleantest

* fixqueries

* fixtests

* storemimetype

* buffer

* tests

* setup

* passingtests

* pdfHElloWorld

* rm

* test

* notodos

* tests

* clean

* testpdf
This commit is contained in:
Michael McGuinness
2025-02-28 13:11:53 +00:00
parent a49084ebce
commit f11f4def43
198 changed files with 71722 additions and 146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/document/clean/clean.go
+92 -10
View File
@@ -2,36 +2,100 @@ package documentclean
import (
"context"
"fmt"
"log/slog"
"queryorchestration/internal/database"
"queryorchestration/internal/database/repository"
"queryorchestration/internal/document"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/google/uuid"
)
type CleanParams struct {
ID uuid.UUID
Hash string
Location document.Location
}
func (s *Service) executeCleanTasks(params *CleanParams) (*document.Location, error) {
// TODO - various cleaning tasks
return &params.Location, nil
type InvalidDocumentReason string
const (
InvalidDocumentMimeType InvalidDocumentReason = "invalid_mimetype"
InvalidDocumentRead InvalidDocumentReason = "invalid_read"
InvalidDocumentReadPages InvalidDocumentReason = "invalid_read_pages"
InvalidDocumentZeroPageCount InvalidDocumentReason = "zero_page_count"
InvalidDocumentLargePageCount InvalidDocumentReason = "large_page_count"
InvalidDocumentLargeFile InvalidDocumentReason = "large_file"
InvalidDocumentQuality InvalidDocumentReason = "invalid_quality"
)
type ExecuteCleanResponse struct {
location *document.Location
mimetype *MimeType
failReason *InvalidDocumentReason
}
func (s *Service) executeCleanTasks(ctx context.Context, params *CleanParams) (*ExecuteCleanResponse, error) {
out, err := s.cfg.GetStoreClient().HeadObject(ctx, &s3.HeadObjectInput{
Bucket: &params.Location.Bucket,
Key: &params.Location.Key,
IfMatch: &params.Hash,
})
if err != nil {
return nil, err
}
length := int64(1)
if out.ContentLength != nil {
length = *out.ContentLength
}
mimeType, err := s.getAcceptedMimeType(ctx, params, out.ContentType, length)
if err != nil {
return nil, err
} else if mimeType == MimeTypeInvalid {
reason := InvalidDocumentMimeType
return &ExecuteCleanResponse{
failReason: &reason,
}, nil
}
content, err := s.getContent(ctx, params, mimeType)
if err != nil {
return nil, err
}
corruptReason := content.isCorrupted(ctx)
if corruptReason != nil {
return &ExecuteCleanResponse{
failReason: corruptReason,
}, nil
}
return &ExecuteCleanResponse{
location: &params.Location,
mimetype: &mimeType,
}, nil
}
func (s *Service) clean(ctx context.Context, id uuid.UUID) error {
slog.Debug("cleaning document", "id", id.String())
docId := database.MustToDBUUID(id)
doc, err := s.cfg.GetDBQueries().GetDocument(ctx, docId)
if err != nil {
return err
}
entry, err := s.cfg.GetDBQueries().GetDocumentEntry(ctx, docId)
if err != nil {
return err
}
outLocation, err := s.executeCleanTasks(&CleanParams{
ID: id,
out, err := s.executeCleanTasks(ctx, &CleanParams{
ID: id,
Hash: doc.Hash,
Location: document.Location{
Bucket: entry.Bucket,
Key: entry.Key,
@@ -43,15 +107,33 @@ func (s *Service) clean(ctx context.Context, id uuid.UUID) error {
version := s.svc.Version.GetVersion()
err = s.cfg.GetDBQueries().AddDocumentCleanEntry(ctx, &repository.AddDocumentCleanEntryParams{
params := &repository.AddDocumentCleanEntryParams{
Documentid: docId,
Version: version,
Bucket: outLocation.Bucket,
Key: outLocation.Key,
})
}
if out.failReason != nil {
slog.Info("Failed document", "id", id, "reason", *out.failReason)
params.Fail = repository.NullCleanfailtype{
Cleanfailtype: repository.Cleanfailtype(*out.failReason),
Valid: true,
}
} else {
params.Bucket = &out.location.Bucket
params.Key = &out.location.Key
params.Mimetype = repository.NullCleanmimetypes{
Cleanmimetypes: repository.Cleanmimetypes(*out.mimetype),
Valid: true,
}
}
err = s.cfg.GetDBQueries().AddDocumentCleanEntry(ctx, params)
if err != nil {
return err
}
if out.failReason != nil {
return fmt.Errorf("%s", *out.failReason)
}
return nil
}
internal/document/clean/clean_test.go
+169 -18
View File
@@ -2,15 +2,20 @@ package documentclean
import (
"context"
"io"
"queryorchestration/internal/database"
"queryorchestration/internal/database/repository"
"queryorchestration/internal/document"
cleanversion "queryorchestration/internal/document/clean/version"
objectstoremock "queryorchestration/mocks/objectstore"
"strings"
"testing"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/google/uuid"
"github.com/pashagolub/pgxmock/v3"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
)
func TestClean(t *testing.T) {
@@ -23,6 +28,8 @@ func TestClean(t *testing.T) {
cfg := &DocCleanConfig{}
cfg.DBPool = pool
cfg.DBQueries = repository.New(pool)
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
@@ -31,33 +38,177 @@ func TestClean(t *testing.T) {
},
}
doc := document.Document{
ID: uuid.New(),
JobID: uuid.New(),
Hash: "example_hash",
}
inloc := document.Location{
Bucket: "bucket_name",
Key: "/i/am/here",
}
dbid := database.MustToDBUUID(doc.ID)
pool.ExpectQuery("name: GetDocument :one").WithArgs(dbid).
WillReturnRows(
pgxmock.NewRows([]string{"id", "jobId", "hash"}).
AddRow(dbid, database.MustToDBUUID(doc.JobID), doc.Hash),
)
pool.ExpectQuery("name: GetDocumentEntry :one").WithArgs(dbid).WillReturnRows(
pgxmock.NewRows([]string{"documentId", "bucket", "key"}).
AddRow(dbid, inloc.Bucket, inloc.Key),
)
mimeType := "application/pdf"
dbmimetype := repository.NullCleanmimetypes{
Valid: true,
Cleanmimetypes: repository.Cleanmimetypes(mimeType),
}
pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(dbid, int32(1), &inloc.Bucket, &inloc.Key, dbmimetype, repository.NullCleanfailtype{}).
WillReturnResult(pgxmock.NewResult("", 1))
var length int64 = 10
mockS3.EXPECT().
HeadObject(
mock.Anything,
mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
return *in.Bucket == inloc.Bucket && *in.Key == inloc.Key
}),
mock.Anything,
).
Return(&s3.HeadObjectOutput{
ContentType: &mimeType,
ContentLength: &length,
}, nil)
body := io.NopCloser(strings.NewReader(pdfHelloWorld))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == inloc.Bucket && *in.Key == inloc.Key
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
err = svc.clean(ctx, doc.ID)
assert.NoError(t, err)
}
func TestExecuteCleanTasks(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
id := uuid.New()
hash := "example_hash"
inloc := document.Location{
Bucket: "bucket_name",
Key: "/i/am/here",
}
pool.ExpectQuery("name: GetDocumentEntry :one").WithArgs(database.MustToDBUUID(id)).WillReturnRows(
pgxmock.NewRows([]string{"documentId", "bucket", "key"}).
AddRow(database.MustToDBUUID(id), inloc.Bucket, inloc.Key),
)
pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(database.MustToDBUUID(id), int32(1), inloc.Bucket, inloc.Key).
WillReturnResult(pgxmock.NewResult("", 1))
mimeType := "application/pdf"
mockS3.EXPECT().
HeadObject(
mock.Anything,
mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
return *in.Bucket == inloc.Bucket && *in.Key == inloc.Key
}),
mock.Anything,
).
Return(&s3.HeadObjectOutput{
ContentType: &mimeType,
}, nil)
err = svc.clean(ctx, id)
assert.NoError(t, err)
}
body := io.NopCloser(strings.NewReader(pdfHelloWorld))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == inloc.Bucket && *in.Key == inloc.Key
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
func TestExecuteCleanTasks(t *testing.T) {
svc := Service{}
id := uuid.New()
location := document.Location{}
outloc, err := svc.executeCleanTasks(&CleanParams{
outloc, err := svc.executeCleanTasks(ctx, &CleanParams{
ID: id,
Location: location,
Hash: hash,
Location: inloc,
})
assert.NoError(t, err)
assert.EqualExportedValues(t, location, *outloc)
assert.EqualExportedValues(t, ExecuteCleanResponse{
location: &inloc,
}, *outloc)
}
const pdfHelloWorld = `%PDF-1.4
%����
1 0 obj
<<
/Type /Catalog
/Pages 2 0 R
/Version /1.4
>>
endobj
2 0 obj
<<
/Type /Pages
/Kids [3 0 R]
/Count 1
>>
endobj
3 0 obj
<<
/Type /Page
/Parent 2 0 R
/MediaBox [0 0 612 792]
/Resources <<
/Font <<
/F1 <<
/Type /Font
/Subtype /Type1
/BaseFont /Helvetica
>>
>>
>>
/Contents 4 0 R
>>
endobj
4 0 obj
<<
/Length
44
>>
stream
BT
/F1 24 Tf
100 700 Td
(Hello World!) Tj
ET
endstream
endobj
xref
0 5
0000000000 65535 f
0000000015 00000 n
0000000086 00000 n
0000000151 00000 n
0000000376 00000 n
trailer
<<
/Size 5
/Root 1 0 R
>>
startxref
472
%%EOF`
internal/document/clean/content.go
+52
View File
@@ -0,0 +1,52 @@
package documentclean
import (
"bytes"
"context"
"errors"
"io"
"github.com/aws/aws-sdk-go-v2/service/s3"
)
type Content interface {
isCorrupted(context.Context) *InvalidDocumentReason
}
func (s *Service) getContent(ctx context.Context, params *CleanParams, mimeType MimeType) (Content, error) {
resp, err := s.cfg.GetStoreClient().GetObject(ctx, &s3.GetObjectInput{
Bucket: &params.Location.Bucket,
Key: &params.Location.Key,
IfMatch: &params.Hash,
})
if err != nil {
return nil, err
}
defer resp.Body.Close()
seek, err := ReaderToSeeker(resp.Body)
if err != nil {
return nil, err
}
var content Content
switch mimeType {
case MimeTypePDF:
content = NewPDF(seek)
default:
return nil, errors.New("invalid mimetype")
}
return content, nil
}
func ReaderToSeeker(readCloser io.ReadCloser) (io.ReadSeeker, error) {
data, err := io.ReadAll(readCloser)
if err != nil {
return nil, err
}
readCloser.Close()
return bytes.NewReader(data), nil
}
internal/document/clean/contentType.go
+110
View File
@@ -0,0 +1,110 @@
package documentclean
import (
"bytes"
"context"
"fmt"
"github.com/aws/aws-sdk-go-v2/service/s3"
)
type MimeType string
const (
MimeTypePDF MimeType = "application/pdf"
MimeTypeBinaryOctetStream MimeType = "binary/octet-stream"
MimeTypeInvalid MimeType = "invalid"
)
var AcceptMimeTypes = []MimeType{
MimeTypePDF,
}
const (
PDFSignatureStr = "%PDF-"
)
const (
PDFSignatureEndStr = "%%EOF"
)
var (
PDFSignature = []byte(PDFSignatureStr)
)
var (
PDFSignatureEnd = []byte(PDFSignatureEndStr)
)
func (s *Service) getAcceptedMimeType(ctx context.Context, params *CleanParams, contentType *string, length int64) (MimeType, error) {
mimeType := s.getMetadataMimeType(contentType)
if mimeType != MimeTypeInvalid && mimeType != MimeTypeBinaryOctetStream {
return mimeType, nil
}
return s.getBodyMimeType(ctx, params, length)
}
func (s *Service) getMetadataMimeType(contentType *string) MimeType {
if contentType == nil {
return MimeTypeInvalid
}
parsedType := MimeType(*contentType)
for _, accepted := range AcceptMimeTypes {
if accepted == parsedType {
return accepted
}
}
return MimeTypeInvalid
}
func (s *Service) getBodyMimeType(ctx context.Context, params *CleanParams, length int64) (MimeType, error) {
startBuffer, err := s.getBytesBuffer(ctx, params, 0, 5)
if err != nil {
return MimeTypeInvalid, err
}
endBuffer, err := s.getBytesBuffer(ctx, params, length-6, length)
if err != nil {
return MimeTypeInvalid, err
}
prefixMimeType := MimeTypeInvalid
if bytes.HasPrefix(startBuffer, PDFSignature) && bytes.HasSuffix(endBuffer, PDFSignatureEnd) {
prefixMimeType = MimeTypePDF
}
return prefixMimeType, nil
}
func (s *Service) getBytesBuffer(ctx context.Context, params *CleanParams, start int64, length int64) ([]byte, error) {
if start < 0 {
start = 0
}
if length < 1 {
length = 1
}
byteRange := fmt.Sprintf("bytes=%d-%d", start, length-1)
out, err := s.cfg.GetStoreClient().GetObject(ctx, &s3.GetObjectInput{
Bucket: &params.Location.Bucket,
Key: &params.Location.Key,
IfMatch: &params.Hash,
Range: &byteRange,
})
if err != nil {
return nil, err
}
defer out.Body.Close()
buffer := make([]byte, length)
n, err := out.Body.Read(buffer)
if err != nil && n == 0 {
return nil, fmt.Errorf("unable to read object body: %v", err)
}
return buffer[:n], nil
}
internal/document/clean/contentType_test.go
+459
View File
@@ -0,0 +1,459 @@
package documentclean
import (
"context"
"fmt"
"io"
"queryorchestration/internal/document"
objectstoremock "queryorchestration/mocks/objectstore"
"strings"
"testing"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
)
func TestGetBodyMimeType(t *testing.T) {
t.Run("invalid type", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
params := &CleanParams{
Hash: "example_hash",
Location: document.Location{
Bucket: "buck",
Key: "Key",
},
}
bodyStr := "invalid"
body := io.NopCloser(strings.NewReader(bodyStr))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
contentType, err := svc.getBodyMimeType(ctx, params, int64(len(bodyStr)))
assert.NoError(t, err)
assert.Equal(t, MimeTypeInvalid, contentType)
})
t.Run("pdf", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
params := &CleanParams{
Hash: "example_hash",
Location: document.Location{
Bucket: "buck",
Key: "Key",
},
}
bodyStr := fmt.Sprintf("%s %s", PDFSignatureStr, PDFSignatureEndStr)
body := io.NopCloser(strings.NewReader(bodyStr))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
contentType, err := svc.getBodyMimeType(ctx, params, int64(len(bodyStr)))
assert.NoError(t, err)
assert.Equal(t, MimeTypePDF, contentType)
})
t.Run("pdf start", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
params := &CleanParams{
Hash: "example_hash",
Location: document.Location{
Bucket: "buck",
Key: "Key",
},
}
bodyStr := fmt.Sprintf("%s---------------", PDFSignatureStr)
body := io.NopCloser(strings.NewReader(bodyStr))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
contentType, err := svc.getBodyMimeType(ctx, params, int64(len(bodyStr)))
assert.NoError(t, err)
assert.Equal(t, MimeTypeInvalid, contentType)
})
t.Run("pdf end", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
params := &CleanParams{
Hash: "example_hash",
Location: document.Location{
Bucket: "buck",
Key: "Key",
},
}
bodyStr := fmt.Sprintf("%s---------------", PDFSignatureEndStr)
body := io.NopCloser(strings.NewReader(bodyStr))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
contentType, err := svc.getBodyMimeType(ctx, params, int64(len(bodyStr)))
assert.NoError(t, err)
assert.Equal(t, MimeTypeInvalid, contentType)
})
t.Run("short", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
params := &CleanParams{
Hash: "example_hash",
Location: document.Location{
Bucket: "buck",
Key: "Key",
},
}
bodyStr := ""
body := io.NopCloser(strings.NewReader(bodyStr))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
_, err := svc.getBodyMimeType(ctx, params, int64(len(bodyStr)))
assert.Error(t, err)
})
}
func TestGetMetadataMimeType(t *testing.T) {
t.Run("nil type", func(t *testing.T) {
svc := Service{}
contentType := svc.getMetadataMimeType(nil)
assert.Equal(t, MimeTypeInvalid, contentType)
})
t.Run("invalid type", func(t *testing.T) {
svc := Service{}
inType := "invalid_type"
contentType := svc.getMetadataMimeType(&inType)
assert.Equal(t, MimeTypeInvalid, contentType)
})
t.Run("pdf", func(t *testing.T) {
svc := Service{}
inType := "application/pdf"
contentType := svc.getMetadataMimeType(&inType)
assert.Equal(t, MimeTypePDF, contentType)
})
}
func TestGetAcceptedMimeType(t *testing.T) {
t.Run("invalid type", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
params := &CleanParams{
Location: document.Location{
Bucket: "buck",
Key: "Key",
},
}
bodyStr := "invalid"
body := io.NopCloser(strings.NewReader(bodyStr))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
inType := "invalid_type"
contentType, err := svc.getAcceptedMimeType(ctx, params, &inType, int64(len(bodyStr)))
assert.NoError(t, err)
assert.Equal(t, MimeTypeInvalid, contentType)
})
t.Run("pdf mimetype", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
inType := "application/pdf"
contentType, err := svc.getAcceptedMimeType(ctx, &CleanParams{}, &inType, 0)
assert.NoError(t, err)
assert.Equal(t, MimeTypePDF, contentType)
})
t.Run("pdf format", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
params := &CleanParams{
Location: document.Location{
Bucket: "buck",
Key: "Key",
},
}
bodyStr := fmt.Sprintf("%s %s", PDFSignatureStr, PDFSignatureEndStr)
body := io.NopCloser(strings.NewReader(bodyStr))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
inType := "invalid_type"
contentType, err := svc.getAcceptedMimeType(ctx, params, &inType, int64(len(bodyStr)))
assert.NoError(t, err)
assert.Equal(t, MimeTypePDF, contentType)
})
}
func TestGetBytesBuffer(t *testing.T) {
t.Run("valid", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
params := &CleanParams{
Hash: "example_hash",
Location: document.Location{
Bucket: "buck",
Key: "Key",
},
}
start := int64(0)
bodyStr := "invalid"
length := int64(len(bodyStr))
body := io.NopCloser(strings.NewReader(bodyStr))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash && *in.Range == "bytes=0-6"
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
buffer, err := svc.getBytesBuffer(ctx, params, start, length)
assert.NoError(t, err)
assert.NotNil(t, buffer)
})
t.Run("negative start", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
params := &CleanParams{
Hash: "example_hash",
Location: document.Location{
Bucket: "buck",
Key: "Key",
},
}
start := int64(-1)
bodyStr := "invalid"
length := int64(len(bodyStr))
body := io.NopCloser(strings.NewReader(bodyStr))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash && *in.Range == "bytes=0-6"
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
buffer, err := svc.getBytesBuffer(ctx, params, start, length)
assert.NoError(t, err)
assert.NotNil(t, buffer)
})
t.Run("zero length", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
}
params := &CleanParams{
Hash: "example_hash",
Location: document.Location{
Bucket: "buck",
Key: "Key",
},
}
start := int64(0)
bodyStr := ""
length := int64(0)
body := io.NopCloser(strings.NewReader(bodyStr))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == params.Location.Bucket && *in.Key == params.Location.Key && *in.IfMatch == params.Hash && *in.Range == "bytes=0-0"
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
buffer, err := svc.getBytesBuffer(ctx, params, start, length)
assert.Error(t, err)
assert.Nil(t, buffer)
})
}
internal/document/clean/content_test.go
+59
View File
@@ -0,0 +1,59 @@
package documentclean
import (
"context"
"io"
"queryorchestration/internal/document"
cleanversion "queryorchestration/internal/document/clean/version"
objectstoremock "queryorchestration/mocks/objectstore"
"strings"
"testing"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
)
func TestGetContent(t *testing.T) {
t.Run("Working", func(t *testing.T) {
ctx := context.Background()
cfg := &DocCleanConfig{}
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
svc: &Services{
Version: cleanversion.New(cfg),
},
}
inloc := document.Location{
Bucket: "bucket_name",
Key: "/i/am/here",
}
params := &CleanParams{
Hash: "hash",
Location: inloc,
}
mimeType := MimeTypePDF
body := io.NopCloser(strings.NewReader(pdfHelloWorld))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == inloc.Bucket && *in.Key == inloc.Key
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
content, err := svc.getContent(ctx, params, mimeType)
assert.NoError(t, err)
assert.NotNil(t, content)
})
}
internal/document/clean/create.go
+1 -1
View File
@@ -10,7 +10,7 @@ import (
)
func (s *Service) Clean(ctx context.Context, id uuid.UUID) error {
isclean, err := s.cfg.GetDBQueries().IsDocumentClean(ctx, database.MustToDBUUID(id))
isclean, err := s.cfg.GetDBQueries().HasDocumentCleanEntry(ctx, database.MustToDBUUID(id))
if err != nil {
return err
}
internal/document/clean/create_test.go
+55 -7
View File
@@ -3,15 +3,20 @@ package documentclean
import (
"context"
"fmt"
"io"
"queryorchestration/internal/database"
"queryorchestration/internal/database/repository"
"queryorchestration/internal/document"
cleanversion "queryorchestration/internal/document/clean/version"
"queryorchestration/internal/serviceconfig"
"queryorchestration/internal/serviceconfig/objectstore"
"queryorchestration/internal/serviceconfig/queue/documenttext"
objectstoremock "queryorchestration/mocks/objectstore"
queuemock "queryorchestration/mocks/queue"
"strings"
"testing"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/aws/aws-sdk-go-v2/service/sqs"
"github.com/google/uuid"
"github.com/pashagolub/pgxmock/v3"
@@ -22,6 +27,7 @@ import (
type DocCleanConfig struct {
serviceconfig.BaseConfig
documenttext.DocTextConfig
objectstore.ObjectStoreConfig
}
func TestCreate(t *testing.T) {
@@ -38,6 +44,8 @@ func TestCreate(t *testing.T) {
cfg.DocumentTextURL = "/i/am/here"
cfg.DBPool = pool
cfg.DBQueries = repository.New(pool)
mockS3 := objectstoremock.NewMockS3Client(t)
cfg.StoreClient = mockS3
svc := Service{
cfg: cfg,
@@ -46,34 +54,74 @@ func TestCreate(t *testing.T) {
},
}
id := uuid.New()
doc := document.Document{
ID: uuid.New(),
JobID: uuid.New(),
Hash: "example_hash",
}
inloc := document.Location{
Bucket: "bucket_name",
Key: "/i/am/here",
}
dbid := database.MustToDBUUID(doc.ID)
pool.ExpectQuery("name: IsDocumentClean :one").WithArgs(database.MustToDBUUID(id)).WillReturnRows(
pool.ExpectQuery("name: HasDocumentCleanEntry :one").WithArgs(dbid).WillReturnRows(
pgxmock.NewRows([]string{"isclean"}).
AddRow(false),
)
pool.ExpectQuery("name: GetDocumentEntry :one").WithArgs(database.MustToDBUUID(id)).WillReturnRows(
pool.ExpectQuery("name: GetDocument :one").WithArgs(dbid).
WillReturnRows(
pgxmock.NewRows([]string{"id", "jobId", "hash"}).
AddRow(dbid, database.MustToDBUUID(doc.JobID), doc.Hash),
)
pool.ExpectQuery("name: GetDocumentEntry :one").WithArgs(dbid).WillReturnRows(
pgxmock.NewRows([]string{"documentId", "bucket", "key"}).
AddRow(database.MustToDBUUID(id), inloc.Bucket, inloc.Key),
AddRow(dbid, inloc.Bucket, inloc.Key),
)
pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(database.MustToDBUUID(id), int32(1), inloc.Bucket, inloc.Key).
mimeType := "application/pdf"
dbmimetype := repository.NullCleanmimetypes{
Valid: true,
Cleanmimetypes: repository.Cleanmimetypes(mimeType),
}
pool.ExpectExec("name: AddDocumentCleanEntry :exec").WithArgs(dbid, int32(1), &inloc.Bucket, &inloc.Key, dbmimetype, repository.NullCleanfailtype{}).
WillReturnResult(pgxmock.NewResult("", 1))
mockSQS.EXPECT().
SendMessage(
mock.Anything,
mock.MatchedBy(func(in *sqs.SendMessageInput) bool {
return *in.QueueUrl == cfg.DocumentTextURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", id)
return *in.QueueUrl == cfg.DocumentTextURL && *in.MessageBody == fmt.Sprintf("{\"id\":\"%s\"}", doc.ID)
}),
mock.Anything,
).
Return(&sqs.SendMessageOutput{}, nil)
err = svc.Clean(ctx, id)
mockS3.EXPECT().
HeadObject(
mock.Anything,
mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
return *in.Bucket == inloc.Bucket && *in.Key == inloc.Key
}),
mock.Anything,
).
Return(&s3.HeadObjectOutput{
ContentType: &mimeType,
}, nil)
body := io.NopCloser(strings.NewReader(pdfHelloWorld))
mockS3.EXPECT().
GetObject(
mock.Anything,
mock.MatchedBy(func(in *s3.GetObjectInput) bool {
return *in.Bucket == inloc.Bucket && *in.Key == inloc.Key
}),
mock.Anything,
).
Return(&s3.GetObjectOutput{
Body: body,
}, nil)
err = svc.Clean(ctx, doc.ID)
assert.NoError(t, err)
}
internal/document/clean/pdf.go
+54
View File
@@ -0,0 +1,54 @@
package documentclean
import (
"context"
"io"
"github.com/pdfcpu/pdfcpu/pkg/pdfcpu"
"github.com/pdfcpu/pdfcpu/pkg/pdfcpu/model"
)
type PDF struct {
pdf io.ReadSeeker
config *model.Configuration
maxPages int
maxBytes int64
}
const (
TEXTRACT_SYNC_LIMIT = 10 * 1024 * 1024
TEXTRACT_ASYNC_LIMIT = 500 * 1024 * 1024
TEXTRACT_PAGE_LIMIT = 3000
)
func NewPDF(buf io.ReadSeeker) *PDF {
config := &model.Configuration{
ValidationMode: model.ValidationRelaxed,
Reader15: true,
Offline: true,
}
return &PDF{
pdf: buf,
config: config,
maxPages: TEXTRACT_PAGE_LIMIT,
maxBytes: TEXTRACT_SYNC_LIMIT,
}
}
func (s *PDF) isCorrupted(ctx context.Context) *InvalidDocumentReason {
var reason InvalidDocumentReason
pdfCtx, err := pdfcpu.ReadWithContext(ctx, s.pdf, s.config)
if err != nil {
reason = InvalidDocumentRead
return &reason
}
if int64(pdfCtx.Read.ReadFileSize()) > s.maxBytes {
reason = InvalidDocumentLargeFile
return &reason
}
return nil
}
internal/document/clean/pdf_test.go
+43
View File
@@ -0,0 +1,43 @@
package documentclean
import (
"context"
"strings"
"testing"
"github.com/stretchr/testify/assert"
)
func TestIsCorrupt(t *testing.T) {
t.Run("valid", func(t *testing.T) {
ctx := context.Background()
pdf := NewPDF(strings.NewReader(pdfHelloWorld))
reason := pdf.isCorrupted(ctx)
assert.Nil(t, reason)
})
t.Run("no content", func(t *testing.T) {
ctx := context.Background()
pdf := NewPDF(strings.NewReader("%PDF-1.4 %%EOF"))
reason := pdf.isCorrupted(ctx)
assert.Equal(t, InvalidDocumentRead, *reason)
})
t.Run("large file size", func(t *testing.T) {
ctx := context.Background()
pdf := NewPDF(strings.NewReader(pdfHelloWorld))
pdf.maxBytes = 1
reason := pdf.isCorrupted(ctx)
assert.Equal(t, InvalidDocumentLargeFile, *reason)
})
}
func TestNewPDF(t *testing.T) {
t.Run("valid", func(t *testing.T) {
pdf := NewPDF(strings.NewReader(pdfHelloWorld))
assert.NotNil(t, pdf.pdf)
assert.NotNil(t, pdf.config)
assert.Equal(t, 3000, pdf.maxPages)
})
}
internal/document/clean/service.go
+2
View File
@@ -3,12 +3,14 @@ package documentclean
import (
cleanversion "queryorchestration/internal/document/clean/version"
"queryorchestration/internal/serviceconfig"
"queryorchestration/internal/serviceconfig/objectstore"
"queryorchestration/internal/serviceconfig/queue/documenttext"
)
type ConfigProvider interface {
serviceconfig.ConfigProvider
documenttext.ConfigProvider
objectstore.ConfigProvider
}
type Services struct {
internal/document/clean/service_test.go
+2
View File
@@ -3,6 +3,7 @@ package documentclean_test
import (
documentclean "queryorchestration/internal/document/clean"
"queryorchestration/internal/serviceconfig"
"queryorchestration/internal/serviceconfig/objectstore"
"queryorchestration/internal/serviceconfig/queue/documenttext"
"testing"
@@ -12,6 +13,7 @@ import (
type DocCleanConfig struct {
serviceconfig.BaseConfig
documenttext.DocTextConfig
objectstore.ConfigProvider
}
func TestService(t *testing.T) {
internal/document/clean/version/version.go
-1
View File
@@ -5,7 +5,6 @@ import (
)
func (s *Service) GetVersion() int32 {
// TODO - actual version
return 1
}