Merged in feature/track-filesize (pull request #200)

track file sizes for all documents in system

* feature complete

needs dev testing

internal/document/init/create.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"errors"
+	"fmt"
 	"log/slog"
 
 	docsyncrunner "queryorchestration/api/docSyncRunner"
@@ -11,6 +12,8 @@ import (
 	"queryorchestration/internal/serviceconfig/objectstore"
 	"queryorchestration/internal/serviceconfig/queue"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/google/uuid"
 )
 
@@ -85,16 +88,24 @@ func (s *Service) getCreateParams(ctx context.Context, doc *Create) (*createDocu
 func (s *Service) submitCreate(ctx context.Context, params *createDocumentParams) (uuid.UUID, error) {
 	var id uuid.UUID
 
-	err := s.cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error {
+	// Measure file size from S3 (trusted source - we measure it ourselves)
+	// The S3 client must be initialized - file size tracking is required
+	fileSizeBytes, err := s.measureFileSize(ctx, params.Bucket, params.Key.String())
+	if err != nil {
+		return uuid.Nil, fmt.Errorf("failed to measure file size: %w", err)
+	}
+
+	err = s.cfg.ExecuteDBTransaction(ctx, func(ctx context.Context, q *repository.Queries) error {
 		var dbid uuid.UUID
 		if params.ID == nil {
 			createid, err := s.cfg.GetDBQueries().CreateDocument(ctx, &repository.CreateDocumentParams{
-				Clientid:     params.Key.ClientID,
-				Hash:         params.Hash,
-				BatchID:      params.Key.BatchID,
-				Filename:     params.Filename,
-				Folderid:     params.FolderID,
-				Originalpath: params.OriginalPath,
+				Clientid:      params.Key.ClientID,
+				Hash:          params.Hash,
+				BatchID:       params.Key.BatchID,
+				Filename:      params.Filename,
+				Folderid:      params.FolderID,
+				Originalpath:  params.OriginalPath,
+				FileSizeBytes: fileSizeBytes,
 			})
 			if err != nil {
 				return err
@@ -127,3 +138,38 @@ func (s *Service) submitCreate(ctx context.Context, params *createDocumentParams
 
 	return id, nil
 }
+
+// measureFileSize retrieves the file size from S3 using HeadObject.
+// Returns the file size in bytes, or an error if the S3 client is not configured
+// or if the HeadObject call fails.
+//
+// Parameters:
+//   - ctx: Context for the S3 operation
+//   - bucket: S3 bucket name
+//   - key: S3 object key
+//
+// Returns:
+//   - *int64: Pointer to file size in bytes
+//   - error: Error if S3 client is nil or HeadObject fails
+//
+// Note: This operation requires s3:GetObject permission on the bucket/key.
+func (s *Service) measureFileSize(ctx context.Context, bucket, key string) (*int64, error) {
+	storeClient := s.cfg.GetStoreClient()
+	if storeClient == nil {
+		return nil, errors.New("S3 client not initialized - required for file size measurement")
+	}
+
+	headOutput, err := storeClient.HeadObject(ctx, &s3.HeadObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(key),
+	})
+	if err != nil {
+		return nil, fmt.Errorf("HeadObject failed for s3://%s/%s: %w", bucket, key, err)
+	}
+
+	if headOutput.ContentLength == nil {
+		return nil, fmt.Errorf("HeadObject returned nil ContentLength for s3://%s/%s", bucket, key)
+	}
+
+	return headOutput.ContentLength, nil
+}

internal/document/init/create_test.go

@@ -1,6 +1,7 @@
 package documentinit
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"testing"
@@ -9,10 +10,14 @@ import (
 	"queryorchestration/internal/database/repository"
 	"queryorchestration/internal/document"
 	"queryorchestration/internal/serviceconfig/objectstore"
+	"queryorchestration/internal/test"
+	objectstoremock "queryorchestration/mocks/objectstore"
 	queuemock "queryorchestration/mocks/queue"
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"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"
@@ -30,6 +35,11 @@ func TestCreate(t *testing.T) {
 	cfg.DBQueries = repository.New(pool)
 	mockSQS := queuemock.NewMockSQSClient(t)
 	cfg.QueueClient = mockSQS
+
+	// Set up mock S3 client for file size measurement
+	mockS3 := objectstoremock.NewMockS3Client(t)
+	cfg.StoreClient = mockS3
+
 	cfg.DocumentSyncURL = "/i/am/here"
 
 	svc := New(cfg)
@@ -48,11 +58,19 @@ func TestCreate(t *testing.T) {
 		Location:  objectstore.Import,
 	}
 
+	// Mock HeadObject to return file size
+	fileSize := int64(1024)
+	mockS3.EXPECT().
+		HeadObject(mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
+			return *in.Bucket == bucket && *in.Key == key.String()
+		})).
+		Return(&s3.HeadObjectOutput{ContentLength: &fileSize}, nil)
+
 	pool.ExpectQuery("name: GetDocumentIDByHash :one").WithArgs(doc.Hash, doc.ClientID).WillReturnRows(
 		pgxmock.NewRows([]string{"id"}),
 	)
 	pool.ExpectBegin()
-	pool.ExpectQuery("name: CreateDocument :one").WithArgs(doc.ClientID, doc.Hash, (*uuid.UUID)(nil), (*string)(nil), (*uuid.UUID)(nil), (*string)(nil)).
+	pool.ExpectQuery("name: CreateDocument :one").WithArgs(doc.ClientID, doc.Hash, (*uuid.UUID)(nil), (*string)(nil), (*uuid.UUID)(nil), (*string)(nil), &fileSize).
 		WillReturnRows(
 			pgxmock.NewRows([]string{"id"}).
 				AddRow(doc.ID),
@@ -220,6 +238,10 @@ func TestSubmitCreate(t *testing.T) {
 	cfg.DBPool = pool
 	cfg.DBQueries = repository.New(pool)
 
+	// Set up mock S3 client for file size measurement
+	mockS3 := objectstoremock.NewMockS3Client(t)
+	cfg.StoreClient = mockS3
+
 	svc := New(cfg)
 
 	doc := document.DocumentSummary{
@@ -235,8 +257,16 @@ func TestSubmitCreate(t *testing.T) {
 		Location:  objectstore.Import,
 	}
 
+	// Mock HeadObject to return file size
+	fileSize := int64(2048)
+	mockS3.EXPECT().
+		HeadObject(mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
+			return *in.Bucket == bucket && *in.Key == key.String()
+		})).
+		Return(&s3.HeadObjectOutput{ContentLength: &fileSize}, nil)
+
 	pool.ExpectBegin()
-	pool.ExpectQuery("name: CreateDocument :one").WithArgs(doc.ClientID, doc.Hash, (*uuid.UUID)(nil), (*string)(nil), (*uuid.UUID)(nil), (*string)(nil)).
+	pool.ExpectQuery("name: CreateDocument :one").WithArgs(doc.ClientID, doc.Hash, (*uuid.UUID)(nil), (*string)(nil), (*uuid.UUID)(nil), (*string)(nil), &fileSize).
 		WillReturnRows(
 			pgxmock.NewRows([]string{"id"}).
 				AddRow(doc.ID),
@@ -264,6 +294,10 @@ func TestSubmitCreateExists(t *testing.T) {
 	cfg.DBPool = pool
 	cfg.DBQueries = repository.New(pool)
 
+	// Set up mock S3 client for file size measurement
+	mockS3 := objectstoremock.NewMockS3Client(t)
+	cfg.StoreClient = mockS3
+
 	svc := New(cfg)
 
 	doc := document.DocumentSummary{
@@ -279,6 +313,14 @@ func TestSubmitCreateExists(t *testing.T) {
 		Location:  objectstore.Import,
 	}
 
+	// Mock HeadObject to return file size (still called even for existing docs)
+	fileSize := int64(4096)
+	mockS3.EXPECT().
+		HeadObject(mock.Anything, mock.MatchedBy(func(in *s3.HeadObjectInput) bool {
+			return *in.Bucket == bucket && *in.Key == key.String()
+		})).
+		Return(&s3.HeadObjectOutput{ContentLength: &fileSize}, nil)
+
 	pool.ExpectBegin()
 	pool.ExpectExec("name: AddDocumentEntry :exec").WithArgs(doc.ID, bucket, key.String()).
 		WillReturnResult(pgxmock.NewResult("", 1))
@@ -294,3 +336,200 @@ func TestSubmitCreateExists(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, doc.ID, id)
 }
+
+// TestSubmitCreate_HeadObjectError verifies that HeadObject errors are properly propagated.
+func TestSubmitCreate_HeadObjectError(t *testing.T) {
+	ctx := t.Context()
+
+	pool, err := pgxmock.NewPool()
+	require.NoError(t, err)
+	cfg := &DocInitConfig{}
+	cfg.DBPool = pool
+	cfg.DBQueries = repository.New(pool)
+
+	// Set up mock S3 client that returns an error
+	mockS3 := objectstoremock.NewMockS3Client(t)
+	cfg.StoreClient = mockS3
+
+	svc := New(cfg)
+
+	doc := document.DocumentSummary{
+		ID:       uuid.New(),
+		ClientID: "hello",
+		Hash:     "example_hash",
+	}
+	bucket := "example_bucket"
+	key := objectstore.BucketKey{
+		ClientID:  doc.ClientID,
+		CreatedAt: time.Now().UTC(),
+		EntityID:  uuid.New(),
+		Location:  objectstore.Import,
+	}
+
+	// Mock HeadObject to return an error
+	mockS3.EXPECT().
+		HeadObject(mock.Anything, mock.Anything).
+		Return(nil, errors.New("access denied"))
+
+	_, err = svc.submitCreate(ctx, &createDocumentParams{
+		Hash:   doc.Hash,
+		Bucket: bucket,
+		Key:    key,
+	})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to measure file size")
+	assert.Contains(t, err.Error(), "access denied")
+}
+
+// TestSubmitCreate_NoS3Client verifies that missing S3 client returns an error.
+func TestSubmitCreate_NoS3Client(t *testing.T) {
+	ctx := t.Context()
+
+	pool, err := pgxmock.NewPool()
+	require.NoError(t, err)
+	cfg := &DocInitConfig{}
+	cfg.DBPool = pool
+	cfg.DBQueries = repository.New(pool)
+	// Intentionally not setting StoreClient
+
+	svc := New(cfg)
+
+	doc := document.DocumentSummary{
+		ID:       uuid.New(),
+		ClientID: "hello",
+		Hash:     "example_hash",
+	}
+	bucket := "example_bucket"
+	key := objectstore.BucketKey{
+		ClientID:  doc.ClientID,
+		CreatedAt: time.Now().UTC(),
+		EntityID:  uuid.New(),
+		Location:  objectstore.Import,
+	}
+
+	_, err = svc.submitCreate(ctx, &createDocumentParams{
+		Hash:   doc.Hash,
+		Bucket: bucket,
+		Key:    key,
+	})
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "S3 client not initialized")
+}
+
+// TestMeasureFileSize_NilContentLength verifies handling of nil ContentLength.
+func TestMeasureFileSize_NilContentLength(t *testing.T) {
+	ctx := t.Context()
+
+	pool, err := pgxmock.NewPool()
+	require.NoError(t, err)
+	cfg := &DocInitConfig{}
+	cfg.DBPool = pool
+	cfg.DBQueries = repository.New(pool)
+
+	// Set up mock S3 client that returns nil ContentLength
+	mockS3 := objectstoremock.NewMockS3Client(t)
+	cfg.StoreClient = mockS3
+
+	svc := New(cfg)
+
+	bucket := "example_bucket"
+	key := "example_key"
+
+	mockS3.EXPECT().
+		HeadObject(mock.Anything, mock.Anything).
+		Return(&s3.HeadObjectOutput{ContentLength: nil}, nil)
+
+	_, err = svc.measureFileSize(ctx, bucket, key)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "nil ContentLength")
+}
+
+// TestMeasureFileSize_Integration tests the measureFileSize function using a real
+// S3 client (localstack) to verify that HeadObject correctly returns file sizes.
+// This is an integration test that requires Docker.
+func TestMeasureFileSize_Integration(t *testing.T) {
+	ctx := t.Context()
+
+	// Set up AWS container (localstack) and S3 client
+	cfg := &DocInitConfig{}
+	test.CreateAWSResources(t, cfg)
+
+	svc := New(cfg)
+
+	// Upload a test file with known content
+	bucket := cfg.GetBucket()
+	testKey := fmt.Sprintf("test/filesize/%s.txt", uuid.New().String())
+	testContent := []byte("Hello, World! This is a test file for measuring file size.")
+	expectedSize := int64(len(testContent))
+
+	// Upload the test file to S3
+	_, err := cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(testKey),
+		Body:   bytes.NewReader(testContent),
+	})
+	require.NoError(t, err, "failed to upload test file to S3")
+
+	// Measure the file size using our function
+	measuredSize, err := svc.measureFileSize(ctx, bucket, testKey)
+	require.NoError(t, err, "measureFileSize should not return an error")
+	require.NotNil(t, measuredSize, "measured size should not be nil")
+
+	// Verify the measured size matches expected
+	assert.Equal(t, expectedSize, *measuredSize, "measured file size should match uploaded content size")
+}
+
+// TestMeasureFileSize_Integration_LargeFile tests file size measurement for larger files.
+func TestMeasureFileSize_Integration_LargeFile(t *testing.T) {
+	ctx := t.Context()
+
+	// Set up AWS container (localstack) and S3 client
+	cfg := &DocInitConfig{}
+	test.CreateAWSResources(t, cfg)
+
+	svc := New(cfg)
+
+	// Upload a larger test file (1MB)
+	bucket := cfg.GetBucket()
+	testKey := fmt.Sprintf("test/filesize/large/%s.bin", uuid.New().String())
+	expectedSize := int64(1024 * 1024) // 1MB
+	testContent := make([]byte, expectedSize)
+	for i := range testContent {
+		testContent[i] = byte(i % 256)
+	}
+
+	// Upload the test file to S3
+	_, err := cfg.GetStoreClient().PutObject(ctx, &s3.PutObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(testKey),
+		Body:   bytes.NewReader(testContent),
+	})
+	require.NoError(t, err, "failed to upload large test file to S3")
+
+	// Measure the file size using our function
+	measuredSize, err := svc.measureFileSize(ctx, bucket, testKey)
+	require.NoError(t, err, "measureFileSize should not return an error for large files")
+	require.NotNil(t, measuredSize, "measured size should not be nil")
+
+	// Verify the measured size matches expected
+	assert.Equal(t, expectedSize, *measuredSize, "measured file size should match 1MB")
+}
+
+// TestMeasureFileSize_Integration_NonExistentKey tests error handling for non-existent S3 keys.
+func TestMeasureFileSize_Integration_NonExistentKey(t *testing.T) {
+	ctx := t.Context()
+
+	// Set up AWS container (localstack) and S3 client
+	cfg := &DocInitConfig{}
+	test.CreateAWSResources(t, cfg)
+
+	svc := New(cfg)
+
+	bucket := cfg.GetBucket()
+	nonExistentKey := fmt.Sprintf("test/nonexistent/%s.txt", uuid.New().String())
+
+	// Try to measure size of non-existent file
+	_, err := svc.measureFileSize(ctx, bucket, nonExistentKey)
+	assert.Error(t, err, "measureFileSize should return an error for non-existent key")
+	assert.Contains(t, err.Error(), "HeadObject failed", "error should indicate HeadObject failure")
+}

internal/document/init/service.go

@@ -2,12 +2,14 @@ package documentinit
 
 import (
 	"queryorchestration/internal/serviceconfig"
+	"queryorchestration/internal/serviceconfig/objectstore"
 	documentsync "queryorchestration/internal/serviceconfig/queue/documentsync"
 )
 
 type ConfigProvider interface {
 	serviceconfig.ConfigProvider
 	documentsync.ConfigProvider
+	objectstore.ConfigProvider // Provides GetStoreClient() for S3 access to measure file size
 }
 
 type Service struct {

internal/document/init/service_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"queryorchestration/internal/serviceconfig"
+	"queryorchestration/internal/serviceconfig/objectstore"
 	"queryorchestration/internal/serviceconfig/queue/documentsync"
 
 	"github.com/stretchr/testify/assert"
@@ -12,6 +13,7 @@ import (
 type DocInitConfig struct {
 	serviceconfig.BaseConfig
 	documentsync.DocSyncConfig
+	objectstore.ObjectStoreConfig
 }
 
 func TestNewDocumentInitService(t *testing.T) {