|
|
|
@@ -0,0 +1,246 @@
|
|
|
|
|
package main
|
|
|
|
|
|
|
|
|
|
import (
|
|
|
|
|
"context"
|
|
|
|
|
"crypto/sha256"
|
|
|
|
|
"encoding/hex"
|
|
|
|
|
"encoding/json"
|
|
|
|
|
"fmt"
|
|
|
|
|
"io"
|
|
|
|
|
"log"
|
|
|
|
|
"os"
|
|
|
|
|
"os/signal"
|
|
|
|
|
"syscall"
|
|
|
|
|
"time"
|
|
|
|
|
|
|
|
|
|
"github.com/joho/godotenv"
|
|
|
|
|
"github.com/streadway/amqp"
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
type JobCompletionMessage struct {
|
|
|
|
|
JobID string `json:"job_id"`
|
|
|
|
|
InputPath string `json:"input_path"`
|
|
|
|
|
OutputPath string `json:"output_path"`
|
|
|
|
|
ProcessingTime float64 `json:"processing_time_seconds"`
|
|
|
|
|
Status string `json:"status"`
|
|
|
|
|
Timestamp string `json:"timestamp"`
|
|
|
|
|
FileSize int64 `json:"file_size,omitempty"`
|
|
|
|
|
ImageDimensions string `json:"image_dimensions,omitempty"`
|
|
|
|
|
ErrorMessage string `json:"error_message,omitempty"`
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
type ProcessedImageRecord struct {
|
|
|
|
|
JobID string `json:"job_id"`
|
|
|
|
|
InputPath string `json:"input_path"`
|
|
|
|
|
OutputPath string `json:"output_path"`
|
|
|
|
|
SHA256Hash string `json:"sha256_hash"`
|
|
|
|
|
FileSize int64 `json:"file_size"`
|
|
|
|
|
ImageDimensions string `json:"image_dimensions"`
|
|
|
|
|
ProcessingTime float64 `json:"processing_time_seconds"`
|
|
|
|
|
Status string `json:"status"`
|
|
|
|
|
CreatedAt time.Time `json:"created_at"`
|
|
|
|
|
UpdatedAt time.Time `json:"updated_at"`
|
|
|
|
|
ErrorMessage string `json:"error_message,omitempty"`
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func calculateSHA256(filePath string) (string, error) {
|
|
|
|
|
file, err := os.Open(filePath)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return "", fmt.Errorf("failed to open file: %v", err)
|
|
|
|
|
}
|
|
|
|
|
defer file.Close()
|
|
|
|
|
|
|
|
|
|
hasher := sha256.New()
|
|
|
|
|
if _, err := io.Copy(hasher, file); err != nil {
|
|
|
|
|
return "", fmt.Errorf("failed to calculate hash: %v", err)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
return hex.EncodeToString(hasher.Sum(nil)), nil
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func processCompletionMessage(msg JobCompletionMessage) error {
|
|
|
|
|
log.Printf("Processing completion message for job: %s", msg.JobID)
|
|
|
|
|
|
|
|
|
|
// Create the processed image record
|
|
|
|
|
record := ProcessedImageRecord{
|
|
|
|
|
JobID: msg.JobID,
|
|
|
|
|
InputPath: msg.InputPath,
|
|
|
|
|
OutputPath: msg.OutputPath,
|
|
|
|
|
FileSize: msg.FileSize,
|
|
|
|
|
ImageDimensions: msg.ImageDimensions,
|
|
|
|
|
ProcessingTime: msg.ProcessingTime,
|
|
|
|
|
Status: msg.Status,
|
|
|
|
|
CreatedAt: time.Now(),
|
|
|
|
|
UpdatedAt: time.Now(),
|
|
|
|
|
ErrorMessage: msg.ErrorMessage,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// If the job completed successfully, calculate SHA256
|
|
|
|
|
if msg.Status == "completed" && msg.OutputPath != "" {
|
|
|
|
|
hash, err := calculateSHA256(msg.OutputPath)
|
|
|
|
|
if err != nil {
|
|
|
|
|
log.Printf("Failed to calculate SHA256 for %s: %v", msg.OutputPath, err)
|
|
|
|
|
record.Status = "hash_calculation_failed"
|
|
|
|
|
record.ErrorMessage = fmt.Sprintf("SHA256 calculation failed: %v", err)
|
|
|
|
|
} else {
|
|
|
|
|
record.SHA256Hash = hash
|
|
|
|
|
log.Printf("Calculated SHA256 for %s: %s", msg.OutputPath, hash)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Here you would typically save to a database
|
|
|
|
|
// For now, we'll just log and save to a JSON file as an example
|
|
|
|
|
return saveToDatabase(record)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func saveToDatabase(record ProcessedImageRecord) error {
|
|
|
|
|
// TODO: Replace this with actual database integration
|
|
|
|
|
// This is just a placeholder that saves to a JSON file
|
|
|
|
|
|
|
|
|
|
// Create logs directory if it doesn't exist
|
|
|
|
|
if err := os.MkdirAll("./logs", 0755); err != nil {
|
|
|
|
|
return fmt.Errorf("failed to create logs directory: %v", err)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Append to a JSON lines file
|
|
|
|
|
file, err := os.OpenFile("./logs/processed_images.jsonl", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return fmt.Errorf("failed to open log file: %v", err)
|
|
|
|
|
}
|
|
|
|
|
defer file.Close()
|
|
|
|
|
|
|
|
|
|
// Marshal record to JSON
|
|
|
|
|
data, err := json.Marshal(record)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return fmt.Errorf("failed to marshal record: %v", err)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Write to file
|
|
|
|
|
if _, err := file.Write(append(data, '\n')); err != nil {
|
|
|
|
|
return fmt.Errorf("failed to write record: %v", err)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
log.Printf("Saved record for job %s to database", record.JobID)
|
|
|
|
|
return nil
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func consumeJobCompletions() {
|
|
|
|
|
rabbitURL := os.Getenv("RABBITMQ_URL")
|
|
|
|
|
if rabbitURL == "" {
|
|
|
|
|
log.Println("RABBITMQ_URL not set, using default")
|
|
|
|
|
rabbitURL = "amqp://guest:guest@rabbitmq:5672/"
|
|
|
|
|
} else {
|
|
|
|
|
log.Println("Using RabbitMQ URL:", rabbitURL)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
|
|
|
|
|
defer stop()
|
|
|
|
|
|
|
|
|
|
for {
|
|
|
|
|
select {
|
|
|
|
|
case <-ctx.Done():
|
|
|
|
|
fmt.Println("Shutting down job completion consumer...")
|
|
|
|
|
return
|
|
|
|
|
default:
|
|
|
|
|
// Try to connect to RabbitMQ
|
|
|
|
|
conn, err := amqp.Dial(rabbitURL)
|
|
|
|
|
if err != nil {
|
|
|
|
|
fmt.Println("Failed to connect to RabbitMQ:", err)
|
|
|
|
|
time.Sleep(5 * time.Second)
|
|
|
|
|
continue
|
|
|
|
|
}
|
|
|
|
|
defer conn.Close()
|
|
|
|
|
|
|
|
|
|
ch, err := conn.Channel()
|
|
|
|
|
if err != nil {
|
|
|
|
|
fmt.Println("Failed to open RabbitMQ channel:", err)
|
|
|
|
|
time.Sleep(5 * time.Second)
|
|
|
|
|
continue
|
|
|
|
|
}
|
|
|
|
|
defer ch.Close()
|
|
|
|
|
|
|
|
|
|
// Declare the job_completions queue
|
|
|
|
|
q, err := ch.QueueDeclare(
|
|
|
|
|
"job_completions",
|
|
|
|
|
true, // durable
|
|
|
|
|
false, // auto-delete
|
|
|
|
|
false, // exclusive
|
|
|
|
|
false, // no-wait
|
|
|
|
|
nil,
|
|
|
|
|
)
|
|
|
|
|
if err != nil {
|
|
|
|
|
fmt.Println("Queue declare failed:", err)
|
|
|
|
|
time.Sleep(5 * time.Second)
|
|
|
|
|
continue
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Set QoS to process one message at a time
|
|
|
|
|
err = ch.Qos(
|
|
|
|
|
1, // prefetch count
|
|
|
|
|
0, // prefetch size
|
|
|
|
|
false, // global
|
|
|
|
|
)
|
|
|
|
|
if err != nil {
|
|
|
|
|
fmt.Println("Failed to set QoS:", err)
|
|
|
|
|
time.Sleep(5 * time.Second)
|
|
|
|
|
continue
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
msgs, err := ch.Consume(
|
|
|
|
|
q.Name,
|
|
|
|
|
"completion-consumer", // consumer tag
|
|
|
|
|
false, // manual ack
|
|
|
|
|
false, // not exclusive
|
|
|
|
|
false, // no-local
|
|
|
|
|
false, // no-wait
|
|
|
|
|
nil,
|
|
|
|
|
)
|
|
|
|
|
if err != nil {
|
|
|
|
|
fmt.Println("Queue consume failed:", err)
|
|
|
|
|
time.Sleep(5 * time.Second)
|
|
|
|
|
continue
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Start processing
|
|
|
|
|
fmt.Println("Job completion consumer started...")
|
|
|
|
|
for d := range msgs {
|
|
|
|
|
var completionMsg JobCompletionMessage
|
|
|
|
|
if err := json.Unmarshal(d.Body, &completionMsg); err != nil {
|
|
|
|
|
log.Printf("Invalid completion message payload: %v", err)
|
|
|
|
|
d.Nack(false, false) // discard the message
|
|
|
|
|
continue
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Process the completion message
|
|
|
|
|
if err := processCompletionMessage(completionMsg); err != nil {
|
|
|
|
|
log.Printf("Failed to process completion message: %v", err)
|
|
|
|
|
d.Nack(false, true) // requeue the message
|
|
|
|
|
continue
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Acknowledge successful processing
|
|
|
|
|
d.Ack(false)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// In case the channel is closed for any reason, loop will reconnect
|
|
|
|
|
fmt.Println("Job completion consumer disconnected. Reconnecting...")
|
|
|
|
|
time.Sleep(5 * time.Second)
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func main() {
|
|
|
|
|
// Load environment variables
|
|
|
|
|
err := godotenv.Load()
|
|
|
|
|
if err != nil {
|
|
|
|
|
log.Println(".env file not found or could not be loaded")
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
fmt.Println("Starting Image Processing Completion Consumer")
|
|
|
|
|
fmt.Println("===========================================")
|
|
|
|
|
fmt.Printf("Consumer will process job completion messages from 'job_completions' queue\n")
|
|
|
|
|
fmt.Printf("Processed records will be saved to: ./logs/processed_images.jsonl\n\n")
|
|
|
|
|
|
|
|
|
|
// Start consuming job completion messages
|
|
|
|
|
consumeJobCompletions()
|
|
|
|
|
}
|
F04C