Merge pull request #70 from sev-2/feature/support-probe-controller

feature(state) support endpoint for import state

middleware_test.go

@@ -88,6 +88,9 @@ func Test_Tracer(t *testing.T) {
 		RequestContextFn: func() *fasthttp.RequestCtx {
 			return &fasthttp.RequestCtx{}
 		},
+		SendJsonFn: func(data any) error {
+			return nil
+		},
 	}
 
 	breakerFn := breakerMiddleware(fn)

pkg/controllers/state_controller.go

@@ -0,0 +1,31 @@
+package controllers
+
+import (
+	"net/http"
+
+	"github.com/sev-2/raiden"
+	"github.com/sev-2/raiden/pkg/resource"
+)
+
+type StateReadyRequest struct {
+}
+
+type StateReadyResponse struct {
+	Message string `json:"message"`
+}
+
+type StateReadyController struct {
+	raiden.ControllerBase
+	Http    string `path:"/state/ready" type:"custom"`
+	Payload *StateReadyRequest
+	Result  StateReadyResponse
+}
+
+func (c *StateReadyController) Post(ctx raiden.Context) error {
+	if err := resource.Import(&resource.Flags{UpdateStateOnly: true}, ctx.Config()); err != nil {
+		return ctx.SendErrorWithCode(http.StatusInternalServerError, err)
+	}
+
+	c.Result.Message = "success initiate local state"
+	return ctx.SendJson(c.Result)
+}

pkg/controllers/state_controller_test.go

@@ -0,0 +1,100 @@
+package controllers_test
+
+import (
+	"os"
+	"testing"
+
+	"github.com/sev-2/raiden"
+	"github.com/sev-2/raiden/pkg/controllers"
+	"github.com/sev-2/raiden/pkg/logger"
+	"github.com/sev-2/raiden/pkg/mock"
+	"github.com/sev-2/raiden/pkg/resource"
+	"github.com/sev-2/raiden/pkg/state"
+	"github.com/sev-2/raiden/pkg/supabase/objects"
+	"github.com/stretchr/testify/assert"
+	"github.com/valyala/fasthttp"
+)
+
+func loadConfig() *raiden.Config {
+	return &raiden.Config{
+		DeploymentTarget:    raiden.DeploymentTargetCloud,
+		ProjectId:           "test-project-id",
+		ProjectName:         "test-project",
+		SupabaseApiBasePath: "/v1",
+		SupabaseApiUrl:      "http://supabase.cloud.com",
+		SupabasePublicUrl:   "http://supabase.cloud.com",
+	}
+}
+
+func TestStateController_Post(t *testing.T) {
+	flags := &resource.Flags{
+		ProjectPath:     "test_project",
+		DryRun:          false,
+		UpdateStateOnly: true,
+	}
+	config := loadConfig()
+	resource.ImportLogger = logger.HcLog().Named("import")
+
+	mockContext := &mock.MockContext{
+		RequestCtx: &fasthttp.RequestCtx{},
+		ConfigFn: func() *raiden.Config {
+			return config
+		},
+		SendJsonFn: func(data any) error {
+			return nil
+		},
+	}
+
+	mock := &mock.MockSupabase{Cfg: config}
+	mock.Activate()
+	defer mock.Deactivate()
+
+	dir, errDir := os.MkdirTemp("", "import")
+	assert.NoError(t, errDir)
+	flags.ProjectPath = dir
+
+	err0 := mock.MockGetBucketsWithExpectedResponse(200, []objects.Bucket{
+		{Name: "some_bucket"},
+	})
+	assert.NoError(t, err0)
+
+	err1 := mock.MockGetTablesWithExpectedResponse(200, []objects.Table{
+		{ID: 1, Name: "some_table", Schema: "public"},
+	})
+	assert.NoError(t, err1)
+
+	err2 := mock.MockGetFunctionsWithExpectedResponse(200, []objects.Function{
+		{ID: 1, Schema: "public", Name: "some_function", Definition: "SELECT * FROM some_table;end $function$", ReturnType: "json"},
+	})
+	assert.NoError(t, err2)
+
+	err3 := mock.MockGetRolesWithExpectedResponse(200, []objects.Role{
+		{
+			ID:              1,
+			ConnectionLimit: 10,
+			Name:            "mock_other_role",
+			InheritRole:     true,
+			CanLogin:        true,
+			CanCreateDB:     true,
+			CanCreateRole:   true,
+			CanBypassRLS:    true,
+		},
+	})
+	assert.NoError(t, err3)
+
+	c := controllers.StateReadyController{
+		Result: controllers.StateReadyResponse{},
+	}
+
+	err := c.Post(mockContext)
+	assert.NoError(t, err)
+
+	// validate state
+	localState, err := state.Load()
+	assert.NoError(t, err)
+	assert.NotNil(t, localState)
+	assert.Len(t, localState.Tables, 1)
+	assert.Len(t, localState.Roles, 1)
+	assert.Len(t, localState.Rpc, 1)
+	assert.Len(t, localState.Storage, 1)
+}

pkg/generator/model_register.go

@@ -72,8 +72,10 @@ func GenerateModelRegister(basePath string, projectName string, generateFn Gener
 		return err
 	}
 
+	writer := &FileWriter{FilePath: input.OutputPath}
+
 	ModelRegisterLogger.Debug("generate model register", "path", input.OutputPath)
-	return generateFn(input, nil)
+	return generateFn(input, writer)
 }
 
 func createModelRegisterInput(projectName string, modelRegisterDir string, modelList []string) (input GenerateInput, err error) {

pkg/generator/role_register.go

@@ -72,8 +72,11 @@ func GenerateRoleRegister(basePath string, projectName string, generateFn Genera
 		return err
 	}
 
+	// setup writer
+	writer := &FileWriter{FilePath: input.OutputPath}
+
 	RoleRegisterLogger.Debug("generate role register", "path", input.OutputPath)
-	return generateFn(input, nil)
+	return generateFn(input, writer)
 }
 
 func createRoleRegisterInput(projectName string, roleRegisterDir string, roleList []string) (input GenerateInput, err error) {

pkg/generator/router.go

@@ -62,6 +62,7 @@ import (
 {{end }}
 func RegisterRoute(server *raiden.Server) {
 	server.RegisterRoute([]*raiden.Route{
+		raiden.NewRouteFromController(&raiden_controllers.StateReadyController{}, []string{fasthttp.MethodPost}),
 		{{- range .Routes}}
 		{
 			Type:       {{ .Type }},
@@ -106,8 +107,11 @@ func GenerateRoute(basePath string, projectName string, generateFn GenerateFn) e
 		return err
 	}
 
+	// setup writer
+	writer := &FileWriter{FilePath: input.OutputPath}
+
 	RouterLogger.Debug("generate route", "path", input.OutputPath)
-	return generateFn(input, nil)
+	return generateFn(input, writer)
 }
 
 func WalkScanControllers(controllerPath string) ([]GenerateRouteItem, error) {
@@ -350,6 +354,8 @@ func createRouteInput(projectName string, routePath string, routes []GenerateRou
 	// set imports path
 	imports := []string{
 		fmt.Sprintf("%q", "github.com/sev-2/raiden"),
+		fmt.Sprintf("%q", "github.com/valyala/fasthttp"),
+		fmt.Sprintf("%s %q", "raiden_controllers", "github.com/sev-2/raiden/pkg/controllers"),
 	}
 
 	if len(routes) > 0 {
@@ -358,7 +364,6 @@ func createRouteInput(projectName string, routePath string, routes []GenerateRou
 	}
 
 	isHaveModel := false
-	isHaveMethods := false
 	isHaveStorage := false
 	for i := range routes {
 		r := routes[i]
@@ -367,10 +372,6 @@ func createRouteInput(projectName string, routePath string, routes []GenerateRou
 			isHaveModel = true
 		}
 
-		if r.Methods != "" && r.Methods != "[]string{}" && !isHaveMethods {
-			isHaveMethods = true
-		}
-
 		if r.Storage != "" && !isHaveStorage {
 			isHaveStorage = true
 		}
@@ -381,10 +382,6 @@ func createRouteInput(projectName string, routePath string, routes []GenerateRou
 		imports = append(imports, fmt.Sprintf("%q", modelImportPath))
 	}
 
-	if isHaveMethods {
-		imports = append(imports, fmt.Sprintf("%q", "github.com/valyala/fasthttp"))
-	}
-
 	if isHaveStorage {
 		storageImportPath := fmt.Sprintf("%s/internal/storages", utils.ToGoModuleName(projectName))
 		imports = append(imports, fmt.Sprintf("%q", storageImportPath))

pkg/generator/rpc_register.go

@@ -75,8 +75,11 @@ func GenerateRpcRegister(basePath string, projectName string, generateFn Generat
 		return err
 	}
 
+	// setup writer
+	writer := &FileWriter{FilePath: input.OutputPath}
+
 	RpcRegisterLogger.Debug("generate rpc", "path", input.OutputPath)
-	return generateFn(input, nil)
+	return generateFn(input, writer)
 }
 
 func createRegisterRpcInput(projectName string, rpcRegisterDir string, rpcList []string) (input GenerateInput, err error) {

pkg/generator/storage_register.go

@@ -73,8 +73,11 @@ func GenerateStoragesRegister(basePath string, projectName string, generateFn Ge
 		return err
 	}
 
+	// setup writer
+	writer := &FileWriter{FilePath: input.OutputPath}
+
 	StorageRegisterLogger.Debug("generate storage register", "path", input.OutputPath)
-	return generateFn(input, nil)
+	return generateFn(input, writer)
 }
 
 func createStorageRegisterInput(projectName string, storageRegisterDir string, storageList []string) (input GenerateInput, err error) {

pkg/resource/common.go

@@ -16,16 +16,17 @@ import (
 
 // Flags is struct to binding options when import and apply is run binart
 type Flags struct {
-	ProjectPath   string
-	RpcOnly       bool
-	RolesOnly     bool
-	ModelsOnly    bool
-	StoragesOnly  bool
-	AllowedSchema string
-	DebugMode     bool
-	TraceMode     bool
-	Generate      generate.Flags
-	DryRun        bool
+	ProjectPath     string
+	RpcOnly         bool
+	RolesOnly       bool
+	ModelsOnly      bool
+	StoragesOnly    bool
+	AllowedSchema   string
+	DebugMode       bool
+	TraceMode       bool
+	Generate        generate.Flags
+	UpdateStateOnly bool
+	DryRun          bool
 }
 
 // LoadAll is function to check is all resource need to import or apply

pkg/resource/import.go

@@ -185,12 +185,18 @@ func Import(flags *Flags, config *raiden.Config) error {
 		Rpc:     rpc.GetNewCountData(spResource.Functions, appRpcFunctions),
 	}
 	if !flags.DryRun {
-		// generate resource
-		if err := generateImportResource(config, &importState, flags.ProjectPath, spResource, mapModelValidationTags); err != nil {
-			return err
+		if flags.UpdateStateOnly {
+			return updateStateOnly(&importState, spResource, mapModelValidationTags)
+		} else {
+			// generate resource
+			if err := generateImportResource(config, &importState, flags.ProjectPath, spResource, mapModelValidationTags); err != nil {
+				return err
+			}
+			PrintImportReport(importReport, false)
 		}
-		PrintImportReport(importReport, false)
+
 	} else {
+
 		if len(dryRunError) > 0 {
 			errMessage := strings.Join(dryRunError, "\n")
 			ImportLogger.Error("got error", "err-msg", errMessage)
@@ -303,6 +309,59 @@ func generateImportResource(config *raiden.Config, importState *state.LocalState
 	}
 }
 
+func updateStateOnly(importState *state.LocalState, resource *Resource, mapModelValidationTags map[string]state.ModelValidationTag) error {
+	if len(resource.Tables) > 0 {
+		tableInputs := tables.BuildGenerateModelInputs(resource.Tables, resource.Policies, mapModelValidationTags)
+		for i := range tableInputs {
+			t := tableInputs[i]
+			importState.AddTable(state.TableState{
+				Table:       t.Table,
+				Relation:    t.Relations,
+				Policies:    t.Policies,
+				ModelStruct: utils.SnakeCaseToPascalCase(t.Table.Name),
+				LastUpdate:  time.Now(),
+			})
+		}
+	}
+
+	if len(resource.Roles) > 0 {
+		for i := range resource.Roles {
+			r := resource.Roles[i]
+			importState.AddRole(state.RoleState{
+				Role:       r,
+				IsNative:   false,
+				RoleStruct: utils.SnakeCaseToPascalCase(r.Name),
+				LastUpdate: time.Now(),
+			})
+		}
+	}
+
+	if len(resource.Functions) > 0 {
+		for i := range resource.Functions {
+			f := resource.Functions[i]
+			importState.AddRpc(state.RpcState{
+				Function:   f,
+				RpcStruct:  utils.SnakeCaseToPascalCase(f.Name),
+				LastUpdate: time.Now(),
+			})
+		}
+	}
+
+	if len(resource.Storages) > 0 {
+		storageInputs := storages.BuildGenerateStorageInput(resource.Storages, resource.Policies)
+		for i := range storageInputs {
+			s := storageInputs[i]
+			importState.AddStorage(state.StorageState{
+				Storage:       s.Bucket,
+				StorageStruct: utils.SnakeCaseToPascalCase(s.Bucket.Name),
+				Policies:      s.Policies,
+				LastUpdate:    time.Now(),
+			})
+		}
+	}
+	return importState.Persist()
+}
+
 func ImportDecorateFunc[T any](data []T, findFunc func(T, generator.GenerateInput) bool, stateChan chan any) generator.GenerateFn {
 	return func(input generator.GenerateInput, writer io.Writer) error {
 		if err := generator.Generate(input, writer); err != nil {