Enhance company document management with file ID sanitization and detachment functionality

- Updated the `UploadDocuments` method to sanitize document file IDs before saving, ensuring only valid references are stored.
- Introduced `DetachDocumentFileID` method in the `companyService` to remove file IDs from all companies referencing a deleted file, improving data integrity.
- Enhanced the `companyRepository` with a new method to handle the removal of document file IDs from the database.
- Updated the `filestore` handler to utilize the new detachment functionality when files are deleted, ensuring consistent state across domain entities.

This update improves the management of document file IDs within the company domain, enhancing data integrity and reference handling.

cmd/web/main.go

@@ -253,7 +253,7 @@ func main() {
 	contactHandler := contact.NewHandler(contactService, hcaptchaVerifier)
 	cmsHandler := cms.NewHandler(cmsService)
 	kanbanHandler := kanban.NewHandler(kanbanService)
-	fileStoreHandler := filestore.NewHandler(fileStoreService, logger)
+	fileStoreHandler := filestore.NewHandler(fileStoreService, logger, companyService)
 	documentScraperHandler := document_scraper.NewHandler(documentScraperService, logger)
 	dashboardHandler := dashboard.NewHandler(dashboardService)
 	aiPipelineHandler := ai_pipeline.NewHandler(aiPipelineService)

internal/company/documents.go

@@ -75,12 +75,13 @@ func (s *companyService) UploadDocuments(ctx context.Context, companyID, uploade
 			"file_count": len(files),
 		})
 		return &UploadDocumentsResponse{
-			DocumentFileIDs: company.DocumentFileIDs,
+			DocumentFileIDs: s.sanitizeDocumentFileIDs(company.DocumentFileIDs),
 			Uploaded:        uploaded,
 			Failed:          failed,
 		}, errors.New("all file uploads failed")
 	}
 
+	company.DocumentFileIDs = s.sanitizeDocumentFileIDs(company.DocumentFileIDs)
 	company.DocumentFileIDs = mergeDocumentFileIDs(company.DocumentFileIDs, newIDs)
 	if err := s.repository.Update(ctx, company); err != nil {
 		s.logger.Error("Failed to save company documents", map[string]interface{}{
@@ -125,3 +126,40 @@ func mergeDocumentFileIDs(existing, additional []string) []string {
 	appendUnique(additional)
 	return merged
 }
+
+// DetachDocumentFileID removes a file ID from every company that still references it.
+func (s *companyService) DetachDocumentFileID(ctx context.Context, fileID string) error {
+	if err := s.repository.RemoveDocumentFileID(ctx, fileID); err != nil {
+		s.logger.Error("Failed to detach document file ID from companies", map[string]interface{}{
+			"error":   err.Error(),
+			"file_id": fileID,
+		})
+		return err
+	}
+	return nil
+}
+
+func (s *companyService) sanitizeDocumentFileIDs(ids []string) []string {
+	if len(ids) == 0 || s.fileStore == nil {
+		return ids
+	}
+
+	sanitized := make([]string, 0, len(ids))
+	for _, id := range ids {
+		if id == "" {
+			continue
+		}
+		exists, err := s.fileStore.Exists(id)
+		if err != nil {
+			s.logger.Warn("Failed to verify company document file", map[string]interface{}{
+				"error":   err.Error(),
+				"file_id": id,
+			})
+			continue
+		}
+		if exists {
+			sanitized = append(sanitized, id)
+		}
+	}
+	return sanitized
+}

internal/company/repository.go

@@ -11,6 +11,7 @@ import (
 	"tm/pkg/response"
 
 	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/mongo"
 )
 
 // Repository defines the interface for company data operations
@@ -26,12 +27,14 @@ type Repository interface {
 	Delete(ctx context.Context, id string) error
 	UpdateStatus(ctx context.Context, id string, status CompanyStatus) error
 	UpdateVerification(ctx context.Context, id string, isVerified, isCompliant bool, complianceNotes *string) error
+	RemoveDocumentFileID(ctx context.Context, fileID string) error
 }
 
 // companyRepository implements the Repository interface using the MongoDB ORM
 type companyRepository struct {
-	ormRepo orm.Repository[Company]
-	logger  logger.Logger
+	ormRepo    orm.Repository[Company]
+	collection *mongo.Collection
+	logger     logger.Logger
 }
 
 // NewCompanyRepository creates a new company repository
@@ -49,12 +52,15 @@ func NewRepository(mongoManager *orm.ConnectionManager, logger logger.Logger) Re
 		})
 	}
 
+	collection := mongoManager.GetCollection("companies")
+
 	// Create ORM repository
-	ormRepo := orm.NewRepository[Company](mongoManager.GetCollection("companies"), logger)
+	ormRepo := orm.NewRepository[Company](collection, logger)
 
 	return &companyRepository{
-		ormRepo: ormRepo,
-		logger:  logger,
+		ormRepo:    ormRepo,
+		collection: collection,
+		logger:     logger,
 	}
 }
 
@@ -413,3 +419,34 @@ func (r *companyRepository) UpdateVerification(ctx context.Context, id string, i
 
 	return nil
 }
+
+// RemoveDocumentFileID removes a file ID from all companies that reference it.
+func (r *companyRepository) RemoveDocumentFileID(ctx context.Context, fileID string) error {
+	if fileID == "" {
+		return nil
+	}
+
+	filter := bson.M{"document_file_ids": fileID}
+	update := bson.M{
+		"$pull": bson.M{"document_file_ids": fileID},
+		"$set":  bson.M{"updated_at": time.Now().Unix()},
+	}
+
+	result, err := r.collection.UpdateMany(ctx, filter, update)
+	if err != nil {
+		r.logger.Error("Failed to remove document file ID from companies", map[string]interface{}{
+			"error":   err.Error(),
+			"file_id": fileID,
+		})
+		return err
+	}
+
+	if result.ModifiedCount > 0 {
+		r.logger.Info("Removed document file ID from companies", map[string]interface{}{
+			"file_id":        fileID,
+			"modified_count": result.ModifiedCount,
+		})
+	}
+
+	return nil
+}

internal/company/service.go

@@ -35,6 +35,9 @@ type Service interface {
 	// UploadDocuments uploads files and attaches them to a company
 	UploadDocuments(ctx context.Context, companyID, uploadedBy string, files []*multipart.FileHeader) (*UploadDocumentsResponse, error)
 
+	// DetachDocumentFileID removes a file ID from companies when the file is deleted
+	DetachDocumentFileID(ctx context.Context, fileID string) error
+
 	// StartAIOnboarding sends company data to the AI team for recommendation onboarding
 	StartAIOnboarding(ctx context.Context, companyID string) (*OnboardingResponse, error)
 
@@ -121,7 +124,7 @@ func (s *companyService) Create(ctx context.Context, form *CompanyForm) (*Compan
 		Address:            s.convertAddressForm(form.Address),
 		Phone:              form.Phone,
 		Email:              form.Email,
-		DocumentFileIDs:    form.DocumentFileIDs,
+		DocumentFileIDs:    s.sanitizeDocumentFileIDs(form.DocumentFileIDs),
 		Tags:               s.convertCompanyTagsForm(form.Tags),
 		IsVerified:         false,
 		IsCompliant:        false,
@@ -183,6 +186,8 @@ func (s *companyService) GetByID(ctx context.Context, id string) (*CompanyRespon
 		})
 	}
 
+	company.DocumentFileIDs = s.sanitizeDocumentFileIDs(company.DocumentFileIDs)
+
 	return company.ToResponse(convertedCategories), nil
 }
 
@@ -284,7 +289,7 @@ func (s *companyService) Update(ctx context.Context, id string, form *CompanyFor
 	}
 
 	if form.DocumentFileIDs != nil {
-		company.DocumentFileIDs = form.DocumentFileIDs
+		company.DocumentFileIDs = s.sanitizeDocumentFileIDs(form.DocumentFileIDs)
 	}
 
 	if form.Tags != nil {
@@ -497,6 +502,8 @@ func (s *companyService) GetMyCompany(ctx context.Context, id string) (*CompanyP
 		return nil, errors.New("company not found")
 	}
 
+	company.DocumentFileIDs = s.sanitizeDocumentFileIDs(company.DocumentFileIDs)
+
 	return company.ToProfileResponse(), nil
 }
 

pkg/filestore/handler.go

@@ -1,6 +1,7 @@
 package filestore
 
 import (
+	"context"
 	"errors"
 	"net/http"
 	"strconv"
@@ -11,17 +12,24 @@ import (
 	"github.com/labstack/echo/v4"
 )
 
+// DocumentReferenceCleaner removes stored file IDs from domain entities when a file is deleted.
+type DocumentReferenceCleaner interface {
+	DetachDocumentFileID(ctx context.Context, fileID string) error
+}
+
 // Handler handles HTTP requests for file operations
 type Handler struct {
-	service FileStoreService
-	logger  logger.Logger
+	service           FileStoreService
+	logger            logger.Logger
+	referenceCleaner  DocumentReferenceCleaner
 }
 
 // NewHandler creates a new file handler
-func NewHandler(service FileStoreService, logger logger.Logger) *Handler {
+func NewHandler(service FileStoreService, logger logger.Logger, referenceCleaner DocumentReferenceCleaner) *Handler {
 	return &Handler{
-		service: service,
-		logger:  logger,
+		service:          service,
+		logger:           logger,
+		referenceCleaner: referenceCleaner,
 	}
 }
 
@@ -417,6 +425,15 @@ func (h *Handler) Delete(c echo.Context) error {
 		})
 	}
 
+	if h.referenceCleaner != nil {
+		if cleanErr := h.referenceCleaner.DetachDocumentFileID(c.Request().Context(), fileID); cleanErr != nil {
+			h.logger.Warn("Failed to detach deleted file from domain references", map[string]interface{}{
+				"file_id": fileID,
+				"error":   cleanErr.Error(),
+			})
+		}
+	}
+
 	return c.JSON(http.StatusOK, map[string]string{
 		"message": "File deleted successfully",
 	})