From b2233cf4deb094842370967a726f0d51fe3b6f3c Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Wed, 2 Jun 2021 13:38:39 -0600
Subject: [PATCH] Support Azure account key via environment variable

---
 abs/replica_client.go  |  9 ++++++++-
 cmd/litestream/main.go | 31 ++++++++++++++++---------------
 2 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/abs/replica_client.go b/abs/replica_client.go
index 3c180ba4..576e4662 100644
--- a/abs/replica_client.go
+++ b/abs/replica_client.go
@@ -56,8 +56,14 @@ func (c *ReplicaClient) Init(ctx context.Context) (err error) {
 		return nil
 	}
 
+	// Read account key from environment, if available.
+	accountKey := c.AccountKey
+	if accountKey == "" {
+		accountKey = os.Getenv("LITESTREAM_AZURE_ACCOUNT_KEY")
+	}
+
 	// Authenticate to ACS.
-	credential, err := azblob.NewSharedKeyCredential(c.AccountName, c.AccountKey)
+	credential, err := azblob.NewSharedKeyCredential(c.AccountName, accountKey)
 	if err != nil {
 		return err
 	}
@@ -487,6 +493,7 @@ func (itr *walSegmentIterator) fetch() error {
 		}
 		marker = resp.NextMarker
 
+		println("dbg/wal.fetch", len(resp.Segment.BlobItems))
 		for _, item := range resp.Segment.BlobItems {
 			key := path.Base(item.Name)
 			index, offset, err := litestream.ParseWALSegmentPath(key)
diff --git a/cmd/litestream/main.go b/cmd/litestream/main.go
index bf92e551..f3bd9854 100644
--- a/cmd/litestream/main.go
+++ b/cmd/litestream/main.go
@@ -502,36 +502,37 @@ func newABSReplicaClientFromConfig(c *ReplicaConfig, r *litestream.Replica) (_ *
 		return nil, fmt.Errorf("cannot specify url & bucket for abs replica")
 	}
 
-	bucket, path := c.Bucket, c.Path
+	// Build replica.
+	client := abs.NewReplicaClient()
+	client.AccountName = c.AccountName
+	client.AccountKey = c.AccountKey
+	client.Bucket = c.Bucket
+	client.Path = c.Path
+	client.Endpoint = c.Endpoint
 
 	// Apply settings from URL, if specified.
 	if c.URL != "" {
-		_, uhost, upath, err := ParseReplicaURL(c.URL)
+		u, err := url.Parse(c.URL)
 		if err != nil {
 			return nil, err
 		}
 
-		// Only apply URL parts to field that have not been overridden.
-		if path == "" {
-			path = upath
+		if client.AccountName == "" && u.User != nil {
+			client.AccountName = u.User.Username()
 		}
-		if bucket == "" {
-			bucket = uhost
+		if client.Bucket == "" {
+			client.Bucket = u.Host
+		}
+		if client.Path == "" {
+			client.Path = strings.TrimPrefix(path.Clean(u.Path), "/")
 		}
 	}
 
 	// Ensure required settings are set.
-	if bucket == "" {
+	if client.Bucket == "" {
 		return nil, fmt.Errorf("bucket required for abs replica")
 	}
 
-	// Build replica.
-	client := abs.NewReplicaClient()
-	client.AccountName = c.AccountName
-	client.AccountKey = c.AccountKey
-	client.Bucket = bucket
-	client.Path = path
-	client.Endpoint = c.Endpoint
 	return client, nil
 }