fs go plugin

pkg/lang/golang/arguments.go

@@ -0,0 +1,49 @@
+package golang
+
+import (
+	"fmt"
+
+	sitter "github.com/smacker/go-tree-sitter"
+)
+
+type Argument struct {
+	Content string
+	Type    string
+}
+
+// GetArguements is passed a tree-sitter node, which is of type argument_list, and returns a list of in order Arguments
+func GetArguements(args *sitter.Node) []Argument {
+
+	arguments := []Argument{}
+	nextMatch := doQuery(args, findArgs)
+	for {
+		match, found := nextMatch()
+		if !found {
+			break
+		}
+
+		arg := match["arg"]
+		if arg == nil {
+			continue
+		}
+
+		arguments = append(arguments, Argument{Content: arg.Content(), Type: arg.Type()})
+	}
+	return arguments
+}
+
+func ArgumentListToString(args []Argument) string {
+	result := "("
+	for index, arg := range args {
+		if index < len(args)-1 {
+			result += fmt.Sprintf("%s, ", arg.Content)
+		} else {
+			result += arg.Content + ")"
+		}
+	}
+	return result
+}
+
+func (a *Argument) IsString() bool {
+	return a.Type == "interpreted_string_literal"
+}

pkg/lang/golang/aws_runtime/aws.go

@@ -57,3 +57,10 @@ func (r *AwsRuntime) AddExecRuntimeFiles(unit *core.ExecutionUnit, result *core.
 
 	return nil
 }
+
+func (r *AwsRuntime) GetFsImports() []string {
+	return []string{
+		"gocloud.dev/blob",
+		"_ \"gocloud.dev/blob/s3blob\"",
+	}
+}

pkg/lang/golang/imports.go

@@ -47,6 +47,16 @@ func GetImportsInFile(f *core.SourceFile) []Import {
 	return imports
 }
 
+func GetNamedImportInFile(f *core.SourceFile, namedImport string) Import {
+	imports := GetImportsInFile(f)
+	for _, i := range imports {
+		if i.Package == namedImport {
+			return i
+		}
+	}
+	return Import{}
+}
+
 func UpdateImportsInFile(f *core.SourceFile, importsToAdd []Import, importsToRemove []Import) error {
 
 	imports := GetImportsInFile(f)

pkg/lang/golang/imports_test.go

@@ -85,6 +85,52 @@ func Test_GetImportsInFile(t *testing.T) {
 	}
 }
 
+func Test_GetNamedImportInFile(t *testing.T) {
+	tests := []struct {
+		name        string
+		source      string
+		importToGet string
+		want        Import
+	}{
+		{
+			name: "simple import",
+			source: `import(
+				"os"
+)`,
+			want:        Import{Package: "os"},
+			importToGet: "os",
+		},
+		{
+			name: "no match",
+			source: `import(
+				"os"
+				"github.com/go-chi/chi/v5"
+)`,
+			importToGet: "NotReal",
+		},
+		{
+			name: "multiple imports with an alias",
+			source: `import(
+				"os"
+				chi "github.com/go-chi/chi/v5"
+)`,
+			importToGet: "github.com/go-chi/chi/v5",
+			want:        Import{Package: "github.com/go-chi/chi/v5", Alias: "chi"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			assert := assert.New(t)
+			f, err := core.NewSourceFile("test.go", strings.NewReader(tt.source), Language)
+			if !assert.NoError(err) {
+				return
+			}
+			i := GetNamedImportInFile(f, tt.importToGet)
+			assert.Equal(tt.want, i)
+		})
+	}
+}
+
 func Test_UpdateImportsInFile(t *testing.T) {
 	tests := []struct {
 		name            string

pkg/lang/golang/plugin_fs.go

@@ -0,0 +1,159 @@
+package golang
+
+import (
+	"errors"
+	"fmt"
+	"strings"
+
+	"github.com/klothoplatform/klotho/pkg/annotation"
+	"github.com/klothoplatform/klotho/pkg/core"
+	"github.com/klothoplatform/klotho/pkg/logging"
+	"github.com/klothoplatform/klotho/pkg/multierr"
+	"github.com/klothoplatform/klotho/pkg/query"
+	sitter "github.com/smacker/go-tree-sitter"
+	"go.uber.org/zap"
+)
+
+type PersistFsPlugin struct {
+	runtime Runtime
+}
+
+func (p PersistFsPlugin) Name() string { return "Persist" }
+
+func (p PersistFsPlugin) Transform(result *core.CompilationResult, deps *core.Dependencies) error {
+
+	var errs multierr.Error
+	for _, res := range result.Resources() {
+		unit, ok := res.(*core.ExecutionUnit)
+		if !ok {
+			continue
+		}
+		for _, f := range unit.Files() {
+			goSource, ok := Language.ID.CastFile(f)
+			if !ok {
+				continue
+			}
+
+			resources, err := p.handleFile(goSource, unit)
+			if err != nil {
+				errs.Append(core.WrapErrf(err, "failed to handle persist in unit %s", unit.Name))
+				continue
+			}
+
+			for _, r := range resources {
+				result.Add(r)
+
+				deps.Add(core.ResourceKey{
+					Name: unit.Name,
+					Kind: core.ExecutionUnitKind,
+				}, r.Key())
+			}
+		}
+	}
+
+	return errs.ErrOrNil()
+}
+
+func (p *PersistFsPlugin) handleFile(f *core.SourceFile, unit *core.ExecutionUnit) ([]core.CloudResource, error) {
+	resources := []core.CloudResource{}
+	var errs multierr.Error
+	annots := f.Annotations()
+	for _, annot := range annots {
+		cap := annot.Capability
+		if cap.Name != annotation.PersistCapability {
+			continue
+		}
+		fsResult := queryFS(f, annot)
+		if fsResult != nil {
+			persistResource, err := p.transformFS(f, annot, fsResult, unit)
+			if err != nil {
+				errs.Append(err)
+			}
+			resources = append(resources, persistResource)
+
+		}
+	}
+	return resources, errs.ErrOrNil()
+}
+
+func (p *PersistFsPlugin) transformFS(f *core.SourceFile, cap *core.Annotation, result *persistResult, unit *core.ExecutionUnit) (core.CloudResource, error) {
+
+	fsEnvVar := core.EnvironmentVariable{
+		Name:       cap.Capability.ID + "_fs_bucket",
+		Kind:       string(core.PersistFileKind),
+		ResourceID: cap.Capability.ID,
+		Value:      "bucket_url",
+	}
+
+	unit.EnvironmentVariables = append(unit.EnvironmentVariables, fsEnvVar)
+
+	args := GetArguements(result.args)
+
+	// We need to check to make sure the path supplied to the original node content is a static string. This is because it will get erased and we dont want to leave os level orphaned code
+	if !args[0].IsString() {
+		return nil, errors.New("must supply static string for secret path")
+	}
+
+	args[0].Content = "nil"
+	args[1].Content = fmt.Sprintf(`os.Getenv("%s")`, fsEnvVar.Name)
+
+	newNodeContent := strings.Replace(result.args.Content(), result.args.Content(), ArgumentListToString(args), 1)
+
+	f.ReplaceNodeContent(result.args, newNodeContent)
+	f.ReplaceNodeContent(result.operator, "blob")
+
+	err := UpdateImportsInFile(f, p.runtime.GetFsImports(), []Import{{Package: "gocloud.dev/blob/fileblob"}})
+	if err != nil {
+		return nil, err
+	}
+
+	persist := &core.Persist{
+		Kind: core.PersistFileKind,
+		Name: cap.Capability.ID,
+	}
+	return persist, nil
+}
+
+type persistResult struct {
+	varName  string
+	operator *sitter.Node
+	args     *sitter.Node
+}
+
+func queryFS(file *core.SourceFile, annotation *core.Annotation) *persistResult {
+	log := zap.L().With(logging.FileField(file), logging.AnnotationField(annotation))
+
+	fileBlobImport := GetNamedImportInFile(file, "gocloud.dev/blob/fileblob")
+
+	nextMatch := doQuery(annotation.Node, fileBucket)
+
+	match, found := nextMatch()
+	if !found {
+		return nil
+	}
+
+	varName, args, id := match["varName"], match["args"], match["id"]
+
+	if id != nil {
+		if fileBlobImport.Alias != "" {
+			if !query.NodeContentEquals(id, fileBlobImport.Alias) {
+				return nil
+			}
+		} else {
+			if !query.NodeContentEquals(id, "fileblob") {
+				return nil
+			}
+		}
+	}
+
+	if _, found := nextMatch(); found {
+		log.Warn("too many assignments for fs_storage")
+		return nil
+	}
+
+	return &persistResult{
+		varName:  varName.Content(),
+		operator: id,
+		args:     args,
+	}
+}