Implement context archive handling and metadata extraction for render jobs. Add functionality to check for Blender availability, create context archives, and extract metadata from .blend files. Update job creation and retrieval processes to support new metadata structure and context file management. Enhance client-side components to display context files and integrate new API endpoints for context handling.

2025-11-24 10:02:13 -06:00
parent f9ff4d0138
commit a029714e08
13 changed files with 3887 additions and 856 deletions
--- a/internal/storage/storage.go
+++ b/internal/storage/storage.go
@@ -1,7 +1,9 @@
 package storage

 import (
+	"archive/tar"
 	"archive/zip"
+	"compress/gzip"
 	"fmt"
 	"io"
 	"os"
@@ -194,3 +196,377 @@ func (s *Storage) ExtractZip(zipPath, destDir string) ([]string, error) {
 	return extractedFiles, nil
 }

+// findCommonPrefix finds the common leading directory prefix if all paths share the same first-level directory
+// Returns the prefix to strip (with trailing slash) or empty string if no common prefix
+func findCommonPrefix(relPaths []string) string {
+	if len(relPaths) == 0 {
+		return ""
+	}
+
+	// Get the first path component of each path
+	firstComponents := make([]string, 0, len(relPaths))
+	for _, path := range relPaths {
+		parts := strings.Split(filepath.ToSlash(path), "/")
+		if len(parts) > 0 && parts[0] != "" {
+			firstComponents = append(firstComponents, parts[0])
+		} else {
+			// If any path is at root level, no common prefix
+			return ""
+		}
+	}
+
+	// Check if all first components are the same
+	if len(firstComponents) == 0 {
+		return ""
+	}
+
+	commonFirst := firstComponents[0]
+	for _, comp := range firstComponents {
+		if comp != commonFirst {
+			// Not all paths share the same first directory
+			return ""
+		}
+	}
+
+	// All paths share the same first directory - return it with trailing slash
+	return commonFirst + "/"
+}
+
+// isBlenderSaveFile checks if a filename is a Blender save file (.blend1, .blend2, etc.)
+// Returns true for files like "file.blend1", "file.blend2", but false for "file.blend"
+func isBlenderSaveFile(filename string) bool {
+	lower := strings.ToLower(filename)
+	// Check if it ends with .blend followed by one or more digits
+	// Pattern: *.blend[digits]
+	if !strings.HasSuffix(lower, ".blend") {
+		// Doesn't end with .blend, check if it ends with .blend + digits
+		idx := strings.LastIndex(lower, ".blend")
+		if idx == -1 {
+			return false
+		}
+		// Check if there are digits after .blend
+		suffix := lower[idx+len(".blend"):]
+		if len(suffix) == 0 {
+			return false
+		}
+		// All remaining characters must be digits
+		for _, r := range suffix {
+			if r < '0' || r > '9' {
+				return false
+			}
+		}
+		return true
+	}
+	// Ends with .blend exactly - this is a regular blend file, not a save file
+	return false
+}
+
+// CreateJobContext creates a tar.gz archive containing all job input files
+// Filters out Blender save files (.blend1, .blend2, etc.)
+// Uses temporary directories and streaming to handle large files efficiently
+func (s *Storage) CreateJobContext(jobID int64) (string, error) {
+	jobPath := s.JobPath(jobID)
+	contextPath := filepath.Join(jobPath, "context.tar.gz")
+
+	// Create temporary directory for staging
+	tmpDir, err := os.MkdirTemp("", "fuego-context-*")
+	if err != nil {
+		return "", fmt.Errorf("failed to create temporary directory: %w", err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	// Collect all files from job directory, excluding the context file itself and Blender save files
+	var filesToInclude []string
+	err = filepath.Walk(jobPath, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+
+		// Skip directories
+		if info.IsDir() {
+			return nil
+		}
+
+		// Skip the context file itself if it exists
+		if path == contextPath {
+			return nil
+		}
+
+		// Skip Blender save files
+		if isBlenderSaveFile(info.Name()) {
+			return nil
+		}
+
+		// Get relative path from job directory
+		relPath, err := filepath.Rel(jobPath, path)
+		if err != nil {
+			return err
+		}
+
+		// Sanitize path - ensure it doesn't escape the job directory
+		cleanRelPath := filepath.Clean(relPath)
+		if strings.HasPrefix(cleanRelPath, "..") {
+			return fmt.Errorf("invalid file path: %s", relPath)
+		}
+
+		filesToInclude = append(filesToInclude, path)
+		return nil
+	})
+	if err != nil {
+		return "", fmt.Errorf("failed to walk job directory: %w", err)
+	}
+
+	if len(filesToInclude) == 0 {
+		return "", fmt.Errorf("no files found to include in context")
+	}
+
+	// Create the tar.gz file using streaming
+	contextFile, err := os.Create(contextPath)
+	if err != nil {
+		return "", fmt.Errorf("failed to create context file: %w", err)
+	}
+	defer contextFile.Close()
+
+	gzWriter := gzip.NewWriter(contextFile)
+	defer gzWriter.Close()
+
+	tarWriter := tar.NewWriter(gzWriter)
+	defer tarWriter.Close()
+
+	// Add each file to the tar archive
+	for _, filePath := range filesToInclude {
+		file, err := os.Open(filePath)
+		if err != nil {
+			return "", fmt.Errorf("failed to open file %s: %w", filePath, err)
+		}
+
+		info, err := file.Stat()
+		if err != nil {
+			file.Close()
+			return "", fmt.Errorf("failed to stat file %s: %w", filePath, err)
+		}
+
+		// Get relative path for tar header
+		relPath, err := filepath.Rel(jobPath, filePath)
+		if err != nil {
+			file.Close()
+			return "", fmt.Errorf("failed to get relative path for %s: %w", filePath, err)
+		}
+
+		// Normalize path separators for tar (use forward slashes)
+		tarPath := filepath.ToSlash(relPath)
+
+		// Create tar header
+		header, err := tar.FileInfoHeader(info, "")
+		if err != nil {
+			file.Close()
+			return "", fmt.Errorf("failed to create tar header for %s: %w", filePath, err)
+		}
+		header.Name = tarPath
+
+		// Write header
+		if err := tarWriter.WriteHeader(header); err != nil {
+			file.Close()
+			return "", fmt.Errorf("failed to write tar header for %s: %w", filePath, err)
+		}
+
+		// Copy file contents using streaming
+		if _, err := io.Copy(tarWriter, file); err != nil {
+			file.Close()
+			return "", fmt.Errorf("failed to write file %s to tar: %w", filePath, err)
+		}
+
+		file.Close()
+	}
+
+	// Ensure all data is flushed
+	if err := tarWriter.Close(); err != nil {
+		return "", fmt.Errorf("failed to close tar writer: %w", err)
+	}
+	if err := gzWriter.Close(); err != nil {
+		return "", fmt.Errorf("failed to close gzip writer: %w", err)
+	}
+	if err := contextFile.Close(); err != nil {
+		return "", fmt.Errorf("failed to close context file: %w", err)
+	}
+
+	return contextPath, nil
+}
+
+// CreateJobContextFromDir creates a context archive (tar.gz) from files in a source directory
+// This is used during upload to immediately create the context archive as the primary artifact
+// excludeFiles is a set of relative paths (from sourceDir) to exclude from the context
+func (s *Storage) CreateJobContextFromDir(sourceDir string, jobID int64, excludeFiles ...string) (string, error) {
+	jobPath := s.JobPath(jobID)
+	contextPath := filepath.Join(jobPath, "context.tar.gz")
+
+	// Ensure job directory exists
+	if err := os.MkdirAll(jobPath, 0755); err != nil {
+		return "", fmt.Errorf("failed to create job directory: %w", err)
+	}
+
+	// Build set of files to exclude (normalize paths)
+	excludeSet := make(map[string]bool)
+	for _, excludeFile := range excludeFiles {
+		// Normalize the exclude path
+		excludePath := filepath.Clean(excludeFile)
+		excludeSet[excludePath] = true
+		// Also add with forward slash for cross-platform compatibility
+		excludeSet[filepath.ToSlash(excludePath)] = true
+	}
+
+	// Collect all files from source directory, excluding Blender save files and excluded files
+	var filesToInclude []string
+	err := filepath.Walk(sourceDir, func(path string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+
+		// Skip directories
+		if info.IsDir() {
+			return nil
+		}
+
+		// Skip Blender save files
+		if isBlenderSaveFile(info.Name()) {
+			return nil
+		}
+
+		// Get relative path from source directory
+		relPath, err := filepath.Rel(sourceDir, path)
+		if err != nil {
+			return err
+		}
+
+		// Sanitize path - ensure it doesn't escape the source directory
+		cleanRelPath := filepath.Clean(relPath)
+		if strings.HasPrefix(cleanRelPath, "..") {
+			return fmt.Errorf("invalid file path: %s", relPath)
+		}
+
+		// Check if this file should be excluded
+		if excludeSet[cleanRelPath] || excludeSet[filepath.ToSlash(cleanRelPath)] {
+			return nil
+		}
+
+		filesToInclude = append(filesToInclude, path)
+		return nil
+	})
+	if err != nil {
+		return "", fmt.Errorf("failed to walk source directory: %w", err)
+	}
+
+	if len(filesToInclude) == 0 {
+		return "", fmt.Errorf("no files found to include in context archive")
+	}
+
+	// Collect relative paths to find common prefix
+	relPaths := make([]string, 0, len(filesToInclude))
+	for _, filePath := range filesToInclude {
+		relPath, err := filepath.Rel(sourceDir, filePath)
+		if err != nil {
+			return "", fmt.Errorf("failed to get relative path for %s: %w", filePath, err)
+		}
+		relPaths = append(relPaths, relPath)
+	}
+
+	// Find and strip common leading directory if all files share one
+	commonPrefix := findCommonPrefix(relPaths)
+
+	// Validate that there's exactly one .blend file at the root level after prefix stripping
+	blendFilesAtRoot := 0
+	for _, relPath := range relPaths {
+		tarPath := filepath.ToSlash(relPath)
+		// Strip common prefix if present
+		if commonPrefix != "" && strings.HasPrefix(tarPath, commonPrefix) {
+			tarPath = strings.TrimPrefix(tarPath, commonPrefix)
+		}
+		
+		// Check if it's a .blend file at root (no path separators after prefix stripping)
+		if strings.HasSuffix(strings.ToLower(tarPath), ".blend") {
+			// Check if it's at root level (no directory separators)
+			if !strings.Contains(tarPath, "/") {
+				blendFilesAtRoot++
+			}
+		}
+	}
+
+	if blendFilesAtRoot == 0 {
+		return "", fmt.Errorf("no .blend file found at root level in context archive")
+	}
+	if blendFilesAtRoot > 1 {
+		return "", fmt.Errorf("multiple .blend files found at root level in context archive (found %d, expected 1)", blendFilesAtRoot)
+	}
+
+	// Create the tar.gz file using streaming
+	contextFile, err := os.Create(contextPath)
+	if err != nil {
+		return "", fmt.Errorf("failed to create context file: %w", err)
+	}
+	defer contextFile.Close()
+
+	gzWriter := gzip.NewWriter(contextFile)
+	defer gzWriter.Close()
+
+	tarWriter := tar.NewWriter(gzWriter)
+	defer tarWriter.Close()
+
+	// Add each file to the tar archive
+	for i, filePath := range filesToInclude {
+		file, err := os.Open(filePath)
+		if err != nil {
+			return "", fmt.Errorf("failed to open file %s: %w", filePath, err)
+		}
+
+		info, err := file.Stat()
+		if err != nil {
+			file.Close()
+			return "", fmt.Errorf("failed to stat file %s: %w", filePath, err)
+		}
+
+		// Get relative path and strip common prefix if present
+		relPath := relPaths[i]
+		tarPath := filepath.ToSlash(relPath)
+		
+		// Strip common prefix if found
+		if commonPrefix != "" && strings.HasPrefix(tarPath, commonPrefix) {
+			tarPath = strings.TrimPrefix(tarPath, commonPrefix)
+		}
+
+		// Create tar header
+		header, err := tar.FileInfoHeader(info, "")
+		if err != nil {
+			file.Close()
+			return "", fmt.Errorf("failed to create tar header for %s: %w", filePath, err)
+		}
+		header.Name = tarPath
+
+		// Write header
+		if err := tarWriter.WriteHeader(header); err != nil {
+			file.Close()
+			return "", fmt.Errorf("failed to write tar header for %s: %w", filePath, err)
+		}
+
+		// Copy file contents using streaming
+		if _, err := io.Copy(tarWriter, file); err != nil {
+			file.Close()
+			return "", fmt.Errorf("failed to write file %s to tar: %w", filePath, err)
+		}
+
+		file.Close()
+	}
+
+	// Ensure all data is flushed
+	if err := tarWriter.Close(); err != nil {
+		return "", fmt.Errorf("failed to close tar writer: %w", err)
+	}
+	if err := gzWriter.Close(); err != nil {
+		return "", fmt.Errorf("failed to close gzip writer: %w", err)
+	}
+	if err := contextFile.Close(); err != nil {
+		return "", fmt.Errorf("failed to close context file: %w", err)
+	}
+
+	return contextPath, nil
+}
+