1 files changed, 52 insertions, 23 deletions

diff --git a/internal/storage/b2.go b/internal/storage/b2.go
index 0df321f..7d7ad0b 100644
--- a/internal/storage/b2.go
+++ b/internal/storage/b2.go
@@ -16,12 +16,13 @@ import (
 var _ StorageClient = &B2StorageClient{}
 
 type B2StorageClient struct {
-	prefix string
-	bucket *b2.Bucket
-	b2cl   *b2.Client
+	prefix       string
+	bucket       *b2.Bucket
+	b2cl         *b2.Client
+	draftModeCfg *config.DraftModeConfig
 }
 
-func NewB2StorageClient(cfg *config.StorageConfig) (StorageClient, error) {
+func NewB2StorageClient(cfg *config.StorageConfig, draftModeCfg *config.DraftModeConfig) (StorageClient, error) {
 	if cfg.Type != "b2" {
 		return nil, fmt.Errorf("invalid storage type for B2InputClient")
 	}
@@ -37,7 +38,9 @@ func NewB2StorageClient(cfg *config.StorageConfig) (StorageClient, error) {
 		return nil, err
 	}
 
-	return &B2StorageClient{b2cl: b2cl, bucket: bucket, prefix: b2cfg.Prefix}, nil
+	draftModeCfgCopy := *draftModeCfg
+
+	return &B2StorageClient{b2cl: b2cl, bucket: bucket, prefix: b2cfg.Prefix, draftModeCfg: &draftModeCfgCopy}, nil
 }
 
 func (sc *B2StorageClient) Scan() ([]string, error) {
@@ -60,11 +63,16 @@ func (sc *B2StorageClient) Scan() ([]string, error) {
 			continue
 		}
 
-		if !strings.HasSuffix(obj.Name(), ".md") {
+		name := obj.Name()
+		if !strings.HasSuffix(name, ".md") {
+			continue
+		}
+
+		if sc.draftModeCfg.Enabled && !strings.HasSuffix(name, sc.draftModeCfg.DraftSuffix) {
 			continue
 		}
 
-		filePaths = append(filePaths, strings.TrimPrefix(obj.Name(), sc.prefix))
+		filePaths = append(filePaths, strings.TrimPrefix(name, sc.prefix))
 	}
 
 	if err := iter.Err(); err != nil {
@@ -88,13 +96,13 @@ func (sc *B2StorageClient) GetReader(path string) (io.ReadCloser, int64, error)
 }
 
 func (sc *B2StorageClient) WriteMetadata(path string, metadata *frontmatter.Metadata) error {
-	obj := sc.bucket.Object(sc.prefix + path)
-	if obj == nil {
-		return fmt.Errorf("failed to reference object in B2 bucket")
+	draft := sc.bucket.Object(sc.prefix + path)
+	if draft == nil {
+		return fmt.Errorf("failed to reference draft object in B2 bucket")
 	}
-	oldAttrs, err := obj.Attrs(context.Background())
+	draftAttrs, err := draft.Attrs(context.Background())
 	if err != nil {
-		return fmt.Errorf("error getting attributes of an object: %w", err)
+		return fmt.Errorf("error getting attributes of a draft object: %w", err)
 	}
 
 	geolocationParts := strings.Split(metadata.Geolocation, " ")
@@ -125,16 +133,27 @@ func (sc *B2StorageClient) WriteMetadata(path string, metadata *frontmatter.Meta
 			"thumbnail":                 metadata.Thumbnail,
 			"tags":                      strings.Join(metadata.Tags, ","),
 			"geolocation":               metadata.Geolocation,
-			"metadata-last-update-sha1": oldAttrs.SHA1,
+			"metadata-last-update-sha1": draftAttrs.SHA1,
 		}}
 
-	reader := obj.NewReader(context.Background())
-	content := make([]byte, oldAttrs.Size)
+	reader := draft.NewReader(context.Background())
+	content := make([]byte, draftAttrs.Size)
 	if _, err = reader.Read(content); err != nil {
-		return fmt.Errorf("failed to read an object back for writing (required for attribute setting): %w", err)
+		return fmt.Errorf("failed to read a draft object back for writing (required for attribute setting): %w", err)
+	}
+
+	var prod *b2.Object
+	if sc.draftModeCfg.Enabled {
+		prodPath := strings.TrimSuffix(path, sc.draftModeCfg.DraftSuffix) + sc.draftModeCfg.ProdSuffix
+		prod = sc.bucket.Object(sc.prefix + prodPath)
+		if prod == nil {
+			return fmt.Errorf("failed to reference prod object in B2 bucket")
+		}
+	} else {
+		prod = draft
 	}
 
-	writer := obj.NewWriter(context.Background(), b2.WithAttrsOption(attrs))
+	writer := prod.NewWriter(context.Background(), b2.WithAttrsOption(attrs))
 	defer writer.Close()
 	if _, err := writer.Write(content); err != nil {
 		return fmt.Errorf("failed to write an object back after attribute settings: %w", err)
@@ -143,21 +162,31 @@ func (sc *B2StorageClient) WriteMetadata(path string, metadata *frontmatter.Meta
 	return nil
 }
 
-func (sc *B2StorageClient) FileHasChanged(path string) bool {
-	obj := sc.bucket.Object(sc.prefix + path)
-	if obj == nil {
+func (sc *B2StorageClient) CompareDraftAndProd(path string) (changed bool) {
+	draft := sc.bucket.Object(sc.prefix + path)
+	if draft == nil {
+		return false
+	}
+
+	prodPath := strings.TrimSuffix(path, sc.draftModeCfg.DraftSuffix) + sc.draftModeCfg.ProdSuffix
+	prod := sc.bucket.Object(sc.prefix + prodPath)
+	if prod == nil {
 		return true
 	}
 
-	attrs, err := obj.Attrs(context.Background())
+	draftAttrs, err := draft.Attrs(context.Background())
+	if err != nil {
+		return false
+	}
+	prodAttrs, err := prod.Attrs(context.Background())
 	if err != nil {
 		return true
 	}
 
-	lastUpdateSha1, ok := attrs.Info["metadata-last-update-sha1"]
+	lastUpdateSha1, ok := prodAttrs.Info["metadata-last-update-sha1"]
 	if !ok {
 		return true
 	}
 
-	return attrs.SHA1 == lastUpdateSha1
+	return draftAttrs.SHA1 != lastUpdateSha1
 }