Skip to content

Commit d91328f

Browse files
RoseSecurityRoseSecurity
authored
Merge pull request #100 from RoseSecurity/add-timeouts
DEV-19 Incorporate `context.Context` for Commands
2 parents c53cfd7 + 8945302 commit d91328f

File tree

3 files changed

+54
-27
lines changed

3 files changed

+54
-27
lines changed

‎cmd/run.go

Lines changed: 46 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@
44
package cmd
55

66
import (
7+
"context"
78
"fmt"
89
"os"
910
"os/exec"
11+
"time"
1012

1113
"github.com/RoseSecurity/terramaid/internal"
1214
"github.com/RoseSecurity/terramaid/pkg/utils"
@@ -15,14 +17,15 @@ import (
1517
)
1618

1719
type options struct {
18-
WorkingDir string `env:"WORKING_DIR" envDefault:"."`
19-
TFPlan string `env:"TF_PLAN"`
20-
TFBinary string `env:"TF_BINARY"`
21-
Output string `env:"OUTPUT" envDefault:"Terramaid.md"`
22-
Direction string `env:"DIRECTION" envDefault:"TD"`
23-
SubgraphName string `env:"SUBGRAPH_NAME" envDefault:"Terraform"`
24-
ChartType string `env:"CHART_TYPE" envDefault:"flowchart"`
25-
Verbose bool `env:"VERBOSE" envDefault:"false"`
20+
WorkingDir string `env:"WORKING_DIR" envDefault:"."`
21+
TFPlan string `env:"TF_PLAN"`
22+
TFBinary string `env:"TF_BINARY"`
23+
Output string `env:"OUTPUT" envDefault:"Terramaid.md"`
24+
Direction string `env:"DIRECTION" envDefault:"TD"`
25+
SubgraphName string `env:"SUBGRAPH_NAME" envDefault:"Terraform"`
26+
ChartType string `env:"CHART_TYPE" envDefault:"flowchart"`
27+
Verbose bool `env:"VERBOSE" envDefault:"false"`
28+
Timeout time.Duration `env:"TIMEOUT" envDefault:"0"`
2629
}
2730

2831
var opts options // Global variable for flags and env variables
@@ -33,12 +36,19 @@ var runCmd = &cobra.Command{
3336
SilenceUsage: true,
3437
SilenceErrors: true,
3538
RunE: func(cmd *cobra.Command, args []string) error {
36-
// The opts variable is automatically populated with flags here
37-
return generateDiagrams(&opts)
39+
ctx := cmd.Context()
40+
41+
if opts.Timeout > 0 {
42+
var cancel context.CancelFunc
43+
ctx, cancel = context.WithTimeout(ctx, opts.Timeout)
44+
defer cancel()
45+
}
46+
47+
return generateDiagrams(ctx, &opts)
3848
},
3949
}
4050

41-
func generateDiagrams(opts *options) error {
51+
func generateDiagrams(ctx context.Context, opts *options) error {
4252
if opts.Verbose {
4353
utils.LogVerbose("Starting Terramaid with the following options:")
4454
utils.LogVerbose("- Working Directory: %s", opts.WorkingDir)
@@ -48,6 +58,16 @@ func generateDiagrams(opts *options) error {
4858
utils.LogVerbose("- Direction: %s", opts.Direction)
4959
utils.LogVerbose("- Subgraph Name: %s", opts.SubgraphName)
5060
utils.LogVerbose("- Chart Type: %s", opts.ChartType)
61+
if opts.Timeout > 0 {
62+
utils.LogVerbose("- Timeout: %s", opts.Timeout)
63+
}
64+
}
65+
66+
// Early cancellation check
67+
select {
68+
case <-ctx.Done():
69+
return ctx.Err()
70+
default:
5171
}
5272

5373
if opts.WorkingDir != "" {
@@ -56,7 +76,7 @@ func generateDiagrams(opts *options) error {
5676
return fmt.Errorf("error checking Terraform files in directory \"%s\": %v", opts.WorkingDir, err)
5777
}
5878
if !exists {
59-
return fmt.Errorf("Terraform files do not exist in directory \"%s\"", opts.WorkingDir)
79+
return fmt.Errorf("terraform files do not exist in directory \"%s\"", opts.WorkingDir)
6080
}
6181
if opts.Verbose {
6282
utils.LogVerbose("Confirmed Terraform files exist in %s", opts.WorkingDir)
@@ -83,36 +103,43 @@ func generateDiagrams(opts *options) error {
83103
// Spinner initialization and graph parsing
84104
sp := utils.NewSpinner("Generating Terramaid Diagrams")
85105
sp.Start()
106+
defer sp.Stop()
86107

87108
if opts.Verbose {
88109
utils.LogVerbose("Initializing Terraform and building graph...")
89110
}
90-
graph, err := internal.ParseTerraform(opts.WorkingDir, opts.TFBinary, opts.TFPlan, opts.Verbose)
111+
112+
graph, err := internal.ParseTerraform(ctx, opts.WorkingDir, opts.TFBinary, opts.TFPlan, opts.Verbose)
91113
if err != nil {
92-
sp.Stop()
93114
return fmt.Errorf("error parsing Terraform: %w", err)
94115
}
95116

117+
// Respect context cancellation after heavy parsing
118+
if err := ctx.Err(); err != nil {
119+
return err
120+
}
121+
96122
// Generate the Mermaid diagram
97123
if opts.Verbose {
98124
utils.LogVerbose("Generating Mermaid flowchart...")
99125
}
100-
mermaidDiagram, err := internal.GenerateMermaidFlowchart(graph, opts.Direction, opts.SubgraphName, opts.Verbose)
126+
mermaidDiagram, err := internal.GenerateMermaidFlowchart(ctx, graph, opts.Direction, opts.SubgraphName, opts.Verbose)
101127
if err != nil {
102-
sp.Stop()
103128
return fmt.Errorf("error generating Mermaid diagram: %w", err)
104129
}
105130

131+
if err := ctx.Err(); err != nil {
132+
return err
133+
}
134+
106135
// Write the Mermaid diagram to the specified output file
107136
if opts.Verbose {
108137
utils.LogVerbose("Writing Mermaid diagram to %s", opts.Output)
109138
}
110139
if err := os.WriteFile(opts.Output, []byte(mermaidDiagram), 0o644); err != nil {
111-
sp.Stop()
112140
return fmt.Errorf("error writing to file: %w", err)
113141
}
114142

115-
sp.Stop()
116143
fmt.Printf("Mermaid diagram successfully written to %s\n", opts.Output)
117144

118145
return nil
@@ -133,6 +160,7 @@ func init() {
133160
runCmd.Flags().StringVarP(&opts.TFBinary, "tf-binary", "b", opts.TFBinary, "Path to Terraform binary (env: TERRAMAID_TF_BINARY)")
134161
runCmd.Flags().StringVarP(&opts.WorkingDir, "working-dir", "w", opts.WorkingDir, "Working directory for Terraform (env: TERRAMAID_WORKING_DIR)")
135162
runCmd.Flags().BoolVarP(&opts.Verbose, "verbose", "v", opts.Verbose, "Enable verbose output (env: TERRAMAID_VERBOSE)")
163+
runCmd.Flags().DurationVarP(&opts.Timeout, "timeout", "t", opts.Timeout, "Timeout for the entire run (e.g. 5m) (env: TERRAMAID_TIMEOUT)")
136164

137165
// Disable auto-generated string from documentation so that documentation is cleanly built and updated
138166
runCmd.DisableAutoGenTag = true

‎internal/flowchart.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package internal
55

66
import (
7+
"context"
78
"fmt"
89
"regexp"
910
"strings"
@@ -40,7 +41,7 @@ func CleanLabel(label string) string {
4041
}
4142

4243
// GenerateMermaidFlowchart generates a Mermaid diagram from a gographviz graph
43-
func GenerateMermaidFlowchart(graph *gographviz.Graph, direction string, subgraphName string, verbose bool) (string, error) {
44+
func GenerateMermaidFlowchart(ctx context.Context, graph *gographviz.Graph, direction string, subgraphName string, verbose bool) (string, error) {
4445
validDirections := map[string]bool{"TB": true, "TD": true, "BT": true, "RL": true, "LR": true}
4546
if !validDirections[direction] {
4647
return "", fmt.Errorf("invalid direction %s: valid options are TB, TD, BT, RL, LR", direction)
@@ -135,12 +136,12 @@ func GenerateMermaidFlowchart(graph *gographviz.Graph, direction string, subgrap
135136
}
136137

137138
sb.WriteString("```\n")
138-
139+
139140
if verbose {
140141
nodeCount := len(addedNodes)
141142
edgeCount := len(graph.Edges.Edges)
142143
utils.LogVerbose("Mermaid diagram generation complete with %d nodes and %d edges", nodeCount, edgeCount)
143144
}
144-
145+
145146
return sb.String(), nil
146147
}

‎internal/parse.go

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,12 @@ const emptyGraph = `digraph G {
2121
`
2222

2323
// ParseTerraform parses the Terraform plan and returns the generated graph
24-
func ParseTerraform(workingDir, tfPath, planFile string, verbose bool) (*gographviz.Graph, error) {
25-
ctx := context.Background()
26-
24+
func ParseTerraform(ctx context.Context, workingDir, tfPath, planFile string, verbose bool) (*gographviz.Graph, error) {
2725
if verbose {
2826
utils.LogVerbose("Initializing Terraform with working directory: %s", workingDir)
2927
utils.LogVerbose("Using Terraform binary: %s", tfPath)
3028
}
31-
29+
3230
tf, err := tfexec.NewTerraform(workingDir, tfPath)
3331
if err != nil {
3432
return nil, err
@@ -37,7 +35,7 @@ func ParseTerraform(workingDir, tfPath, planFile string, verbose bool) (*gograph
3735
if verbose {
3836
utils.LogVerbose("Running terraform init with upgrade=true")
3937
}
40-
38+
4139
if err := tf.Init(ctx, tfexec.Upgrade(true)); err != nil {
4240
return nil, err
4341
}
@@ -56,7 +54,7 @@ func ParseTerraform(workingDir, tfPath, planFile string, verbose bool) (*gograph
5654
if verbose {
5755
utils.LogVerbose("Running terraform graph command")
5856
}
59-
57+
6058
output, err := tf.Graph(ctx, opts)
6159
if err != nil {
6260
return nil, err

0 commit comments

Comments
 (0)