update with new imports code and find more syntax for how the bucket …

…can be set
klothoplatform · Feb 10, 2023 · 91f3dd9 · 91f3dd9
1 parent 21909d4
commit 91f3dd9
Show file tree

Hide file tree

Showing 10 changed files with 158 additions and 26 deletions.
diff --git a/pkg/lang/golang/arguments.go b/pkg/lang/golang/arguments.go
@@ -45,5 +45,5 @@ func ArgumentListToString(args []Argument) string {
 }
 
 func (a *Argument) IsString() bool {
-	return a.Type == "interpreted_string_literal"
+	return a.Type == "interpreted_string_literal" || a.Type == "raw_string_literal"
 }
diff --git a/pkg/lang/golang/aws_runtime/aws.go b/pkg/lang/golang/aws_runtime/aws.go
@@ -58,9 +58,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\"",
+func (r *AwsRuntime) GetFsImports() []golang.Import {
+	return []golang.Import{
+		{Package: "os"},
+		{Package: "gocloud.dev/blob"},
+		{Package: "gocloud.dev/blob/s3blob", Alias: "_"},
 	}
 }
diff --git a/pkg/lang/golang/imports.go b/pkg/lang/golang/imports.go
@@ -127,7 +127,6 @@ func DeleteImportNodes(f *core.SourceFile) error {
 		if !found {
 			break
 		}
-		fmt.Println(string(f.Program()))
 		err := f.ReplaceNodeContent(match["expression"], "")
 		if err != nil {
 			return err

diff --git a/pkg/lang/golang/imports_test.go b/pkg/lang/golang/imports_test.go
@@ -1,7 +1,6 @@
 package golang
 
 import (
-	"fmt"
 	"strings"
 	"testing"
 
@@ -126,7 +125,8 @@ func Test_GetNamedImportInFile(t *testing.T) {
 				return
 			}
 			i := GetNamedImportInFile(f, tt.importToGet)
-			assert.Equal(tt.want, i)
+			assert.Equal(tt.want.Package, i.Package)
+			assert.Equal(tt.want.Alias, i.Alias)
 		})
 	}
 }
@@ -242,7 +242,6 @@ import (
 			if !assert.NoError(err) {
 				return
 			}
-			fmt.Println(string(f.Program()))
 			assert.Equal(tt.want, string(f.Program()))
 		})
 	}

diff --git a/pkg/lang/golang/plugin_fs.go b/pkg/lang/golang/plugin_fs.go
@@ -28,12 +28,7 @@ func (p PersistFsPlugin) Transform(result *core.CompilationResult, deps *core.De
 		if !ok {
 			continue
 		}
-		for _, f := range unit.Files() {
-			goSource, ok := Language.ID.CastFile(f)
-			if !ok {
-				continue
-			}
-
+		for _, goSource := range unit.FilesOfLang(goLang) {
 			resources, err := p.handleFile(goSource, unit)
 			if err != nil {
 				errs.Append(core.WrapErrf(err, "failed to handle persist in unit %s", unit.Name))

diff --git a/pkg/lang/golang/plugin_fs_test.go b/pkg/lang/golang/plugin_fs_test.go
@@ -31,6 +31,39 @@ bucket, err := fileblob.OpenBucket("myDir", nil)`,
 				varName: "bucket",
 			},
 		},
+		{
+			name: "simple var file blob",
+			source: `
+import (
+	"gocloud.dev/blob/fileblob"
+)
+/**
+* @klotho::persist {
+*	id = "test"
+* }
+*/
+var bucket, err = fileblob.OpenBucket("myDir", nil)`,
+			want: &persistResult{
+				varName: "bucket",
+			},
+		},
+		{
+			name: "simple var file blob",
+			source: `
+import (
+	"gocloud.dev/blob/fileblob"
+)
+/**
+* @klotho::persist {
+*	id = "test"
+* }
+*/
+var bucket, err
+bucket, err = fileblob.OpenBucket("myDir", nil)`,
+			want: &persistResult{
+				varName: "bucket",
+			},
+		},
 		{
 			name: "aliased file blob",
 			source: `
@@ -113,7 +146,8 @@ func Test_Transform(t *testing.T) {
 	}{
 		{
 			name: "simple file blob",
-			source: `import (
+			source: `package fs
+import (
 	"gocloud.dev/blob/fileblob"
 )
 /**
@@ -128,23 +162,95 @@ bucket, err := fileblob.OpenBucket("myDir", nil)
 					Kind: core.PersistFileKind,
 					Name: "test",
 				},
-				content: `
+				content: `package fs
+
 import (
 	"gocloud.dev/blob"
-	_  "gocloud.dev/blob/s3blob"
+	_ "gocloud.dev/blob/s3blob"
 )
+
 /**
 * @klotho::persist {
 *	id = "test"
 * }
 */
 bucket, err := blob.OpenBucket(nil, os.Getenv("test_fs_bucket"))
+`,
+			},
+		},
+		{
+			name: "long var file blob",
+			source: `package fs
+import (
+	"gocloud.dev/blob/fileblob"
+)
+/**
+* @klotho::persist {
+*	id = "test"
+* }
+*/
+var bucket, err = fileblob.OpenBucket("myDir", nil)
+`,
+			want: testResult{
+				resource: core.Persist{
+					Kind: core.PersistFileKind,
+					Name: "test",
+				},
+				content: `package fs
+
+import (
+	"gocloud.dev/blob"
+	_ "gocloud.dev/blob/s3blob"
+)
+
+/**
+* @klotho::persist {
+*	id = "test"
+* }
+*/
+var bucket, err = blob.OpenBucket(nil, os.Getenv("test_fs_bucket"))
+`,
+			},
+		},
+		{
+			name: "var deckaration file blob",
+			source: `package fs
+import (
+	"gocloud.dev/blob/fileblob"
+)
+/**
+* @klotho::persist {
+*	id = "test"
+* }
+*/
+var bucket, err
+bucket, err = fileblob.OpenBucket("myDir", nil)
+`,
+			want: testResult{
+				resource: core.Persist{
+					Kind: core.PersistFileKind,
+					Name: "test",
+				},
+				content: `package fs
+
+import (
+	"gocloud.dev/blob"
+	_ "gocloud.dev/blob/s3blob"
+)
+
+/**
+* @klotho::persist {
+*	id = "test"
+* }
+*/
+var bucket, err
+bucket, err = blob.OpenBucket(nil, os.Getenv("test_fs_bucket"))
 `,
 			},
 		},
 		{
 			name: "non string as path throws err",
-			source: `
+			source: `package fs
 		import (
 			alias "gocloud.dev/blob/fileblob"
 		)

diff --git a/pkg/lang/golang/queries/find_args.scm b/pkg/lang/golang/queries/find_args.scm
@@ -1,4 +1,4 @@
 ; Finds arguments in an arguments_list
-(argument_list [
+(argument_list
   (_) @arg
-])
+)
diff --git a/pkg/lang/golang/queries/gocloud/file_bucket.scm b/pkg/lang/golang/queries/gocloud/file_bucket.scm
@@ -1,4 +1,5 @@
-(short_var_declaration
+[
+  (short_var_declaration
    left: (expression_list
       	(identifier) @varName
        	(identifier)
@@ -13,6 +14,37 @@
       (#match? @method "OpenBucket")
       )
   )
-)@expression
+)@expression ;; bucket, err := fileblob.OpenBucket(myDir, myConnector)
+(assignment_statement
+   left: (expression_list
+      	(identifier) @varName
+       	(identifier)
+        ) @variables
+  right: (expression_list
+    (call_expression
+      function: (selector_expression
+          operand: (identifier) @id
+          field: (field_identifier) @method
+          )
+      arguments: (argument_list) @args
+      (#match? @method "OpenBucket")
+      )
+  )
+)@expression ;; bucket, err = fileblob.OpenBucket(myDir, myConnector)
+(var_declaration
+  	(var_spec
+     name: (identifier) @varName
+	   value: (expression_list
+        (call_expression
+          function: (selector_expression
+            operand: (identifier) @id
+            field: (field_identifier) @method
+          )
+          arguments: (argument_list) @args
+          (#match? @method "OpenBucket")
+        )
+      )
+    )
+)@expression ;; var bucket, err = fileblob.OpenBucket(myDir, myConnector)
+]
 
-;; bucket, err := fileblob.OpenBucket(myDir, myConnector)
diff --git a/pkg/lang/golang/runtime.go b/pkg/lang/golang/runtime.go
@@ -13,7 +13,7 @@ import (
 type (
 	Runtime interface {
 		AddExecRuntimeFiles(unit *core.ExecutionUnit, result *core.CompilationResult, deps *core.Dependencies) error
-		GetFsImports() []string
+		GetFsImports() []Import
 	}
 )
 

diff --git a/pkg/lang/golang/runtime_test.go b/pkg/lang/golang/runtime_test.go
@@ -12,6 +12,6 @@ func (n NoopRuntime) AddExecRuntimeFiles(unit *core.ExecutionUnit, result *core.
 func (n NoopRuntime) GetFsImports() []Import {
 	return []Import{
 		{Package: "gocloud.dev/blob"},
-		{Alias: "_", Package: "gocloud.dev/blob/s3blob\""},
+		{Alias: "_", Package: "gocloud.dev/blob/s3blob"},
 	}
 }