This commit is contained in:
Kristoffer Dalby 2025-07-14 20:56:01 +00:00
parent 67f2c20052
commit 45baead257
9 changed files with 25 additions and 221 deletions

View File

@ -1,82 +0,0 @@
# CLI Simplification - WithClient Pattern
## Problem
Every CLI command has repetitive gRPC client setup boilerplate:
```go
// This pattern appears 25+ times across all commands
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
// ... command logic ...
```
## Solution
Simple closure that handles client lifecycle:
```go
// client.go - 16 lines total
func WithClient(fn func(context.Context, v1.HeadscaleServiceClient) error) error {
ctx, client, conn, cancel := newHeadscaleCLIWithConfig()
defer cancel()
defer conn.Close()
return fn(ctx, client)
}
```
## Usage Example
### Before (users.go listUsersCmd):
```go
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
ctx, client, conn, cancel := newHeadscaleCLIWithConfig() // 4 lines
defer cancel()
defer conn.Close()
request := &v1.ListUsersRequest{}
// ... build request ...
response, err := client.ListUsers(ctx, request)
if err != nil {
ErrorOutput(err, "Cannot get users: "+status.Convert(err).Message(), output)
}
// ... handle response ...
}
```
### After:
```go
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
err := WithClient(func(ctx context.Context, client v1.HeadscaleServiceClient) error {
request := &v1.ListUsersRequest{}
// ... build request ...
response, err := client.ListUsers(ctx, request)
if err != nil {
ErrorOutput(err, "Cannot get users: "+status.Convert(err).Message(), output)
return err
}
// ... handle response ...
return nil
})
if err != nil {
return // Error already handled
}
}
```
## Benefits
- **Removes 4 lines of boilerplate** from every command
- **Ensures proper cleanup** - no forgetting defer statements
- **Simpler error handling** - return from closure, handled centrally
- **Easy to apply** - minimal changes to existing commands
## Rollout
This pattern can be applied to all 25+ commands systematically, removing ~100 lines of repetitive boilerplate.

View File

@ -15,10 +15,6 @@ import (
"google.golang.org/protobuf/types/known/timestamppb"
)
const (
// 90 days.
DefaultAPIKeyExpiry = "90d"
)
func init() {
rootCmd.AddCommand(apiKeysCmd)
@ -53,7 +49,7 @@ var listAPIKeys = &cobra.Command{
Short: "List the Api keys for headscale",
Aliases: []string{"ls", "show"},
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
output := GetOutputFlag(cmd)
err := WithClient(func(ctx context.Context, client v1.HeadscaleServiceClient) error {
request := &v1.ListApiKeysRequest{}
@ -118,7 +114,7 @@ and cannot be retrieved again.
If you loose a key, create a new one and revoke (expire) the old one.`,
Aliases: []string{"c", "new"},
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
output := GetOutputFlag(cmd)
request := &v1.CreateApiKeyRequest{}
@ -164,15 +160,10 @@ var expireAPIKeyCmd = &cobra.Command{
Short: "Expire an ApiKey",
Aliases: []string{"revoke", "exp", "e"},
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
output := GetOutputFlag(cmd)
prefix, err := cmd.Flags().GetString("prefix")
if err != nil {
ErrorOutput(
err,
fmt.Sprintf("Error getting prefix from CLI flag: %s", err),
output,
)
ErrorOutput(err, fmt.Sprintf("Error getting prefix from CLI flag: %s", err), output)
return
}
@ -206,15 +197,10 @@ var deleteAPIKeyCmd = &cobra.Command{
Short: "Delete an ApiKey",
Aliases: []string{"remove", "del"},
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
output := GetOutputFlag(cmd)
prefix, err := cmd.Flags().GetString("prefix")
if err != nil {
ErrorOutput(
err,
fmt.Sprintf("Error getting prefix from CLI flag: %s", err),
output,
)
ErrorOutput(err, fmt.Sprintf("Error getting prefix from CLI flag: %s", err), output)
return
}

View File

@ -1,105 +0,0 @@
#!/usr/bin/env python3
"""Script to convert all commands to use WithClient pattern"""
import re
import sys
import os
def convert_command(content):
"""Convert a single command to use WithClient pattern"""
# Pattern to match the gRPC client setup
pattern = r'(\t+)ctx, client, conn, cancel := newHeadscaleCLIWithConfig\(\)\n\t+defer cancel\(\)\n\t+defer conn\.Close\(\)\n\n'
# Find all occurrences
matches = list(re.finditer(pattern, content))
if not matches:
return content
# Process each match from the end to avoid offset issues
for match in reversed(matches):
indent = match.group(1)
start_pos = match.start()
end_pos = match.end()
# Find the end of the Run function
remaining_content = content[end_pos:]
# Find the matching closing brace for the Run function
brace_count = 0
func_end = -1
for i, char in enumerate(remaining_content):
if char == '{':
brace_count += 1
elif char == '}':
brace_count -= 1
if brace_count < 0: # Found the closing brace
func_end = i
break
if func_end == -1:
continue
# Extract the function body
func_body = remaining_content[:func_end]
# Indent the function body
indented_body = '\n'.join(indent + '\t' + line if line.strip() else line
for line in func_body.split('\n'))
# Create the new function with WithClient
new_func = f"""{indent}err := WithClient(func(ctx context.Context, client v1.HeadscaleServiceClient) error {{
{indented_body}
{indent}\treturn nil
{indent}}})
{indent}
{indent}if err != nil {{
{indent}\treturn
{indent}}}"""
# Replace the old pattern with the new one
content = content[:start_pos] + new_func + '\n' + content[end_pos + func_end:]
return content
def process_file(filepath):
"""Process a single Go file"""
try:
with open(filepath, 'r') as f:
content = f.read()
# Check if context is already imported
if 'import (' in content and '"context"' not in content:
# Add context import
content = content.replace(
'import (',
'import (\n\t"context"'
)
# Convert commands
new_content = convert_command(content)
# Write back if changed
if new_content != content:
with open(filepath, 'w') as f:
f.write(new_content)
print(f"Updated {filepath}")
else:
print(f"No changes needed for {filepath}")
except Exception as e:
print(f"Error processing {filepath}: {e}")
if __name__ == "__main__":
if len(sys.argv) != 2:
print("Usage: python3 convert_commands.py <go_file>")
sys.exit(1)
filepath = sys.argv[1]
if not os.path.exists(filepath):
print(f"File not found: {filepath}")
sys.exit(1)
process_file(filepath)

View File

@ -639,14 +639,14 @@ func nodesToPtables(
var lastSeenTime string
if node.GetLastSeen() != nil {
lastSeen = node.GetLastSeen().AsTime()
lastSeenTime = lastSeen.Format("2006-01-02 15:04:05")
lastSeenTime = lastSeen.Format(HeadscaleDateTimeFormat)
}
var expiry time.Time
var expiryTime string
if node.GetExpiry() != nil {
expiry = node.GetExpiry().AsTime()
expiryTime = expiry.Format("2006-01-02 15:04:05")
expiryTime = expiry.Format(HeadscaleDateTimeFormat)
} else {
expiryTime = "N/A"
}

View File

@ -15,9 +15,6 @@ import (
"google.golang.org/protobuf/types/known/timestamppb"
)
const (
DefaultPreAuthKeyExpiry = "1h"
)
func init() {
rootCmd.AddCommand(preauthkeysCmd)
@ -117,7 +114,7 @@ var listPreAuthKeys = &cobra.Command{
strconv.FormatBool(key.GetEphemeral()),
strconv.FormatBool(key.GetUsed()),
expiration,
key.GetCreatedAt().AsTime().Format("2006-01-02 15:04:05"),
key.GetCreatedAt().AsTime().Format(HeadscaleDateTimeFormat),
aclTags,
})

View File

@ -7,7 +7,7 @@ import (
)
func ColourTime(date time.Time) string {
dateStr := date.Format("2006-01-02 15:04:05")
dateStr := date.Format(HeadscaleDateTimeFormat)
if date.After(time.Now()) {
dateStr = pterm.LightGreen(dateStr)

View File

@ -10,6 +10,8 @@ import (
const (
deprecateNamespaceMessage = "use --user"
HeadscaleDateTimeFormat = "2006-01-02 15:04:05"
DefaultAPIKeyExpiry = "90d"
DefaultPreAuthKeyExpiry = "1h"
)
// FilterTableColumns filters table columns based on --columns flag

View File

@ -52,7 +52,7 @@ func init() {
userCmd.AddCommand(renameUserCmd)
usernameAndIDFlag(renameUserCmd)
renameUserCmd.Flags().StringP("new-name", "r", "", "New username")
renameNodeCmd.MarkFlagRequired("new-name")
renameUserCmd.MarkFlagRequired("new-name")
}
var errMissingParameter = errors.New("missing parameters")
@ -75,8 +75,7 @@ var createUserCmd = &cobra.Command{
return nil
},
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
output := GetOutputFlag(cmd)
userName := args[0]
request := &v1.CreateUserRequest{Name: userName}
@ -133,7 +132,7 @@ var destroyUserCmd = &cobra.Command{
Short: "Destroys a user",
Aliases: []string{"delete"},
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
output := GetOutputFlag(cmd)
id, username := usernameAndIDFromFlag(cmd)
request := &v1.ListUsersRequest{
@ -217,7 +216,7 @@ var listUsersCmd = &cobra.Command{
Short: "List all the users",
Aliases: []string{"ls", "show"},
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
output := GetOutputFlag(cmd)
err := WithClient(func(ctx context.Context, client v1.HeadscaleServiceClient) error {
request := &v1.ListUsersRequest{}
@ -260,7 +259,7 @@ var listUsersCmd = &cobra.Command{
user.GetDisplayName(),
user.GetName(),
user.GetEmail(),
user.GetCreatedAt().AsTime().Format("2006-01-02 15:04:05"),
user.GetCreatedAt().AsTime().Format(HeadscaleDateTimeFormat),
},
)
}
@ -289,7 +288,7 @@ var renameUserCmd = &cobra.Command{
Short: "Renames a user",
Aliases: []string{"mv"},
Run: func(cmd *cobra.Command, args []string) {
output, _ := cmd.Flags().GetString("output")
output := GetOutputFlag(cmd)
id, username := usernameAndIDFromFlag(cmd)
newName, _ := cmd.Flags().GetString("new-name")

View File

@ -12,6 +12,7 @@ import (
"github.com/juanfont/headscale/hscontrol/types"
"github.com/juanfont/headscale/hscontrol/util"
"github.com/rs/zerolog/log"
"github.com/spf13/cobra"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/credentials/insecure"
@ -199,3 +200,9 @@ func (t tokenAuth) GetRequestMetadata(
func (tokenAuth) RequireTransportSecurity() bool {
return true
}
// GetOutputFlag returns the output flag value (never fails)
func GetOutputFlag(cmd *cobra.Command) string {
output, _ := cmd.Flags().GetString("output")
return output
}