Skip to content

Commit 1079bd6

Browse files
vroldanbetvroldanbet
committed
updates bulk-delete to use DeleteRelationships new optional-limit
former implementation of bulk-delete was designed to support deleting large amounts of relationships in a batched way. This commit replaces the implementation that relied on ReadRelationships with the new native DeleteRelationships optional-limit. The contract is kept mostly intact, except the interactive prompt is removed as it was considered bad UX. Instead, the same limits of 1000 before the user is prompted is retained, and instead zed will fail, indicating the user to bump the limit with the new --optional-limit flag, or to use --force. The --optional-limit flag sets the corresponding `DeleteRelationship` request field, and `--force` now sets the `OptionalAllowPartialDeletions` field. If set to true, all relationships will be deleted in batches defined by `--optional-limit`, which continues to default to 1000 elements. It also adds a progress bar to provide visual feedback to the user while deletion takes place.
1 parent 581f04f commit 1079bd6

File tree

12 files changed

+367
-255
lines changed

12 files changed

+367
-255
lines changed

‎go.mod‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ toolchain go1.21.3
66

77
require (
88
github.com/99designs/keyring v1.2.2
9-
github.com/AlecAivazis/survey/v2 v2.3.7
109
github.com/TylerBrock/colorjson v0.0.0-20200706003622-8a50f05110d2
1110
github.com/authzed/authzed-go v0.11.2-0.20240320204618-9622b72a72c6
1211
github.com/authzed/grpcutil v0.0.0-20240123194739-2ea1e3d2d98b
@@ -34,6 +33,7 @@ require (
3433
golang.org/x/mod v0.15.0
3534
golang.org/x/sync v0.6.0
3635
golang.org/x/term v0.17.0
36+
google.golang.org/genproto/googleapis/rpc v0.0.0-20240205150955-31a09d347014
3737
google.golang.org/grpc v1.62.1
3838
google.golang.org/protobuf v1.33.0
3939
gopkg.in/yaml.v3 v3.0.1
@@ -66,7 +66,6 @@ require (
6666
github.com/cloudspannerecosystem/spanner-change-streams-tail v0.3.1 // indirect
6767
github.com/cncf/udpa/go v0.0.0-20220112060539-c52dc94e7fbe // indirect
6868
github.com/cncf/xds/go v0.0.0-20231128003011-0fa0005c9caa // indirect
69-
github.com/creack/pty v1.1.18 // indirect
7069
github.com/creasty/defaults v1.7.0 // indirect
7170
github.com/dalzilio/rudd v1.1.1-0.20230806153452-9e08a6ea8170 // indirect
7271
github.com/danieljoos/wincred v1.2.1 // indirect
@@ -129,7 +128,6 @@ require (
129128
github.com/joho/godotenv v1.5.1 // indirect
130129
github.com/josharian/intern v1.0.0 // indirect
131130
github.com/json-iterator/go v1.1.12 // indirect
132-
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
133131
github.com/klauspost/compress v1.17.6 // indirect
134132
github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 // indirect
135133
github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 // indirect
@@ -139,7 +137,6 @@ require (
139137
github.com/mailru/easyjson v0.7.7 // indirect
140138
github.com/mattn/go-colorable v0.1.13 // indirect
141139
github.com/mattn/go-runewidth v0.0.15 // indirect
142-
github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d // indirect
143140
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
144141
github.com/mitchellh/mapstructure v1.5.0 // indirect
145142
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
@@ -205,7 +202,6 @@ require (
205202
google.golang.org/appengine v1.6.8 // indirect
206203
google.golang.org/genproto v0.0.0-20240205150955-31a09d347014 // indirect
207204
google.golang.org/genproto/googleapis/api v0.0.0-20240205150955-31a09d347014 // indirect
208-
google.golang.org/genproto/googleapis/rpc v0.0.0-20240205150955-31a09d347014 // indirect
209205
gopkg.in/inf.v0 v0.9.1 // indirect
210206
gopkg.in/ini.v1 v1.67.0 // indirect
211207
gopkg.in/yaml.v2 v2.4.0 // indirect

‎go.sum‎

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -53,8 +53,6 @@ github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4 h1:/vQbFIOMb
5353
github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4/go.mod h1:hN7oaIRCjzsZ2dE+yG5k+rsdt3qcwykqK6HVGcKwsw4=
5454
github.com/99designs/keyring v1.2.2 h1:pZd3neh/EmUzWONb35LxQfvuY7kiSXAq3HQd97+XBn0=
5555
github.com/99designs/keyring v1.2.2/go.mod h1:wes/FrByc8j7lFOAGLGSNEg8f/PaI3cgTBqhFkHUrPk=
56-
github.com/AlecAivazis/survey/v2 v2.3.7 h1:6I/u8FvytdGsgonrYsVn2t8t4QiRnh6QSTqkkhIiSjQ=
57-
github.com/AlecAivazis/survey/v2 v2.3.7/go.mod h1:xUTIdE4KCOIjsBAE1JYsUPoCqYdZ1reCfTwbto0Fduo=
5856
github.com/Azure/go-ansiterm v0.0.0-20170929234023-d6e3b3328b78 h1:w+iIsaOQNcT7OZ575w+acHgRric5iCyQh+xv+KJ4HB8=
5957
github.com/Azure/go-ansiterm v0.0.0-20170929234023-d6e3b3328b78/go.mod h1:LmzpDX56iTiv29bbRTIsUNlaFfuhWRQBWjQdVyAevI8=
6058
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
@@ -65,8 +63,6 @@ github.com/Masterminds/squirrel v1.5.4 h1:uUcX/aBc8O7Fg9kaISIUsHXdKuqehiXAMQTYX8
6563
github.com/Masterminds/squirrel v1.5.4/go.mod h1:NNaOrjSoIDfDA40n7sr2tPNZRfjzjA400rg+riTZj10=
6664
github.com/Microsoft/go-winio v0.6.0 h1:slsWYD/zyx7lCXoZVlvQrj0hPTM1HI4+v1sIda2yDvg=
6765
github.com/Microsoft/go-winio v0.6.0/go.mod h1:cTAf44im0RAYeL23bpB+fzCyDH2MJiz2BO69KH/soAE=
68-
github.com/Netflix/go-expect v0.0.0-20220104043353-73e0943537d2 h1:+vx7roKuyA63nhn5WAunQHLTznkw5W8b1Xc0dNjp83s=
69-
github.com/Netflix/go-expect v0.0.0-20220104043353-73e0943537d2/go.mod h1:HBCaDeC1lPdgDeDbhX8XFpy1jqjK0IBG8W5K+xYqA0w=
7066
github.com/Nvveen/Gotty v0.0.0-20120604004816-cd527374f1e5 h1:TngWCqHvy9oXAN6lEVMRuU21PR1EtLVZJmdB18Gu3Rw=
7167
github.com/Nvveen/Gotty v0.0.0-20120604004816-cd527374f1e5/go.mod h1:lmUJ/7eu/Q8D7ML55dXQrVaamCz2vxCfdQBasLZfHKk=
7268
github.com/TylerBrock/colorjson v0.0.0-20200706003622-8a50f05110d2 h1:ZBbLwSJqkHBuFDA6DUhhse0IGJ7T5bemHyNILUjvOq4=
@@ -135,9 +131,6 @@ github.com/containerd/continuity v0.3.0/go.mod h1:wJEAIwKOm/pBZuBd0JmeTvnLquTB1A
135131
github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
136132
github.com/cpuguy83/go-md2man/v2 v2.0.3/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
137133
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
138-
github.com/creack/pty v1.1.17/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
139-
github.com/creack/pty v1.1.18 h1:n56/Zwd5o6whRC5PMGretI4IdRLlmBXYNjScPaBgsbY=
140-
github.com/creack/pty v1.1.18/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=
141134
github.com/creasty/defaults v1.7.0 h1:eNdqZvc5B509z18lD8yc212CAqJNvfT1Jq6L8WowdBA=
142135
github.com/creasty/defaults v1.7.0/go.mod h1:iGzKe6pbEHnpMPtfDXZEr0NVxWnPTjb1bbDy08fPzYM=
143136
github.com/dalzilio/rudd v1.1.1-0.20230806153452-9e08a6ea8170 h1:bHEN1z3EOO/IXHTQ8ZcmGoW4gTJt+mSrH2Sd458uo0E=
@@ -354,8 +347,6 @@ github.com/hashicorp/golang-lru v1.0.2 h1:dV3g9Z/unq5DpblPpw+Oqcv4dU/1omnb4Ok8iP
354347
github.com/hashicorp/golang-lru v1.0.2/go.mod h1:iADmTwqILo4mZ8BN3D2Q6+9jd8WM5uGBxy+E8yxSoD4=
355348
github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
356349
github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
357-
github.com/hinshun/vt10x v0.0.0-20220119200601-820417d04eec h1:qv2VnGeEQHchGaZ/u7lxST/RaJw+cv273q79D81Xbog=
358-
github.com/hinshun/vt10x v0.0.0-20220119200601-820417d04eec/go.mod h1:Q48J4R4DvxnHolD5P8pOtXigYlRuPLGl6moFx3ulM68=
359350
github.com/hokaccha/go-prettyjson v0.0.0-20210113012101-fb4e108d2519 h1:nqAlWFEdqI0ClbTDrhDvE/8LeQ4pftrqKUX9w5k0j3s=
360351
github.com/hokaccha/go-prettyjson v0.0.0-20210113012101-fb4e108d2519/go.mod h1:pFlLw2CfqZiIBOx6BuCeRLCrfxBJipTY0nIOF/VbGcI=
361352
github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc=
@@ -397,8 +388,6 @@ github.com/jzelinskie/cobrautil/v2 v2.0.0-20231016191810-9f8a4f6d038a/go.mod h1:
397388
github.com/jzelinskie/stringz v0.0.3 h1:0GhG3lVMYrYtIvRbxvQI6zqRTT1P1xyQlpa0FhfUXas=
398389
github.com/jzelinskie/stringz v0.0.3/go.mod h1:hHYbgxJuNLRw91CmpuFsYEOyQqpDVFg8pvEh23vy4P0=
399390
github.com/k0kubun/go-ansi v0.0.0-20180517002512-3bf9e2903213/go.mod h1:vNUNkEQ1e29fT/6vq2aBdFsgNPmy8qMdSay1npru+Sw=
400-
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 h1:Z9n2FFNUXsshfwJMBgNA0RU6/i7WVaAegv3PtuIHPMs=
401-
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51/go.mod h1:CzGEWj7cYgsdH8dAjBGEr58BoE7ScuLd+fwFZ44+/x8=
402391
github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
403392
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
404393
github.com/klauspost/compress v1.17.6 h1:60eq2E/jlfwQXtvZEeBUYADs+BwKBWURIY+Gj2eRGjI=
@@ -426,10 +415,8 @@ github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0V
426415
github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0=
427416
github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
428417
github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
429-
github.com/mattn/go-colorable v0.1.2/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
430418
github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
431419
github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
432-
github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
433420
github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
434421
github.com/mattn/go-isatty v0.0.19/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
435422
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
@@ -441,9 +428,6 @@ github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh
441428
github.com/mattn/go-sqlite3 v1.14.6 h1:dNPt6NO46WmLVt2DLNpwczCmdV5boIZ6g/tlDrlRUbg=
442429
github.com/mattn/go-sqlite3 v1.14.6/go.mod h1:NyWgC/yNuGj7Q9rpYnZvas74GogHl5/Z4A/KQRfk6bU=
443430
github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
444-
github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE=
445-
github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d h1:5PJl274Y63IEHC+7izoQE9x6ikvDFZS2mDVS3drnohI=
446-
github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE=
447431
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db h1:62I3jR2EmQ4l5rM/4FEfDWcRD+abF5XlKShorW5LRoQ=
448432
github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db/go.mod h1:l0dey0ia/Uv7NcFFVbCLtqEBQbrT4OCwCSKTEv6enCw=
449433
github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
@@ -780,7 +764,6 @@ golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5h
780764
golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
781765
golang.org/x/sys v0.0.0-20181116152217-5ac8a444bdc5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
782766
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
783-
golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
784767
golang.org/x/sys v0.0.0-20190312061237-fead79001313/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
785768
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
786769
golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
@@ -839,7 +822,6 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
839822
golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
840823
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
841824
golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ=
842-
golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
843825
golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
844826
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
845827
golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=

‎internal/cmd/backup.go‎

Lines changed: 6 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,13 @@ import (
1919
"github.com/mattn/go-isatty"
2020
"github.com/rodaine/table"
2121
"github.com/rs/zerolog/log"
22-
"github.com/schollz/progressbar/v3"
2322
"github.com/spf13/cobra"
2423
"golang.org/x/exp/constraints"
2524
"golang.org/x/exp/maps"
2625

2726
"github.com/authzed/zed/internal/client"
2827
"github.com/authzed/zed/internal/commands"
28+
"github.com/authzed/zed/internal/console"
2929
"github.com/authzed/zed/pkg/backupformat"
3030
)
3131

@@ -231,30 +231,6 @@ func hasRelPrefix(rel *v1.Relationship, prefix string) bool {
231231
strings.HasPrefix(rel.Subject.Object.ObjectType, prefix)
232232
}
233233

234-
func relProgressBar(description string) *progressbar.ProgressBar {
235-
bar := progressbar.NewOptions(-1,
236-
progressbar.OptionSetWidth(10),
237-
progressbar.OptionSetRenderBlankState(true),
238-
progressbar.OptionSetVisibility(false),
239-
)
240-
if isatty.IsTerminal(os.Stderr.Fd()) {
241-
bar = progressbar.NewOptions64(-1,
242-
progressbar.OptionSetDescription(description),
243-
progressbar.OptionSetWriter(os.Stderr),
244-
progressbar.OptionSetWidth(10),
245-
progressbar.OptionThrottle(65*time.Millisecond),
246-
progressbar.OptionShowCount(),
247-
progressbar.OptionShowIts(),
248-
progressbar.OptionSetItsString("relationship"),
249-
progressbar.OptionOnCompletion(func() { _, _ = fmt.Fprint(os.Stderr, "\n") }),
250-
progressbar.OptionSpinnerType(14),
251-
progressbar.OptionFullWidth(),
252-
progressbar.OptionSetRenderBlankState(true),
253-
)
254-
}
255-
return bar
256-
}
257-
258234
func backupCreateCmdFunc(cmd *cobra.Command, args []string) (err error) {
259235
f, err := createBackupFile(args[0])
260236
if err != nil {
@@ -264,13 +240,13 @@ func backupCreateCmdFunc(cmd *cobra.Command, args []string) (err error) {
264240
defer func(e *error) { *e = errors.Join(*e, f.Close()) }(&err)
265241
defer func(e *error) { *e = errors.Join(*e, f.Sync()) }(&err)
266242

267-
client, err := client.NewClient(cmd)
243+
c, err := client.NewClient(cmd)
268244
if err != nil {
269245
return fmt.Errorf("unable to initialize client: %w", err)
270246
}
271247

272248
ctx := cmd.Context()
273-
schemaResp, err := client.ReadSchema(ctx, &v1.ReadSchemaRequest{})
249+
schemaResp, err := c.ReadSchema(ctx, &v1.ReadSchemaRequest{})
274250
if err != nil {
275251
return fmt.Errorf("error reading schema: %w", err)
276252
} else if schemaResp.ReadAt == nil {
@@ -299,7 +275,7 @@ func backupCreateCmdFunc(cmd *cobra.Command, args []string) (err error) {
299275
}
300276
defer func(e *error) { *e = errors.Join(*e, encoder.Close()) }(&err)
301277

302-
relationshipStream, err := client.BulkExportRelationships(ctx, &v1.BulkExportRelationshipsRequest{
278+
relationshipStream, err := c.BulkExportRelationships(ctx, &v1.BulkExportRelationshipsRequest{
303279
Consistency: &v1.Consistency{
304280
Requirement: &v1.Consistency_AtExactSnapshot{
305281
AtExactSnapshot: schemaResp.ReadAt,
@@ -312,7 +288,7 @@ func backupCreateCmdFunc(cmd *cobra.Command, args []string) (err error) {
312288

313289
relationshipReadStart := time.Now()
314290

315-
bar := relProgressBar("processing backup")
291+
bar := console.CreateProgressBar("processing backup")
316292
var relsEncoded, relsProcessed uint
317293
for {
318294
if err := ctx.Err(); err != nil {
@@ -520,7 +496,7 @@ func backupRedactCmdFunc(cmd *cobra.Command, args []string) error {
520496
}
521497

522498
defer func(e *error) { *e = errors.Join(*e, redactor.Close()) }(&err)
523-
bar := relProgressBar("redacting backup")
499+
bar := console.CreateProgressBar("redacting backup")
524500
var written int64
525501
for {
526502
if err := cmd.Context().Err(); err != nil {

‎internal/cmd/backup_test.go‎

Lines changed: 23 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,9 @@ import (
66
"path/filepath"
77
"strings"
88
"testing"
9-
"time"
109

1110
"github.com/authzed/zed/internal/client"
11+
zedtesting "github.com/authzed/zed/internal/testing"
1212

1313
v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
1414
"github.com/authzed/spicedb/pkg/tuple"
@@ -162,7 +162,7 @@ func TestBackupParseRelsCmdFunc(t *testing.T) {
162162
tt := tt
163163
t.Parallel()
164164

165-
cmd := createTestCobraCommandWithFlagValue(t, stringFlag{"prefix-filter", tt.filter})
165+
cmd := zedtesting.CreateTestCobraCommandWithFlagValue(t, zedtesting.StringFlag{FlagName: "prefix-filter", FlagValue: tt.filter})
166166
backupName := createTestBackup(t, tt.schema, tt.relationships)
167167
f, err := os.CreateTemp("", "parse-output")
168168
require.NoError(t, err)
@@ -183,7 +183,7 @@ func TestBackupParseRelsCmdFunc(t *testing.T) {
183183
}
184184

185185
func TestBackupParseRevisionCmdFunc(t *testing.T) {
186-
cmd := createTestCobraCommandWithFlagValue(t, stringFlag{"prefix-filter", "test"})
186+
cmd := zedtesting.CreateTestCobraCommandWithFlagValue(t, zedtesting.StringFlag{FlagName: "prefix-filter", FlagValue: "test"})
187187
backupName := createTestBackup(t, testSchema, testRelationships)
188188
f, err := os.CreateTemp("", "parse-output")
189189
require.NoError(t, err)
@@ -241,9 +241,9 @@ func TestBackupParseSchemaCmdFunc(t *testing.T) {
241241
tt := tt
242242
t.Parallel()
243243

244-
cmd := createTestCobraCommandWithFlagValue(t,
245-
stringFlag{"prefix-filter", tt.filter},
246-
boolFlag{"rewrite-legacy", tt.rewriteLegacy})
244+
cmd := zedtesting.CreateTestCobraCommandWithFlagValue(t,
245+
zedtesting.StringFlag{FlagName: "prefix-filter", FlagValue: tt.filter},
246+
zedtesting.BoolFlag{FlagName: "rewrite-legacy", FlagValue: tt.rewriteLegacy})
247247
backupName := createTestBackup(t, tt.schema, nil)
248248
f, err := os.CreateTemp("", "parse-output")
249249
require.NoError(t, err)
@@ -264,9 +264,9 @@ func TestBackupParseSchemaCmdFunc(t *testing.T) {
264264
}
265265

266266
func TestBackupCreateCmdFunc(t *testing.T) {
267-
cmd := createTestCobraCommandWithFlagValue(t,
268-
stringFlag{"prefix-filter", ""},
269-
boolFlag{"rewrite-legacy", false})
267+
cmd := zedtesting.CreateTestCobraCommandWithFlagValue(t,
268+
zedtesting.StringFlag{FlagName: "prefix-filter"},
269+
zedtesting.BoolFlag{FlagName: "rewrite-legacy"})
270270
f := filepath.Join(os.TempDir(), uuid.NewString())
271271
_, err := os.Stat(f)
272272
require.Error(t, err)
@@ -276,7 +276,7 @@ func TestBackupCreateCmdFunc(t *testing.T) {
276276

277277
ctx, cancel := context.WithCancel(context.Background())
278278
defer cancel()
279-
srv := newServer(ctx, t)
279+
srv := zedtesting.NewTestServer(ctx, t)
280280
go func() {
281281
require.NoError(t, srv.Run(ctx))
282282
}()
@@ -288,9 +288,9 @@ func TestBackupCreateCmdFunc(t *testing.T) {
288288
client.NewClient = originalClient
289289
}()
290290

291-
client.NewClient = clientFromConn(conn)
291+
client.NewClient = zedtesting.ClientFromConn(conn)
292292

293-
c, err := clientFromConn(conn)(cmd)
293+
c, err := zedtesting.ClientFromConn(conn)(cmd)
294294
require.NoError(t, err)
295295

296296
_, err = c.WriteSchema(ctx, &v1.WriteSchemaRequest{Schema: testSchema})
@@ -324,26 +324,21 @@ func TestBackupCreateCmdFunc(t *testing.T) {
324324
require.Equal(t, resp.WrittenAt.Token, d.ZedToken().Token)
325325
}
326326

327-
type durationFlag struct {
328-
flagName string
329-
flagValue time.Duration
330-
}
331-
332327
func TestBackupRestoreCmdFunc(t *testing.T) {
333-
cmd := createTestCobraCommandWithFlagValue(t,
334-
stringFlag{"prefix-filter", "test"},
335-
boolFlag{"rewrite-legacy", false},
336-
stringFlag{"conflict-strategy", "fail"},
337-
boolFlag{"disable-retries", false},
338-
intFlag{"batch-size", 100},
339-
uintFlag{"batches-per-transaction", 10},
340-
durationFlag{"request-timeout", 0},
328+
cmd := zedtesting.CreateTestCobraCommandWithFlagValue(t,
329+
zedtesting.StringFlag{FlagName: "prefix-filter", FlagValue: "test"},
330+
zedtesting.BoolFlag{FlagName: "rewrite-legacy"},
331+
zedtesting.StringFlag{FlagName: "conflict-strategy", FlagValue: "fail"},
332+
zedtesting.BoolFlag{FlagName: "disable-retries"},
333+
zedtesting.IntFlag{FlagName: "batch-size", FlagValue: 100},
334+
zedtesting.UintFlag{FlagName: "batches-per-transaction", FlagValue: 10},
335+
zedtesting.DurationFlag{FlagName: "request-timeout"},
341336
)
342337
backupName := createTestBackup(t, testSchema, testRelationships)
343338

344339
ctx, cancel := context.WithCancel(context.Background())
345340
defer cancel()
346-
srv := newServer(ctx, t)
341+
srv := zedtesting.NewTestServer(ctx, t)
347342
go func() {
348343
require.NoError(t, srv.Run(ctx))
349344
}()
@@ -355,9 +350,9 @@ func TestBackupRestoreCmdFunc(t *testing.T) {
355350
client.NewClient = originalClient
356351
}()
357352

358-
client.NewClient = clientFromConn(conn)
353+
client.NewClient = zedtesting.ClientFromConn(conn)
359354

360-
c, err := clientFromConn(conn)(cmd)
355+
c, err := zedtesting.ClientFromConn(conn)(cmd)
361356
require.NoError(t, err)
362357
err = backupRestoreCmdFunc(cmd, []string{backupName})
363358
require.NoError(t, err)

0 commit comments

Comments
 (0)