diff --git a/api/docTextRunner/runner_test.go b/api/docTextRunner/runner_test.go index 6d79f092..3af9b8b5 100644 --- a/api/docTextRunner/runner_test.go +++ b/api/docTextRunner/runner_test.go @@ -83,12 +83,18 @@ func TestDocCleanRunner(t *testing.T) { } pool.ExpectQuery("name: GetCleanEntryByDocId :one").WithArgs(doc.DocumentID).WillReturnRows( pgxmock.NewRows([]string{"id", "documentId", "bucket", "key", "version", "hash", "mimetype", "fail", "externalClientId"}). - AddRow(cleanId, doc.DocumentID, &bucket, &keyStr, int64(1), &hash, dbmimetype, repository.NullCleanfailtype{}, "clientId"), + AddRow(cleanId, doc.DocumentID, &bucket, &keyStr, int64(1), &hash, dbmimetype, repository.NullCleanfailtype{}, key.ClientID), ) jobId := "hello" triggerId := uuid.New() pool.ExpectBegin() - pool.ExpectQuery("name: AddDocumentTextTrigger :one").WithArgs(cleanId, pgxmock.AnyArg()).WillReturnRows( + pool.ExpectQuery("name: GetTextractOutputCurrentPart :one").WithArgs(&key.ClientID, pgxmock.AnyArg()). + WillReturnRows( + pgxmock.NewRows([]string{"part", "count"}). + AddRow(0, 2), + ) + + pool.ExpectQuery("name: AddDocumentTextTrigger :one").WithArgs(cleanId, pgxmock.AnyArg(), pgxmock.AnyArg(), uint16(0)).WillReturnRows( pgxmock.NewRows([]string{"id"}).AddRow(triggerId), ) diff --git a/internal/database/migrations/00000000000001_extensions.up.sql b/internal/database/migrations/00000000000001_extensions.up.sql index 398fcd52..25ac191b 100644 --- a/internal/database/migrations/00000000000001_extensions.up.sql +++ b/internal/database/migrations/00000000000001_extensions.up.sql @@ -27,8 +27,11 @@ CREATE FUNCTION isInVersion( RETURNS boolean as $$ BEGIN RETURN _version is null OR ( - _version >= _addedVersion AND + _version >= _addedVersion AND (_removedVersion is null OR _version < _removedVersion) ); END; -$$ LANGUAGE plpgsql; \ No newline at end of file +$$ LANGUAGE plpgsql; + +CREATE DOMAIN unsignedsmallint AS SMALLINT + CHECK (VALUE >= 0); diff --git a/internal/database/migrations/00000000000005_documents.up.sql b/internal/database/migrations/00000000000005_documents.up.sql index 9336793b..c6e97caf 100644 --- a/internal/database/migrations/00000000000005_documents.up.sql +++ b/internal/database/migrations/00000000000005_documents.up.sql @@ -54,6 +54,8 @@ CREATE TABLE documentTextTriggerExtractions ( id uuid primary key DEFAULT uuid_generate_v7(), cleanId uuid not null, version bigint not null, + createdAt timestamp not null, + part unsignedsmallint not null, textractJobId text, foreign key (cleanId) references documentCleans(id) ); @@ -62,7 +64,9 @@ CREATE TABLE documentTextExtractions ( id uuid primary key DEFAULT uuid_generate_v7(), bucket text not null, key text not null, - hash text not null + hash text not null, + createdAt timestamp not null, + part unsignedsmallint not null ); CREATE TABLE documentTextExtractionEntries ( diff --git a/internal/database/queries/text.sql b/internal/database/queries/text.sql index a90fa378..4e4b4d46 100644 --- a/internal/database/queries/text.sql +++ b/internal/database/queries/text.sql @@ -7,10 +7,65 @@ SELECT EXISTS( SELECT id FROM currentTextEntries WHERE cleanId = @cleanEntryId and hash = @hash; -- name: AddDocumentTextTrigger :one -INSERT INTO documentTextTriggerExtractions (cleanId, version) VALUES ($1, $2) returning id; +INSERT INTO documentTextTriggerExtractions + (cleanId, version, createdAt, part) VALUES + ($1, $2, $3, $4) + returning id; + +-- name: GetTextractOutputCurrentPart :one +WITH client as ( + select id from clients where id = @clientId +), +parts as ( + SELECT triggers.part + FROM client as c + JOIN documents as docs on docs.clientId = c.id + JOIN documentCleans as clean on docs.id = clean.documentId + JOIN documentTextTriggerExtractions as triggers + on triggers.cleanId = clean.id + WHERE date(triggers.createdAt) = date(@queryDate) +), +max_part as ( + SELECT COALESCE(max(part), 0) as max_part_num FROM parts +) +SELECT + COALESCE(p.part, 0)::unsignedsmallint as part, + COUNT(p.part) as count +FROM parts p +RIGHT JOIN max_part mp ON p.part = mp.max_part_num +GROUP BY p.part; + +-- name: GetTextOutCurrentPart :one +WITH client as ( + select id from clients where id = @clientId +), +parts as ( + SELECT extractions.part + FROM client as c + JOIN documents as docs on docs.clientId = c.id + JOIN documentCleans as clean on docs.id = clean.documentId + JOIN documentTextTriggerExtractions as triggers + on triggers.cleanId = clean.id + JOIN documentTextExtractionEntries as textEntries + on textEntries.triggerId = triggers.id + JOIN documentTextExtractions as extractions + on extractions.id = textEntries.textId + WHERE date(extractions.createdAt) = date(@queryDate) +), +max_part as ( + SELECT COALESCE(max(part), 0) as max_part_num FROM parts +) +SELECT + COALESCE(p.part, 0)::unsignedsmallint as part, + COUNT(p.part) as count +FROM parts p +RIGHT JOIN max_part mp ON p.part = mp.max_part_num +GROUP BY p.part; -- name: AddDocumentTextTriggerJobId :exec -UPDATE documentTextTriggerExtractions SET textractJobId = @textractJobId WHERE id = @id; +UPDATE documentTextTriggerExtractions + SET textractJobId = @textractJobId + WHERE id = @id; -- name: GetDocumentTextTrigger :one SELECT dt.id, dt.cleanId, dt.version, dt.textractJobId, dc.documentId @@ -19,10 +74,15 @@ SELECT dt.id, dt.cleanId, dt.version, dt.textractJobId, dc.documentId where dt.id = @triggerId; -- name: AddDocumentText :one -INSERT INTO documentTextExtractions (bucket, key, hash) VALUES ($1, $2, $3) returning id; +INSERT INTO documentTextExtractions + (bucket, key, hash, createdAt, part) + VALUES ($1, $2, $3, $4, $5) + returning id; -- name: AddDocumentTextEntry :exec -INSERT INTO documentTextExtractionEntries (textId, triggerId, version) VALUES ($1, $2, $3); +INSERT INTO documentTextExtractionEntries + (textId, triggerId, version) + VALUES ($1, $2, $3); -- name: GetTextEntryByDocId :one SELECT id, documentId, bucket, key, hash, cleanId, diff --git a/internal/database/repository/models.go b/internal/database/repository/models.go index 85e4a2dd..77ac336e 100644 --- a/internal/database/repository/models.go +++ b/internal/database/repository/models.go @@ -326,10 +326,12 @@ type Documententry struct { } type Documenttextextraction struct { - ID uuid.UUID `db:"id"` - Bucket string `db:"bucket"` - Key string `db:"key"` - Hash string `db:"hash"` + ID uuid.UUID `db:"id"` + Bucket string `db:"bucket"` + Key string `db:"key"` + Hash string `db:"hash"` + Createdat pgtype.Timestamp `db:"createdat"` + Part uint16 `db:"part"` } type Documenttextextractionentry struct { @@ -340,10 +342,12 @@ type Documenttextextractionentry struct { } type Documenttexttriggerextraction struct { - ID uuid.UUID `db:"id"` - Cleanid uuid.UUID `db:"cleanid"` - Version int64 `db:"version"` - Textractjobid *string `db:"textractjobid"` + ID uuid.UUID `db:"id"` + Cleanid uuid.UUID `db:"cleanid"` + Version int64 `db:"version"` + Createdat pgtype.Timestamp `db:"createdat"` + Part uint16 `db:"part"` + Textractjobid *string `db:"textractjobid"` } type Fullactivecollector struct { diff --git a/internal/database/repository/result_test.go b/internal/database/repository/result_test.go index 217b5fd2..ca3649fb 100644 --- a/internal/database/repository/result_test.go +++ b/internal/database/repository/result_test.go @@ -3,12 +3,14 @@ package repository_test import ( "context" "testing" + "time" "queryorchestration/internal/database/repository" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/test" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -98,12 +100,22 @@ func TestResults(t *testing.T) { triggerId, err := queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ Cleanid: cleanid, Version: 123, + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) textId, err := queries.AddDocumentText(ctx, &repository.AddDocumentTextParams{ Bucket: "hi", Key: "hello", Hash: "example", + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) err = queries.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{ @@ -265,12 +277,22 @@ func TestResultValues(t *testing.T) { triggerId, err := queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ Cleanid: cleanid, Version: 123, + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) textId, err := queries.AddDocumentText(ctx, &repository.AddDocumentTextParams{ Bucket: "hi", Key: "hello", Hash: "example", + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) err = queries.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{ @@ -479,11 +501,21 @@ func TestUnsyncedNoDepsQueries(t *testing.T) { triggerId, err := queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ Cleanid: cleanid, Version: 123, + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) textId, err := queries.AddDocumentText(ctx, &repository.AddDocumentTextParams{ Bucket: "hi", Key: "hello", + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) err = queries.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{ @@ -559,11 +591,21 @@ func TestUnsyncedNoDepsQueries(t *testing.T) { triggerTwoId, err := queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ Cleanid: cleantwoid, Version: 123, + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) textTwoId, err := queries.AddDocumentText(ctx, &repository.AddDocumentTextParams{ Bucket: "hi", Key: "hello", + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) err = queries.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{ diff --git a/internal/database/repository/sync_test.go b/internal/database/repository/sync_test.go index 415e9085..98a916d6 100644 --- a/internal/database/repository/sync_test.go +++ b/internal/database/repository/sync_test.go @@ -3,11 +3,13 @@ package repository_test import ( "context" "testing" + "time" "queryorchestration/internal/database/repository" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/test" + "github.com/jackc/pgx/v5/pgtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -286,12 +288,22 @@ func TestClientSync(t *testing.T) { triggerId, err := queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ Cleanid: cleanId, Version: 123, + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) textId, err := queries.AddDocumentText(ctx, &repository.AddDocumentTextParams{ Bucket: "hi", Key: "hello", Hash: "example", + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) err = queries.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{ @@ -458,11 +470,21 @@ func TestClientSync(t *testing.T) { triggerId, err := queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ Cleanid: cleanId, Version: 123, + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) textId, err := queries.AddDocumentText(ctx, &repository.AddDocumentTextParams{ Bucket: "hi", Key: "hello", + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) err = queries.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{ @@ -575,11 +597,21 @@ func TestClientSync(t *testing.T) { triggerThreeId, err := queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ Cleanid: cleanthreeid, Version: 123, + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) textThreeId, err := queries.AddDocumentText(ctx, &repository.AddDocumentTextParams{ Bucket: "hi", Key: "hello", + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) err = queries.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{ diff --git a/internal/database/repository/text.sql.go b/internal/database/repository/text.sql.go index 1277937a..15c4d90e 100644 --- a/internal/database/repository/text.sql.go +++ b/internal/database/repository/text.sql.go @@ -9,30 +9,47 @@ import ( "context" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" ) const addDocumentText = `-- name: AddDocumentText :one -INSERT INTO documentTextExtractions (bucket, key, hash) VALUES ($1, $2, $3) returning id +INSERT INTO documentTextExtractions + (bucket, key, hash, createdAt, part) + VALUES ($1, $2, $3, $4, $5) + returning id ` type AddDocumentTextParams struct { - Bucket string `db:"bucket"` - Key string `db:"key"` - Hash string `db:"hash"` + Bucket string `db:"bucket"` + Key string `db:"key"` + Hash string `db:"hash"` + Createdat pgtype.Timestamp `db:"createdat"` + Part uint16 `db:"part"` } // AddDocumentText // -// INSERT INTO documentTextExtractions (bucket, key, hash) VALUES ($1, $2, $3) returning id +// INSERT INTO documentTextExtractions +// (bucket, key, hash, createdAt, part) +// VALUES ($1, $2, $3, $4, $5) +// returning id func (q *Queries) AddDocumentText(ctx context.Context, arg *AddDocumentTextParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, addDocumentText, arg.Bucket, arg.Key, arg.Hash) + row := q.db.QueryRow(ctx, addDocumentText, + arg.Bucket, + arg.Key, + arg.Hash, + arg.Createdat, + arg.Part, + ) var id uuid.UUID err := row.Scan(&id) return id, err } const addDocumentTextEntry = `-- name: AddDocumentTextEntry :exec -INSERT INTO documentTextExtractionEntries (textId, triggerId, version) VALUES ($1, $2, $3) +INSERT INTO documentTextExtractionEntries + (textId, triggerId, version) + VALUES ($1, $2, $3) ` type AddDocumentTextEntryParams struct { @@ -43,33 +60,50 @@ type AddDocumentTextEntryParams struct { // AddDocumentTextEntry // -// INSERT INTO documentTextExtractionEntries (textId, triggerId, version) VALUES ($1, $2, $3) +// INSERT INTO documentTextExtractionEntries +// (textId, triggerId, version) +// VALUES ($1, $2, $3) func (q *Queries) AddDocumentTextEntry(ctx context.Context, arg *AddDocumentTextEntryParams) error { _, err := q.db.Exec(ctx, addDocumentTextEntry, arg.Textid, arg.Triggerid, arg.Version) return err } const addDocumentTextTrigger = `-- name: AddDocumentTextTrigger :one -INSERT INTO documentTextTriggerExtractions (cleanId, version) VALUES ($1, $2) returning id +INSERT INTO documentTextTriggerExtractions + (cleanId, version, createdAt, part) VALUES + ($1, $2, $3, $4) + returning id ` type AddDocumentTextTriggerParams struct { - Cleanid uuid.UUID `db:"cleanid"` - Version int64 `db:"version"` + Cleanid uuid.UUID `db:"cleanid"` + Version int64 `db:"version"` + Createdat pgtype.Timestamp `db:"createdat"` + Part uint16 `db:"part"` } // AddDocumentTextTrigger // -// INSERT INTO documentTextTriggerExtractions (cleanId, version) VALUES ($1, $2) returning id +// INSERT INTO documentTextTriggerExtractions +// (cleanId, version, createdAt, part) VALUES +// ($1, $2, $3, $4) +// returning id func (q *Queries) AddDocumentTextTrigger(ctx context.Context, arg *AddDocumentTextTriggerParams) (uuid.UUID, error) { - row := q.db.QueryRow(ctx, addDocumentTextTrigger, arg.Cleanid, arg.Version) + row := q.db.QueryRow(ctx, addDocumentTextTrigger, + arg.Cleanid, + arg.Version, + arg.Createdat, + arg.Part, + ) var id uuid.UUID err := row.Scan(&id) return id, err } const addDocumentTextTriggerJobId = `-- name: AddDocumentTextTriggerJobId :exec -UPDATE documentTextTriggerExtractions SET textractJobId = $1 WHERE id = $2 +UPDATE documentTextTriggerExtractions + SET textractJobId = $1 + WHERE id = $2 ` type AddDocumentTextTriggerJobIdParams struct { @@ -79,7 +113,9 @@ type AddDocumentTextTriggerJobIdParams struct { // AddDocumentTextTriggerJobId // -// UPDATE documentTextTriggerExtractions SET textractJobId = $1 WHERE id = $2 +// UPDATE documentTextTriggerExtractions +// SET textractJobId = $1 +// WHERE id = $2 func (q *Queries) AddDocumentTextTriggerJobId(ctx context.Context, arg *AddDocumentTextTriggerJobIdParams) error { _, err := q.db.Exec(ctx, addDocumentTextTriggerJobId, arg.Textractjobid, arg.ID) return err @@ -169,6 +205,142 @@ func (q *Queries) GetTextEntryByDocId(ctx context.Context, documentid uuid.UUID) return &i, err } +const getTextOutCurrentPart = `-- name: GetTextOutCurrentPart :one +WITH client as ( + select id from clients where id = $1 +), +parts as ( + SELECT extractions.part + FROM client as c + JOIN documents as docs on docs.clientId = c.id + JOIN documentCleans as clean on docs.id = clean.documentId + JOIN documentTextTriggerExtractions as triggers + on triggers.cleanId = clean.id + JOIN documentTextExtractionEntries as textEntries + on textEntries.triggerId = triggers.id + JOIN documentTextExtractions as extractions + on extractions.id = textEntries.textId + WHERE date(extractions.createdAt) = date($2) +), +max_part as ( + SELECT COALESCE(max(part), 0) as max_part_num FROM parts +) +SELECT + COALESCE(p.part, 0)::unsignedsmallint as part, + COUNT(p.part) as count +FROM parts p +RIGHT JOIN max_part mp ON p.part = mp.max_part_num +GROUP BY p.part +` + +type GetTextOutCurrentPartParams struct { + Clientid *string `db:"clientid"` + Querydate pgtype.Timestamptz `db:"querydate"` +} + +type GetTextOutCurrentPartRow struct { + Part uint16 `db:"part"` + Count int64 `db:"count"` +} + +// GetTextOutCurrentPart +// +// WITH client as ( +// select id from clients where id = $1 +// ), +// parts as ( +// SELECT extractions.part +// FROM client as c +// JOIN documents as docs on docs.clientId = c.id +// JOIN documentCleans as clean on docs.id = clean.documentId +// JOIN documentTextTriggerExtractions as triggers +// on triggers.cleanId = clean.id +// JOIN documentTextExtractionEntries as textEntries +// on textEntries.triggerId = triggers.id +// JOIN documentTextExtractions as extractions +// on extractions.id = textEntries.textId +// WHERE date(extractions.createdAt) = date($2) +// ), +// max_part as ( +// SELECT COALESCE(max(part), 0) as max_part_num FROM parts +// ) +// SELECT +// COALESCE(p.part, 0)::unsignedsmallint as part, +// COUNT(p.part) as count +// FROM parts p +// RIGHT JOIN max_part mp ON p.part = mp.max_part_num +// GROUP BY p.part +func (q *Queries) GetTextOutCurrentPart(ctx context.Context, arg *GetTextOutCurrentPartParams) (*GetTextOutCurrentPartRow, error) { + row := q.db.QueryRow(ctx, getTextOutCurrentPart, arg.Clientid, arg.Querydate) + var i GetTextOutCurrentPartRow + err := row.Scan(&i.Part, &i.Count) + return &i, err +} + +const getTextractOutputCurrentPart = `-- name: GetTextractOutputCurrentPart :one +WITH client as ( + select id from clients where id = $1 +), +parts as ( + SELECT triggers.part + FROM client as c + JOIN documents as docs on docs.clientId = c.id + JOIN documentCleans as clean on docs.id = clean.documentId + JOIN documentTextTriggerExtractions as triggers + on triggers.cleanId = clean.id + WHERE date(triggers.createdAt) = date($2) +), +max_part as ( + SELECT COALESCE(max(part), 0) as max_part_num FROM parts +) +SELECT + COALESCE(p.part, 0)::unsignedsmallint as part, + COUNT(p.part) as count +FROM parts p +RIGHT JOIN max_part mp ON p.part = mp.max_part_num +GROUP BY p.part +` + +type GetTextractOutputCurrentPartParams struct { + Clientid *string `db:"clientid"` + Querydate pgtype.Timestamptz `db:"querydate"` +} + +type GetTextractOutputCurrentPartRow struct { + Part uint16 `db:"part"` + Count int64 `db:"count"` +} + +// GetTextractOutputCurrentPart +// +// WITH client as ( +// select id from clients where id = $1 +// ), +// parts as ( +// SELECT triggers.part +// FROM client as c +// JOIN documents as docs on docs.clientId = c.id +// JOIN documentCleans as clean on docs.id = clean.documentId +// JOIN documentTextTriggerExtractions as triggers +// on triggers.cleanId = clean.id +// WHERE date(triggers.createdAt) = date($2) +// ), +// max_part as ( +// SELECT COALESCE(max(part), 0) as max_part_num FROM parts +// ) +// SELECT +// COALESCE(p.part, 0)::unsignedsmallint as part, +// COUNT(p.part) as count +// FROM parts p +// RIGHT JOIN max_part mp ON p.part = mp.max_part_num +// GROUP BY p.part +func (q *Queries) GetTextractOutputCurrentPart(ctx context.Context, arg *GetTextractOutputCurrentPartParams) (*GetTextractOutputCurrentPartRow, error) { + row := q.db.QueryRow(ctx, getTextractOutputCurrentPart, arg.Clientid, arg.Querydate) + var i GetTextractOutputCurrentPartRow + err := row.Scan(&i.Part, &i.Count) + return &i, err +} + const isDocumentTextExtracted = `-- name: IsDocumentTextExtracted :one SELECT EXISTS( SELECT 1 FROM currentTextEntries WHERE documentId = $1 diff --git a/internal/database/repository/text_test.go b/internal/database/repository/text_test.go index c4a59d4f..89ad3c3a 100644 --- a/internal/database/repository/text_test.go +++ b/internal/database/repository/text_test.go @@ -3,12 +3,14 @@ package repository_test import ( "context" "testing" + "time" "queryorchestration/internal/database/repository" "queryorchestration/internal/serviceconfig" "queryorchestration/internal/test" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -69,6 +71,11 @@ func TestTextExtraction(t *testing.T) { triggerId, err := queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ Cleanid: cleanid, Version: 123, + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) @@ -110,6 +117,11 @@ func TestTextExtraction(t *testing.T) { Bucket: bucket, Key: key, Hash: "example", + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, }) require.NoError(t, err) assert.NotNil(t, textId) @@ -157,3 +169,336 @@ func TestTextExtraction(t *testing.T) { }) require.EqualError(t, err, "no rows in result set") } + +func TestTextTextractPart(t *testing.T) { + if testing.Short() { + t.Skip("Skipping long test in short mode") + } + ctx := context.Background() + + cfg := &serviceconfig.BaseConfig{} + test.SetCfgProvider(t, cfg) + _, textup := test.CreateDB(t, ctx, cfg, &test.CreateDatabaseConfig{ + RunMigrations: true, + }) + defer textup() + + queries := cfg.GetDBQueries() + + clientId := "EXAMPLE" + err := queries.CreateClient(ctx, &repository.CreateClientParams{ + Name: "example_client", + ID: clientId, + }) + require.NoError(t, err) + + hash := "example_hash" + id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{ + Clientid: clientId, + Hash: hash, + }) + require.NoError(t, err) + assert.NotEmpty(t, id) + + bucket := "example_bucket" + key := "example_key" + cleanid, err := queries.AddDocumentClean(ctx, &repository.AddDocumentCleanParams{ + Documentid: id, + Bucket: &bucket, + Key: &key, + Hash: &hash, + Mimetype: repository.NullCleanmimetype{ + Valid: true, + Cleanmimetype: repository.CleanmimetypeApplicationPdf, + }, + }) + require.NoError(t, err) + err = queries.AddDocumentCleanEntry(ctx, &repository.AddDocumentCleanEntryParams{ + Cleanid: cleanid, + Version: 1, + }) + require.NoError(t, err) + + part, err := queries.GetTextractOutputCurrentPart(ctx, &repository.GetTextractOutputCurrentPartParams{ + Clientid: &clientId, + Querydate: pgtype.Timestamptz{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + assert.Equal(t, &repository.GetTextractOutputCurrentPartRow{ + Count: 0, + Part: 0, + }, part) + + _, err = queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ + Cleanid: cleanid, + Version: 123, + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + + part, err = queries.GetTextractOutputCurrentPart(ctx, &repository.GetTextractOutputCurrentPartParams{ + Clientid: &clientId, + Querydate: pgtype.Timestamptz{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + assert.Equal(t, &repository.GetTextractOutputCurrentPartRow{ + Count: 1, + Part: 0, + }, part) + + _, err = queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ + Cleanid: cleanid, + Version: 123, + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + + part, err = queries.GetTextractOutputCurrentPart(ctx, &repository.GetTextractOutputCurrentPartParams{ + Clientid: &clientId, + Querydate: pgtype.Timestamptz{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + assert.Equal(t, &repository.GetTextractOutputCurrentPartRow{ + Count: 2, + Part: 0, + }, part) + + _, err = queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ + Cleanid: cleanid, + Version: 123, + Part: 1, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + + part, err = queries.GetTextractOutputCurrentPart(ctx, &repository.GetTextractOutputCurrentPartParams{ + Clientid: &clientId, + Querydate: pgtype.Timestamptz{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + assert.Equal(t, &repository.GetTextractOutputCurrentPartRow{ + Count: 1, + Part: 1, + }, part) +} + +func TestTextOutPart(t *testing.T) { + if testing.Short() { + t.Skip("Skipping long test in short mode") + } + ctx := context.Background() + + cfg := &serviceconfig.BaseConfig{} + test.SetCfgProvider(t, cfg) + _, textup := test.CreateDB(t, ctx, cfg, &test.CreateDatabaseConfig{ + RunMigrations: true, + }) + defer textup() + + queries := cfg.GetDBQueries() + + clientId := "EXAMPLE" + err := queries.CreateClient(ctx, &repository.CreateClientParams{ + Name: "example_client", + ID: clientId, + }) + require.NoError(t, err) + + hash := "example_hash" + id, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{ + Clientid: clientId, + Hash: hash, + }) + require.NoError(t, err) + assert.NotEmpty(t, id) + + bucket := "example_bucket" + key := "example_key" + cleanid, err := queries.AddDocumentClean(ctx, &repository.AddDocumentCleanParams{ + Documentid: id, + Bucket: &bucket, + Key: &key, + Hash: &hash, + Mimetype: repository.NullCleanmimetype{ + Valid: true, + Cleanmimetype: repository.CleanmimetypeApplicationPdf, + }, + }) + require.NoError(t, err) + err = queries.AddDocumentCleanEntry(ctx, &repository.AddDocumentCleanEntryParams{ + Cleanid: cleanid, + Version: 1, + }) + require.NoError(t, err) + + triggerId, err := queries.AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ + Cleanid: cleanid, + Version: 123, + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + + part, err := queries.GetTextOutCurrentPart(ctx, &repository.GetTextOutCurrentPartParams{ + Clientid: &clientId, + Querydate: pgtype.Timestamptz{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + assert.Equal(t, &repository.GetTextOutCurrentPartRow{ + Count: 0, + Part: 0, + }, part) + + textId, err := queries.AddDocumentText(ctx, &repository.AddDocumentTextParams{ + Bucket: "hi", + Key: "hi", + Hash: "hi", + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + err = queries.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{ + Textid: textId, + Triggerid: triggerId, + Version: 123, + }) + require.NoError(t, err) + + part, err = queries.GetTextOutCurrentPart(ctx, &repository.GetTextOutCurrentPartParams{ + Clientid: &clientId, + Querydate: pgtype.Timestamptz{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + assert.Equal(t, &repository.GetTextOutCurrentPartRow{ + Count: 1, + Part: 0, + }, part) + + textId, err = queries.AddDocumentText(ctx, &repository.AddDocumentTextParams{ + Bucket: "hi", + Key: "hi", + Hash: "hi", + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + err = queries.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{ + Textid: textId, + Triggerid: triggerId, + Version: 123, + }) + require.NoError(t, err) + + part, err = queries.GetTextOutCurrentPart(ctx, &repository.GetTextOutCurrentPartParams{ + Clientid: &clientId, + Querydate: pgtype.Timestamptz{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + assert.Equal(t, &repository.GetTextOutCurrentPartRow{ + Count: 2, + Part: 0, + }, part) + + textId, err = queries.AddDocumentText(ctx, &repository.AddDocumentTextParams{ + Bucket: "hi", + Key: "hi", + Hash: "hi", + Part: 0, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + err = queries.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{ + Textid: textId, + Triggerid: triggerId, + Version: 123, + }) + require.NoError(t, err) + + part, err = queries.GetTextOutCurrentPart(ctx, &repository.GetTextOutCurrentPartParams{ + Clientid: &clientId, + Querydate: pgtype.Timestamptz{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + assert.Equal(t, &repository.GetTextOutCurrentPartRow{ + Count: 3, + Part: 0, + }, part) + + textId, err = queries.AddDocumentText(ctx, &repository.AddDocumentTextParams{ + Bucket: "hi", + Key: "hi", + Hash: "hi", + Part: 1, + Createdat: pgtype.Timestamp{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + err = queries.AddDocumentTextEntry(ctx, &repository.AddDocumentTextEntryParams{ + Textid: textId, + Triggerid: triggerId, + Version: 123, + }) + require.NoError(t, err) + + part, err = queries.GetTextOutCurrentPart(ctx, &repository.GetTextOutCurrentPartParams{ + Clientid: &clientId, + Querydate: pgtype.Timestamptz{ + Time: time.Now().UTC(), + Valid: true, + }, + }) + require.NoError(t, err) + assert.Equal(t, &repository.GetTextOutCurrentPartRow{ + Count: 1, + Part: 1, + }, part) +} diff --git a/internal/document/text/extract_test.go b/internal/document/text/extract_test.go index e8f3d794..a49c78c8 100644 --- a/internal/document/text/extract_test.go +++ b/internal/document/text/extract_test.go @@ -78,7 +78,12 @@ func TestCreate(t *testing.T) { jobId := "hello" triggerId := uuid.New() pool.ExpectBegin() - pool.ExpectQuery("name: AddDocumentTextTrigger :one").WithArgs(cleanId, pgxmock.AnyArg()).WillReturnRows( + pool.ExpectQuery("name: GetTextractOutputCurrentPart :one").WithArgs(&clientId, pgxmock.AnyArg()). + WillReturnRows( + pgxmock.NewRows([]string{"part", "count"}). + AddRow(0, 2), + ) + pool.ExpectQuery("name: AddDocumentTextTrigger :one").WithArgs(cleanId, pgxmock.AnyArg(), pgxmock.AnyArg(), uint16(0)).WillReturnRows( pgxmock.NewRows([]string{"id"}).AddRow(triggerId), ) diff --git a/internal/document/text/processExtract.go b/internal/document/text/processExtract.go index c92ffa54..62a8ec01 100644 --- a/internal/document/text/processExtract.go +++ b/internal/document/text/processExtract.go @@ -15,6 +15,7 @@ import ( "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" ) type ProcessParams struct { @@ -103,12 +104,32 @@ func (s *Service) storeProcess(ctx context.Context, text io.Reader, trigger *rep CreatedAt: createdAt, EntityID: trigger.ID, } + + currentPart, err := q.GetTextOutCurrentPart(ctx, &repository.GetTextOutCurrentPartParams{ + Clientid: &key.ClientID, + Querydate: pgtype.Timestamptz{ + Time: key.CreatedAt, + Valid: true, + }, + }) + if err != nil { + return err + } + + part := s.cfg.GetDirectoryPart(currentPart.Part, currentPart.Count) + key.Part = &part + keyStr := key.String() textId, err = q.AddDocumentText(ctx, &repository.AddDocumentTextParams{ Bucket: params.Bucket, Key: keyStr, Hash: hash, + Part: *key.Part, + Createdat: pgtype.Timestamp{ + Time: key.CreatedAt, + Valid: true, + }, }) if err != nil { return err diff --git a/internal/document/text/processExtract_test.go b/internal/document/text/processExtract_test.go index 3accfee9..1b7eab1a 100644 --- a/internal/document/text/processExtract_test.go +++ b/internal/document/text/processExtract_test.go @@ -143,17 +143,24 @@ func TestStoreTrigger(t *testing.T) { require.NoError(t, err) }) t.Run("not exists", func(t *testing.T) { + part := uint16(0) key := objectstore.BucketKey{ ClientID: clientId, Location: objectstore.TextOut, CreatedAt: time.Now().UTC(), EntityID: triggerId, + Part: &part, } hash := `"09644372e99020106946045c6fd2d70b"` pool.ExpectBegin() pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash). WillReturnError(sql.ErrNoRows) - pool.ExpectQuery("name: AddDocumentText :one").WithArgs(bucket, key.String(), hash). + pool.ExpectQuery("name: GetTextOutCurrentPart :one").WithArgs(&clientId, pgxmock.AnyArg()). + WillReturnRows( + pgxmock.NewRows([]string{"part", "count"}). + AddRow(0, 2), + ) + pool.ExpectQuery("name: AddDocumentText :one").WithArgs(bucket, key.String(), hash, pgxmock.AnyArg(), uint16(0)). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(textId), @@ -233,7 +240,12 @@ func TestProcessTrigger(t *testing.T) { pool.ExpectBegin() pool.ExpectQuery("name: GetDocumentTextExtractionByHash :one").WithArgs(cleanId, hash). WillReturnError(sql.ErrNoRows) - pool.ExpectQuery("name: AddDocumentText :one").WithArgs(bucket, key.String(), hash). + pool.ExpectQuery("name: GetTextOutCurrentPart :one").WithArgs(&clientId, pgxmock.AnyArg()). + WillReturnRows( + pgxmock.NewRows([]string{"part", "count"}). + AddRow(0, 2), + ) + pool.ExpectQuery("name: AddDocumentText :one").WithArgs(bucket, key.String(), hash, pgxmock.AnyArg(), uint16(0)). WillReturnRows( pgxmock.NewRows([]string{"id"}). AddRow(textId), diff --git a/internal/document/text/triggerExtract.go b/internal/document/text/triggerExtract.go index 130fd746..e4a50393 100644 --- a/internal/document/text/triggerExtract.go +++ b/internal/document/text/triggerExtract.go @@ -11,28 +11,48 @@ import ( "github.com/aws/aws-sdk-go-v2/service/textract" "github.com/aws/aws-sdk-go-v2/service/textract/types" + "github.com/jackc/pgx/v5/pgtype" ) func (s *Service) triggerExtract(ctx context.Context, clean *repository.Currentcleanentry) error { version := build.GetVersionUnixTimestamp() return s.cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error { - triggerId, err := s.cfg.GetDBQueries().AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ - Cleanid: clean.ID, - Version: version, + key := objectstore.BucketKey{ + CreatedAt: time.Now().UTC(), + Location: objectstore.TextTextract, + ClientID: clean.Clientid, + } + + currentPart, err := q.GetTextractOutputCurrentPart(ctx, &repository.GetTextractOutputCurrentPartParams{ + Clientid: &clean.Clientid, + Querydate: pgtype.Timestamptz{ + Time: key.CreatedAt, + Valid: true, + }, }) if err != nil { return err } - key := objectstore.BucketKey{ - CreatedAt: time.Now().UTC(), - Location: objectstore.TextTextract, - ClientID: clean.Clientid, - EntityID: triggerId, - } - keyStr := key.String() + part := s.cfg.GetDirectoryPart(currentPart.Part, currentPart.Count) + key.Part = &part + triggerId, err := s.cfg.GetDBQueries().AddDocumentTextTrigger(ctx, &repository.AddDocumentTextTriggerParams{ + Cleanid: clean.ID, + Version: version, + Createdat: pgtype.Timestamp{ + Time: key.CreatedAt, + Valid: true, + }, + Part: *key.Part, + }) + if err != nil { + return err + } + + key.EntityID = triggerId + keyStr := key.String() jobTag := triggerId.String() res, err := s.cfg.GetTextractClient().StartDocumentTextDetection(ctx, &textract.StartDocumentTextDetectionInput{ diff --git a/internal/document/text/triggerExtract_test.go b/internal/document/text/triggerExtract_test.go index 7a0ad4fc..97fddb48 100644 --- a/internal/document/text/triggerExtract_test.go +++ b/internal/document/text/triggerExtract_test.go @@ -38,11 +38,18 @@ func TestTriggerExtract(t *testing.T) { bucket := "bucket_name" key := "/i/am/here" hash := "hhasshhh" + clientId := "AA" jobId := "hello" triggerId := uuid.New() pool.ExpectBegin() - pool.ExpectQuery("name: AddDocumentTextTrigger :one").WithArgs(cleanId, pgxmock.AnyArg()).WillReturnRows( + pool.ExpectQuery("name: GetTextractOutputCurrentPart :one").WithArgs(&clientId, pgxmock.AnyArg()). + WillReturnRows( + pgxmock.NewRows([]string{"part", "count"}). + AddRow(0, 2), + ) + + pool.ExpectQuery("name: AddDocumentTextTrigger :one").WithArgs(cleanId, pgxmock.AnyArg(), pgxmock.AnyArg(), uint16(0)).WillReturnRows( pgxmock.NewRows([]string{"id"}).AddRow(triggerId), ) @@ -64,7 +71,7 @@ func TestTriggerExtract(t *testing.T) { err = svc.triggerExtract(ctx, &repository.Currentcleanentry{ ID: cleanId, - Clientid: "AA", + Clientid: clientId, Key: &key, Bucket: &bucket, Hash: &hash, diff --git a/internal/serviceconfig/objectstore/bucketkey.go b/internal/serviceconfig/objectstore/bucketkey.go index ec92b62d..f56564eb 100644 --- a/internal/serviceconfig/objectstore/bucketkey.go +++ b/internal/serviceconfig/objectstore/bucketkey.go @@ -18,7 +18,7 @@ type BucketKey struct { ClientID string Location BucketLocation CreatedAt time.Time - Part *uint8 + Part *uint16 EntityID uuid.UUID FileType *string } @@ -30,22 +30,42 @@ const ( Export BucketLocation = "export" ) -var timeFormat = "2006-01-02T150405Z" -var dateFormat = "20060102" +var PartExclusions = []BucketLocation{ + Export, +} + +const TimeFormat = "2006-01-02T150405Z" +const DateFormat = "20060102" func (c *BucketKey) String() string { - dateStr := c.CreatedAt.Format(dateFormat) + return fmt.Sprintf("%s/%s", c.Prefix(), c.Filename()) +} + +func (c *BucketKey) Prefix() string { + dateStr := c.CreatedAt.Format(DateFormat) + partStr := "" - if c.Part != nil { + isExcluded := false + for _, location := range PartExclusions { + if location == c.Location { + isExcluded = true + break + } + } + if !isExcluded { + if c.Part == nil { + part := uint16(0) + c.Part = &part + } partStr = fmt.Sprintf("/%d", *c.Part) } - return fmt.Sprintf("%s/%s/%s%s/%s", c.ClientID, c.Location, dateStr, partStr, c.Filename()) + return fmt.Sprintf("%s/%s/%s%s", c.ClientID, c.Location, dateStr, partStr) } func (c *BucketKey) Filename() string { location := strings.ReplaceAll(string(c.Location), "/", "-") - name := fmt.Sprintf("%s_%s_%s_%s", c.CreatedAt.Format(timeFormat), c.ClientID, location, c.EntityID) + name := fmt.Sprintf("%s_%s_%s_%s", c.CreatedAt.Format(TimeFormat), c.ClientID, location, c.EntityID) if c.FileType != nil { name = fmt.Sprintf("%s.%s", name, *c.FileType) } @@ -92,7 +112,7 @@ func (c *BucketKey) parseFileType(name string) { c.FileType = &name } func (c *BucketKey) parseTime(name string) error { - parsedTime, err := time.Parse(timeFormat, name) + parsedTime, err := time.Parse(TimeFormat, name) if err != nil { return fmt.Errorf("failed to parse date: %w", err) } @@ -151,7 +171,7 @@ func ParseBucketKey(key string) (BucketKey, error) { if err != nil { return BucketKey{}, err } - safePart := uint8(part) + safePart := uint16(part) formattedKey.Part = &safePart err = formattedKey.parseFilename(subMatches[2]) diff --git a/internal/serviceconfig/objectstore/bucketkey_test.go b/internal/serviceconfig/objectstore/bucketkey_test.go index b79e1a59..e5b3bd5c 100644 --- a/internal/serviceconfig/objectstore/bucketkey_test.go +++ b/internal/serviceconfig/objectstore/bucketkey_test.go @@ -10,27 +10,34 @@ import ( func TestGetBucketKey(t *testing.T) { key := BucketKey{} - assert.Equal(t, "//00010101/0001-01-01T000000Z___00000000-0000-0000-0000-000000000000", key.String()) + assert.Equal(t, "//00010101/0/0001-01-01T000000Z___00000000-0000-0000-0000-000000000000", key.String()) key.CreatedAt = time.Date(2025, time.March, 31, 4, 8, 16, 222, time.UTC) - assert.Equal(t, "//20250331/2025-03-31T040816Z___00000000-0000-0000-0000-000000000000", key.String()) + assert.Equal(t, "//20250331/0/2025-03-31T040816Z___00000000-0000-0000-0000-000000000000", key.String()) key.ClientID = "AAA" - assert.Equal(t, "AAA//20250331/2025-03-31T040816Z_AAA__00000000-0000-0000-0000-000000000000", key.String()) + assert.Equal(t, "AAA//20250331/0/2025-03-31T040816Z_AAA__00000000-0000-0000-0000-000000000000", key.String()) key.EntityID = uuid.MustParse("b745910f-f529-43c5-87e1-25629d46ef40") - assert.Equal(t, "AAA//20250331/2025-03-31T040816Z_AAA__b745910f-f529-43c5-87e1-25629d46ef40", key.String()) + assert.Equal(t, "AAA//20250331/0/2025-03-31T040816Z_AAA__b745910f-f529-43c5-87e1-25629d46ef40", key.String()) key.Location = Import - assert.Equal(t, "AAA/import/20250331/2025-03-31T040816Z_AAA_import_b745910f-f529-43c5-87e1-25629d46ef40", key.String()) - part := uint8(3) + assert.Equal(t, "AAA/import/20250331/0/2025-03-31T040816Z_AAA_import_b745910f-f529-43c5-87e1-25629d46ef40", key.String()) + key.Location = TextTextract + assert.Equal(t, "AAA/text/textract/20250331/0/2025-03-31T040816Z_AAA_text-textract_b745910f-f529-43c5-87e1-25629d46ef40", key.String()) + key.Location = TextOut + assert.Equal(t, "AAA/text/out/20250331/0/2025-03-31T040816Z_AAA_text-out_b745910f-f529-43c5-87e1-25629d46ef40", key.String()) + key.Location = Export + assert.Equal(t, "AAA/export/20250331/2025-03-31T040816Z_AAA_export_b745910f-f529-43c5-87e1-25629d46ef40", key.String()) + part := uint16(3) key.Part = &part + key.Location = Import assert.Equal(t, "AAA/import/20250331/3/2025-03-31T040816Z_AAA_import_b745910f-f529-43c5-87e1-25629d46ef40", key.String()) key.Location = TextTextract assert.Equal(t, "AAA/text/textract/20250331/3/2025-03-31T040816Z_AAA_text-textract_b745910f-f529-43c5-87e1-25629d46ef40", key.String()) key.Location = TextOut assert.Equal(t, "AAA/text/out/20250331/3/2025-03-31T040816Z_AAA_text-out_b745910f-f529-43c5-87e1-25629d46ef40", key.String()) key.Location = Export - assert.Equal(t, "AAA/export/20250331/3/2025-03-31T040816Z_AAA_export_b745910f-f529-43c5-87e1-25629d46ef40", key.String()) + assert.Equal(t, "AAA/export/20250331/2025-03-31T040816Z_AAA_export_b745910f-f529-43c5-87e1-25629d46ef40", key.String()) filetype := "csv" key.FileType = &filetype - assert.Equal(t, "AAA/export/20250331/3/2025-03-31T040816Z_AAA_export_b745910f-f529-43c5-87e1-25629d46ef40.csv", key.String()) + assert.Equal(t, "AAA/export/20250331/2025-03-31T040816Z_AAA_export_b745910f-f529-43c5-87e1-25629d46ef40.csv", key.String()) } func TestParseBucketKey(t *testing.T) { _, err := ParseBucketKey("") @@ -87,7 +94,7 @@ func TestParseBucketKey(t *testing.T) { CreatedAt: time.Date(2025, 3, 31, 0, 0, 0, 0, time.UTC), }, key) - part := uint8(3) + part := uint16(3) key, err = ParseBucketKey("AAA/import/20250331/3/2025-03-31T040816Z_AAA_import_b745910f-f529-43c5-87e1-25629d46ef40") assert.NoError(t, err) assert.EqualExportedValues(t, BucketKey{ diff --git a/internal/serviceconfig/objectstore/config.go b/internal/serviceconfig/objectstore/config.go index 6bb2c633..5114d038 100644 --- a/internal/serviceconfig/objectstore/config.go +++ b/internal/serviceconfig/objectstore/config.go @@ -19,6 +19,7 @@ type ObjectStoreConfig struct { AWSS3UsePathStyle bool `env:"AWS_S3_USE_PATH_STYLE" envDefault:"false"` StoreClient S3Client ChunkSize int + MaxDirSize int64 } type ConfigProvider interface { @@ -31,6 +32,7 @@ type ConfigProvider interface { GetS3Endpoint() string SetS3UsePathStyle(bool) CalculateETag(context.Context, io.Reader) (string, error) + GetDirectoryPart(uint16, int64) uint16 } func (c *ObjectStoreConfig) SetStoreClientAndPingByName(ctx context.Context, name string) error { @@ -98,7 +100,9 @@ func (c *ObjectStoreConfig) SetS3UsePathStyle(e bool) { } func (c *ObjectStoreConfig) CalculateETag(ctx context.Context, doc io.Reader) (string, error) { - c.ChunkSize = 5 * 1024 * 1024 + if c.ChunkSize <= 0 { + c.ChunkSize = 5 * 1024 * 1024 + } buffer := make([]byte, c.ChunkSize) var partHashes [][]byte var totalBytes int @@ -140,3 +144,15 @@ func (c *ObjectStoreConfig) CalculateETag(ctx context.Context, doc io.Reader) (s return fmt.Sprintf(`"%s-%d"`, hex.EncodeToString(finalHash.Sum(nil)), len(partHashes)), nil } + +func (c *ObjectStoreConfig) GetDirectoryPart(part uint16, count int64) uint16 { + if c.MaxDirSize <= 0 { + c.MaxDirSize = 1000 + } + + if count >= c.MaxDirSize { + return part + 1 + } + + return part +} diff --git a/internal/serviceconfig/objectstore/config_test.go b/internal/serviceconfig/objectstore/config_test.go index 5385da39..5ee89e86 100644 --- a/internal/serviceconfig/objectstore/config_test.go +++ b/internal/serviceconfig/objectstore/config_test.go @@ -115,3 +115,26 @@ func TestCalculateETag(t *testing.T) { assert.Equal(t, *res.ETag, hash) } + +func TestGetDirectoryPart(t *testing.T) { + cfg := &StoreConfig{} + + assert.Equal(t, uint16(0), cfg.GetDirectoryPart(0, 0)) + assert.Equal(t, uint16(0), cfg.GetDirectoryPart(0, 1)) + assert.Equal(t, uint16(0), cfg.GetDirectoryPart(0, 2)) + assert.Equal(t, uint16(1), cfg.GetDirectoryPart(0, 1000)) + + cfg.MaxDirSize = 2 + assert.Equal(t, uint16(0), cfg.GetDirectoryPart(0, 0)) + assert.Equal(t, uint16(0), cfg.GetDirectoryPart(0, 1)) + assert.Equal(t, uint16(1), cfg.GetDirectoryPart(0, 2)) + assert.Equal(t, uint16(1), cfg.GetDirectoryPart(1, 0)) + assert.Equal(t, uint16(1), cfg.GetDirectoryPart(1, 1)) + assert.Equal(t, uint16(2), cfg.GetDirectoryPart(1, 2)) + assert.Equal(t, uint16(2), cfg.GetDirectoryPart(2, 0)) + assert.Equal(t, uint16(2), cfg.GetDirectoryPart(2, 1)) + assert.Equal(t, uint16(3), cfg.GetDirectoryPart(2, 2)) + assert.Equal(t, uint16(3), cfg.GetDirectoryPart(3, 0)) + assert.Equal(t, uint16(3), cfg.GetDirectoryPart(3, 1)) + assert.Equal(t, uint16(4), cfg.GetDirectoryPart(3, 2)) +} diff --git a/scripts/openapi-scripts.yml b/scripts/openapi-scripts.yml index 86576423..40c6b5a8 100644 --- a/scripts/openapi-scripts.yml +++ b/scripts/openapi-scripts.yml @@ -32,8 +32,12 @@ tasks: service=$(basename "$file") service=${service%.*} service_lower=${service,,} + + # Client Code generateGo $service_lower ./pkg/$service $file \ "models,client" + + # Server Code generateGo $service_lower ./api/$service $file \ "models,embedded-spec,echo-server" fi diff --git a/sqlc.yml b/sqlc.yml index 9b3910a5..4fe1e05b 100644 --- a/sqlc.yml +++ b/sqlc.yml @@ -39,6 +39,8 @@ sql: import: "github.com/google/uuid" type: "UUID" pointer: true + - db_type: "unsignedsmallint" + go_type: "uint16" rules: - name: no-pg message: "invalid engine: not postgresql" diff --git a/test/process_test.go b/test/process_test.go index a280a44b..c88b055c 100644 --- a/test/process_test.go +++ b/test/process_test.go @@ -65,11 +65,13 @@ func TestProcess(t *testing.T) { queryapitest.WaitForClientStatus(t, ctx, net.Client, client.Id, queryapi.INSYNC) + part := uint16(0) importKey := objectstore.BucketKey{ ClientID: client.Id, EntityID: uuid.New(), Location: objectstore.Import, CreatedAt: time.Now().UTC(), + Part: &part, } expectation := test.CreateStartDocTextDetectionExpectation(t, net.Dependencies.MockServer, test.StartDocTextDetectionExpectationParams{ Bucket: net.Dependencies.BucketName,