From 589a41ceedfa89e1ff334a969d1beb28cb731de9 Mon Sep 17 00:00:00 2001
From: John Cai <jcai@gitlab.com>
Date: Wed, 3 Nov 2021 14:04:06 -0400
Subject: [PATCH] Add Atomic to push options

push --atomic allows a push to succeed or fail atomically. If one ref
fails, the whole push fails. This commit allows the user to set Atomic
as an option for a push.
---
 options.go                                    |  2 ++
 plumbing/protocol/packp/updreq_encode_test.go | 19 +++++++++++++++++++
 remote.go                                     |  4 ++++
 3 files changed, 25 insertions(+)

diff --git a/options.go b/options.go
index e54889f80..77b74e526 100644
--- a/options.go
+++ b/options.go
@@ -230,6 +230,8 @@ type PushOptions struct {
 	FollowTags bool
 	// PushOptions sets options to be transferred to the server during push.
 	Options map[string]string
+	// Atomic sets option to be an atomic push
+	Atomic bool
 }
 
 // Validate validates the fields and sets the default values.
diff --git a/plumbing/protocol/packp/updreq_encode_test.go b/plumbing/protocol/packp/updreq_encode_test.go
index 6ba004310..4370b794f 100644
--- a/plumbing/protocol/packp/updreq_encode_test.go
+++ b/plumbing/protocol/packp/updreq_encode_test.go
@@ -170,3 +170,22 @@ func (s *UpdReqEncodeSuite) TestPushOptions(c *C) {
 
 	s.testEncode(c, r, expected)
 }
+
+func (s *UpdReqEncodeSuite) TestPushAtomic(c *C) {
+	hash1 := plumbing.NewHash("1ecf0ef2c2dffb796033e5a02219af86ec6584e5")
+	hash2 := plumbing.NewHash("2ecf0ef2c2dffb796033e5a02219af86ec6584e5")
+	name := plumbing.ReferenceName("myref")
+
+	r := NewReferenceUpdateRequest()
+	r.Capabilities.Set(capability.Atomic)
+	r.Commands = []*Command{
+		{Name: name, Old: hash1, New: hash2},
+	}
+
+	expected := pktlines(c,
+		"1ecf0ef2c2dffb796033e5a02219af86ec6584e5 2ecf0ef2c2dffb796033e5a02219af86ec6584e5 myref\x00atomic",
+		pktline.FlushString,
+	)
+
+	s.testEncode(c, r, expected)
+}
diff --git a/remote.go b/remote.go
index 426bde928..503ca61f8 100644
--- a/remote.go
+++ b/remote.go
@@ -326,6 +326,10 @@ func (r *Remote) newReferenceUpdateRequest(
 		}
 	}
 
+	if o.Atomic && ar.Capabilities.Supports(capability.Atomic) {
+		_ = req.Capabilities.Set(capability.Atomic)
+	}
+
 	if err := r.addReferencesToUpdate(o.RefSpecs, localRefs, remoteRefs, req, o.Prune); err != nil {
 		return nil, err
 	}