Merge pull request #3 from jvale/azure_retention_policy

Add support for Azure Blob retention policy

README.md

@@ -41,6 +41,7 @@ Usage of consul-snapshotter:
       --azure-blob.storage-access-key string   The Azure Blob storage access key to use
       --azure-blob.storage-sas-token           The Azure Blob storage SAS token to use instead of an access key
       --azure-blob.storage-account string      The Azure Blob storage account to use
+      --azure-blob.retention-period duration   The duration that Azure Blob snapshots need to be retained (default "0s" - keep forever)
       --configdir string                       The path to look for the configuration file (default ".")
       --consul.lock-key string                 The Key to use in the KV lock (default "consul-snapshotter/.lock")
       --consul.lock-timeout duration           The timeout for the session lock (default 10m0s)

azure/azure.go

@@ -47,7 +47,7 @@ func AuthenticateSASToken(containerName string, c *config) pipeline.Pipeline {
 	return azblob.NewPipeline(azblob.NewAnonymousCredential(), azblob.PipelineOptions{})
 }
 
-func UploadBlob(srcFile, destFile, containerName string, c *config) (int, error) {
+func GetContainerURL(containerName string, c *config) (azblob.ContainerURL, error) {
 	var p pipeline.Pipeline
 	var queryParameters string
 
@@ -58,7 +58,7 @@ func UploadBlob(srcFile, destFile, containerName string, c *config) (int, error)
 		var err error
 		p, err = AuthenticateAccountKey(containerName, c)
 		if err != nil {
-			return 0, err
+			return azblob.ContainerURL{}, err
 		}
 	}
 
@@ -68,7 +68,14 @@ func UploadBlob(srcFile, destFile, containerName string, c *config) (int, error)
 		fmt.Sprintf("https://%s.blob.core.windows.net/%s%s", c.accountName, containerName, queryParameters))
 
 	// Create a ContainerURL object using the container URL and a request pipeline
-	containerURL := azblob.NewContainerURL(*URL, p)
+	return azblob.NewContainerURL(*URL, p), nil
+}
+
+func UploadBlob(srcFile, destFile, containerName string, c *config) (int, error) {
+	containerURL, err := GetContainerURL(containerName, c)
+	if err != nil {
+		return 0, err
+	}
 
 	ctx := context.Background()
 
@@ -93,3 +100,47 @@ func UploadBlob(srcFile, destFile, containerName string, c *config) (int, error)
 
 	return 1, nil
 }
+
+func ListBlobs(containerName string, c *config) ([]azblob.BlobItem, error) {
+	var results = make([]azblob.BlobItem, 0)
+
+	containerURL, err := GetContainerURL(containerName, c)
+	if err != nil {
+		return results, err
+	}
+
+	ctx := context.Background()
+
+	for marker := (azblob.Marker{}); marker.NotDone(); {
+		// Get a result segment starting with the blob indicated by the current Marker.
+		listBlob, err := containerURL.ListBlobsFlatSegment(ctx, marker, azblob.ListBlobsSegmentOptions{})
+		if err != nil {
+			return results, fmt.Errorf("Error uploading file: %s", err)
+		}
+
+		// ListBlobs returns the start of the next segment; you MUST use this to get
+		// the next segment (after processing the current result segment).
+		marker = listBlob.NextMarker
+
+		results = append(results, listBlob.Segment.BlobItems...)
+	}
+
+	return results, nil
+}
+
+func DeleteBlob(containerName string, blob azblob.BlobItem, c *config) error {
+	containerURL, err := GetContainerURL(containerName, c)
+	if err != nil {
+		return err
+	}
+
+	blobUrl := containerURL.NewBlobURL(blob.Name)
+	ctx := context.Background()
+
+	_, err = blobUrl.Delete(ctx, azblob.DeleteSnapshotsOptionInclude, azblob.BlobAccessConditions{})
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

config.go

@@ -24,11 +24,12 @@ type localOutputConfig struct {
 }
 
 type azureOutputConfig struct {
-	ContainerName    string `json:"container-name"`
-	ContainerPath    string `json:"container-path"`
-	StorageAccount   string `json:"azure-storage-account"`
-	StorageAccessKey string `json:"azure-storage-access-key"`
-	StorageSASToken  string `json:"azure-storage-sas-token"`
+	ContainerName    string        `json:"container-name"`
+	ContainerPath    string        `json:"container-path"`
+	StorageAccount   string        `json:"azure-storage-account"`
+	StorageAccessKey string        `json:"azure-storage-access-key"`
+	StorageSASToken  string        `json:"azure-storage-sas-token"`
+	RetentionPeriod  time.Duration `json:"retention-period"`
 }
 
 type config struct {
@@ -77,6 +78,7 @@ func (c *config) loadConfig() error {
 	viper.SetDefault("outputs", []string{"local"})
 	viper.SetDefault("local.destination-path", ".")
 	viper.SetDefault("local.retention-period", "0s")
+	viper.SetDefault("azure-blob.retention-period", "0s")
 
 	// read command flags
 	regFlagString("configdir", ".", "The path to look for the configuration file")
@@ -93,6 +95,7 @@ func (c *config) loadConfig() error {
 	regFlagString("azure-blob.storage-account", "", "The Azure Blob storage account to use")
 	regFlagString("azure-blob.storage-access-key", "", "The Azure Blob storage access key to use (mutually exclusive with azure-blob.storage-sas-token)")
 	regFlagString("azure-blob.storage-sas-token", "", "The Azure Blob storage SAS token to use (mutually exclusive with azure-blob.storage-access-key)")
+	regFlagDuration("azure-blob.retention-period", viper.GetDuration("azure-blob.retention-period"), "The duration that Azure Blob snapshots need to be retained (default \"0s\" - keep forever)")
 	regFlagString("local.destination-path", viper.GetString("local.destination-path"), "The local path where to save the snapshots")
 	regFlagDuration("local.retention-period", viper.GetDuration("local.retention-period"), "The duration that Local snapshots need to be retained (default \"0s\" - keep forever)")
 	regFlagBoolP("help", "h", false, "Prints this help message")
@@ -139,6 +142,7 @@ func (c *config) loadConfig() error {
 	azureOutputConfig.StorageAccount = viper.GetString("azure-blob.storage-account")
 	azureOutputConfig.StorageAccessKey = viper.GetString("azure-blob.storage-access-key")
 	azureOutputConfig.StorageSASToken = viper.GetString("azure-blob.storage-sas-token")
+	azureOutputConfig.RetentionPeriod = viper.GetDuration("azure-blob.retention-period")
 
 	// Local output config
 	localOutputConfig := &localOutputConfig{}

main.go

@@ -153,8 +153,13 @@ func processOutputs(snap string, c *config) {
 				StorageAccount:   c.AzureOutputConfig.StorageAccount,
 				StorageAccessKey: c.AzureOutputConfig.StorageAccessKey,
 				StorageSASToken:  c.AzureOutputConfig.StorageSASToken,
+				RetentionPeriod:  c.AzureOutputConfig.RetentionPeriod,
 			}
 			o.Save(snap)
+
+			if err := o.ApplyRetentionPolicy(); err != nil {
+				log.Println("Error applying retention policy: ", err)
+			}
 		}
 	}
 }

outputs/azure_blob.go

@@ -1,8 +1,10 @@
 package outputs
 
 import (
+	"fmt"
 	"log"
 	"path"
+	"time"
 
 	"github.com/ruizink/consul-snapshotter/azure"
 )
@@ -14,6 +16,7 @@ type AzureBlobOutput struct {
 	StorageAccount   string
 	StorageAccessKey string
 	StorageSASToken  string
+	RetentionPeriod  time.Duration
 }
 
 func (o *AzureBlobOutput) Save(snap string) {
@@ -30,3 +33,33 @@ func (o *AzureBlobOutput) Save(snap string) {
 	}
 	log.Println("Uploaded snapshot to:", destFile)
 }
+
+func (o *AzureBlobOutput) ApplyRetentionPolicy() error {
+	log.Println(fmt.Sprintf("Azure Blob Storage retention: %v", o.RetentionPeriod))
+	if o.RetentionPeriod <= 0 {
+		return nil
+	}
+
+	log.Println(fmt.Sprintf("Applying retention policy (remove files older than %v) in Azure Blob Storage", o.RetentionPeriod))
+
+	config, err := azure.AzureConfig(o.StorageAccount, o.StorageAccessKey, o.StorageSASToken)
+	if err != nil {
+		return err
+	}
+
+	blobList, err := azure.ListBlobs(o.ContainerName, config)
+
+	for _, blob := range blobList {
+		if time.Now().Sub(blob.Properties.LastModified) <= o.RetentionPeriod {
+			continue
+		}
+
+		log.Println("Removing from Azure Blob Storage: " + blob.Name)
+		err := azure.DeleteBlob(o.ContainerName, blob, config)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

outputs/local.go

@@ -13,7 +13,6 @@ import (
 type LocalOutput struct {
 	DestinationPath string
 	Filename        string
-	RetentionKeep   int
 	RetentionPeriod time.Duration
 }