Merge pull request #117 from sev-2/tweak

tweak : register postgres Date and rpc register

Author: vani-rf

pkg/cli/generate/command.go

@@ -194,6 +194,12 @@ func Run(flags *Flags, config *raiden.Config, projectPath string, initialize boo
 			}
 			GenerateLogger.Debug("finish generate role register file")
 
+			// generate rpc register
+			GenerateLogger.Debug("start generate rpc register file")
+			if err := generator.GenerateRpcRegister(projectPath, config.ProjectName, generator.Generate); err != nil {
+				errChan <- err
+			}
+			GenerateLogger.Debug("finish generate rpc register file")
 		}
 
 		// generate job register

pkg/generator/model.go

@@ -209,6 +209,8 @@ func MapTableAttributes(projectName string, table objects.Table, mapDataType map
 				"*postgres.Point":    true,
 				"postgres.DateTime":  true,
 				"*postgres.DateTime": true,
+				"postgres.Date":      true,
+				"*postgres.Date":     true,
 			}
 
 			if postgresCustomTypes[column.Type] {

pkg/generator/rpc.go

@@ -637,6 +637,8 @@ func (r *ExtractRpcDataResult) GetParams(mapImports map[string]bool) (columns []
 				importPackageName = "github.com/google/uuid"
 			case "json":
 				importPackageName = "encoding/json"
+			case "postgres":
+				importPackageName = "github.com/sev-2/raiden/pkg/postgres"
 			}
 			key := fmt.Sprintf("%q", importPackageName)
 			mapImports[key] = true
@@ -725,6 +727,8 @@ func (r *ExtractRpcDataResult) GetReturn(mapImports map[string]bool) (returnDecl
 					importPackageName = "github.com/google/uuid"
 				case "json":
 					importPackageName = "encoding/json"
+				case "postgres":
+					importPackageName = "github.com/sev-2/raiden/pkg/postgres"
 				}
 				key := fmt.Sprintf("%q", importPackageName)
 				mapImports[key] = true

pkg/generator/rpc_test.go

@@ -178,8 +178,8 @@ func TestExtractRpcWithPrefix(t *testing.T) {
 		where sc.name = in_scouter_name and c.name = in_candidate_name ; end;
 		`,
 		CompleteStatement: `
-		CREATE OR REPLACE FUNCTION public.get_submissions(in_scouter_name character varying, in_candidate_name character varying)\n 
-		RETURNS TABLE(id integer, created_at timestamp without time zone, sc_name character varying, c_name character varying)\n 
+		CREATE OR REPLACE FUNCTION public.get_submissions(in_scouter_name character varying, in_candidate_name character varying, in_register date)\n 
+		RETURNS TABLE(id integer, created_at timestamp without time zone, sc_name character varying, c_name character varying, in_register date)\n 
 		LANGUAGE plpgsql\n
 		AS $function$ 
 			begin return query 
@@ -202,6 +202,12 @@ func TestExtractRpcWithPrefix(t *testing.T) {
 				TypeId:     1043,
 				HasDefault: false,
 			},
+			{
+				Mode:       "in",
+				Name:       "in_register",
+				TypeId:     1043,
+				HasDefault: false,
+			},
 			{
 				Mode:       "table",
 				Name:       "id",
@@ -227,10 +233,10 @@ func TestExtractRpcWithPrefix(t *testing.T) {
 				HasDefault: false,
 			},
 		},
-		ArgumentTypes:          "in_scouter_name character varying, in_candidate_name character varying",
-		IdentityArgumentTypes:  "in_scouter_name character varying, in_candidate_name character varying",
+		ArgumentTypes:          "in_scouter_name character varying, in_candidate_name character varying, in_register date",
+		IdentityArgumentTypes:  "in_scouter_name character varying, in_candidate_name character varying, in_register date",
 		ReturnTypeID:           2249,
-		ReturnType:             "TABLE(id integer, created_at timestamp without time zone, sc_name character varying, c_name character varying)",
+		ReturnType:             "TABLE(id integer, created_at timestamp without time zone, sc_name character varying, c_name character varying, register date)",
 		ReturnTypeRelationID:   0,
 		IsSetReturningFunction: true,
 		Behavior:               string(raiden.RpcBehaviorVolatile),
@@ -518,6 +524,106 @@ func TestGenerateRpc(t *testing.T) {
 	assert.FileExists(t, dir+"/internal/rpc/get_submissions.go")
 }
 
+func TestGenerateRpc_DateType(t *testing.T) {
+	fns := []objects.Function{
+		{
+			Schema:            "public",
+			Name:              "get_latest_active_rates_by_tenant",
+			Language:          "sql",
+			Definition:        "\n  select distinct on (tr.tax_id)\n    tr.id,\n    tr.tax_id,\n    tr.type,\n    tr.rate,\n    tr.start_date,\n    tr.end_date,\n    tr.rate as applicable_rate,\n    t.name as tax_name,\n    t.tenant::text\n  from tax_rates tr\n  join taxes t on tr.tax_id = t.id\n  where t.tenant = input_tenant::public.tenant\n    and tr.start_date <= now()\n    and (tr.end_date is null or tr.end_date > now())\n  order by tr.tax_id, tr.start_date desc\n",
+			CompleteStatement: "CREATE OR REPLACE FUNCTION public.get_latest_active_rates_by_tenant(input_tenant text)\n RETURNS TABLE(id bigint, tax_id bigint, type text, rate numeric, start_date date, end_date date, applicable_rate numeric, tax_name text, tenant text)\n LANGUAGE sql\nAS $function$\n  select distinct on (tr.tax_id)\n    tr.id,\n    tr.tax_id,\n    tr.type,\n    tr.rate,\n    tr.start_date,\n    tr.end_date,\n    tr.rate as applicable_rate,\n    t.name as tax_name,\n    t.tenant::text\n  from tax_rates tr\n  join taxes t on tr.tax_id = t.id\n  where t.tenant = input_tenant::public.tenant\n    and tr.start_date <= now()\n    and (tr.end_date is null or tr.end_date > now())\n  order by tr.tax_id, tr.start_date desc\n$function$\n",
+			Args: []objects.FunctionArg{
+				{
+					Mode:       "in",
+					Name:       "input_tenant",
+					TypeId:     25,
+					HasDefault: false,
+				},
+				{
+					Mode:       "in",
+					Name:       "input_register",
+					TypeId:     25,
+					HasDefault: false,
+				},
+				{
+					Mode:       "table",
+					Name:       "id",
+					TypeId:     20,
+					HasDefault: false,
+				},
+				{
+					Mode:       "table",
+					Name:       "tax_id",
+					TypeId:     20,
+					HasDefault: false,
+				},
+				{
+					Mode:       "table",
+					Name:       "type",
+					TypeId:     25,
+					HasDefault: false,
+				},
+				{
+					Mode:       "table",
+					Name:       "rate",
+					TypeId:     1700,
+					HasDefault: false,
+				},
+				{
+					Mode:       "table",
+					Name:       "start_date",
+					TypeId:     1082,
+					HasDefault: false,
+				},
+				{
+					Mode:       "table",
+					Name:       "end_date",
+					TypeId:     1082,
+					HasDefault: false,
+				},
+				{
+					Mode:       "table",
+					Name:       "applicable_rate",
+					TypeId:     1700,
+					HasDefault: false,
+				},
+				{
+					Mode:       "table",
+					Name:       "tax_name",
+					TypeId:     25,
+					HasDefault: false,
+				},
+				{
+					Mode:       "table",
+					Name:       "tenant",
+					TypeId:     25,
+					HasDefault: false,
+				},
+			},
+			ArgumentTypes:          "input_tenant text, input_register date",
+			IdentityArgumentTypes:  "input_tenant text, input_register date",
+			ReturnTypeID:           2249,
+			ReturnType:             "TABLE(id bigint, tax_id bigint, type text, rate numeric, start_date date, end_date date, applicable_rate numeric, tax_name text, tenant text)",
+			ReturnTypeRelationID:   0,
+			IsSetReturningFunction: true,
+			Behavior:               string(raiden.RpcBehaviorVolatile),
+			SecurityDefiner:        false,
+			ConfigParams:           nil,
+		},
+	}
+
+	dir, err := os.MkdirTemp("", "rpc")
+	assert.NoError(t, err)
+
+	rpcPath := filepath.Join(dir, "internal")
+	err1 := utils.CreateFolder(rpcPath)
+	assert.NoError(t, err1)
+
+	err2 := generator.GenerateRpc(dir, "test", fns, []objects.Table{}, generator.GenerateFn(generator.Generate))
+	assert.NoError(t, err2)
+	assert.FileExists(t, dir+"/internal/rpc/get_latest_active_rates_by_tenant.go")
+}
+
 func TestRpcWithTrigger(t *testing.T) {
 	fn := objects.Function{
 		Schema:     "public",

pkg/postgres/data_type.go

@@ -114,8 +114,10 @@ func ToGoType(pgType DataType, isNullable bool) (goType string) {
 		goType = "float64"
 	case VarcharType, VarcharTypeAlias, CharType, BpcharType, TextType:
 		goType = "string"
-	case TimestampType, TimestampTypeAlias, TimestampTzType, TimestampTzTypeAlias, TimeType, TimeTypeAlias, TimeTzType, TimeTzTypeAlias, DateType:
+	case TimestampType, TimestampTypeAlias, TimestampTzType, TimestampTzTypeAlias, TimeType, TimeTypeAlias, TimeTzType, TimeTzTypeAlias:
 		goType = "postgres.DateTime"
+	case DateType:
+		goType = "postgres.Date"
 	case IntervalType:
 		goType = "time.Duration"
 	case BooleanType:
@@ -160,6 +162,8 @@ func ToPostgresType(goType string) (pgType DataType) {
 		pgType = TextType
 	case "postgres.DateTime":
 		pgType = TimestampType
+	case "postgres.Date":
+		pgType = DateType
 	case "time.Duration":
 		pgType = IntervalType
 	case "bool":

pkg/postgres/date.go

@@ -0,0 +1,65 @@
+package postgres
+
+import (
+	"database/sql/driver"
+	"fmt"
+	"strings"
+	"time"
+)
+
+type Date struct {
+	time.Time
+}
+
+// Ensure Date implements sql.Scanner and driver.Valuer interfaces
+var _ driver.Valuer = (*Date)(nil)
+var _ fmt.Stringer = (*Date)(nil)
+
+const dateFormat = "2006-01-02"
+
+func (d *Date) Scan(value interface{}) error {
+	switch v := value.(type) {
+	case time.Time:
+		d.Time = v.UTC().Truncate(24 * time.Hour)
+		return nil
+	case string:
+		t, err := time.Parse(dateFormat, v)
+		if err != nil {
+			return err
+		}
+		d.Time = t
+		return nil
+	case []byte:
+		t, err := time.Parse(dateFormat, string(v))
+		if err != nil {
+			return err
+		}
+		d.Time = t
+		return nil
+	default:
+		return fmt.Errorf("cannot scan type %T into Date", value)
+	}
+}
+
+func (d Date) Value() (driver.Value, error) {
+	return d.Format(dateFormat), nil
+}
+
+func (d Date) String() string {
+	return d.Format(dateFormat)
+}
+
+// JSON Marshal/Unmarshal support
+func (d Date) MarshalJSON() ([]byte, error) {
+	return []byte(fmt.Sprintf(`"%s"`, d.Format(dateFormat))), nil
+}
+
+func (d *Date) UnmarshalJSON(b []byte) error {
+	s := strings.Trim(string(b), `"`)
+	t, err := time.Parse(dateFormat, s)
+	if err != nil {
+		return err
+	}
+	d.Time = t
+	return nil
+}

pkg/postgres/date_test.go

@@ -0,0 +1,92 @@
+package postgres_test
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/sev-2/raiden/pkg/postgres"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDate_ScanFromString(t *testing.T) {
+	var d postgres.Date
+	err := d.Scan("2025-05-22")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if d.Format("2006-01-02") != "2025-05-22" {
+		t.Errorf("expected 2025-05-22, got %s", d.String())
+	}
+}
+
+func TestDate_ScanFromString_Err(t *testing.T) {
+	var d postgres.Date
+	err := d.Scan("2025-05-22-01")
+	assert.Error(t, err)
+}
+
+func TestDate_ScanFromBytes(t *testing.T) {
+	var d postgres.Date
+	err := d.Scan([]byte("2025-05-22"))
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if d.Format("2006-01-02") != "2025-05-22" {
+		t.Errorf("expected 2025-05-22, got %s", d.String())
+	}
+}
+
+func TestDate_ScanFromTime(t *testing.T) {
+	expected := time.Date(2025, 5, 22, 12, 0, 0, 0, time.UTC)
+	var d postgres.Date
+	err := d.Scan(expected)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if !d.Time.Equal(expected.Truncate(24 * time.Hour)) {
+		t.Errorf("expected %v, got %v", expected, d.Time)
+	}
+}
+
+func TestDateValue(t *testing.T) {
+	d := postgres.Date{Time: time.Date(2025, 5, 22, 0, 0, 0, 0, time.UTC)}
+	val, err := d.Value()
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if str, ok := val.(string); !ok || str != "2025-05-22" {
+		t.Errorf("expected string 2025-05-22, got %v", val)
+	}
+}
+
+func TestDate_JSONMarshaling(t *testing.T) {
+	d := postgres.Date{Time: time.Date(2025, 5, 22, 0, 0, 0, 0, time.UTC)}
+	jsonBytes, err := json.Marshal(d)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	expected := `"2025-05-22"`
+	if string(jsonBytes) != expected {
+		t.Errorf("expected %s, got %s", expected, string(jsonBytes))
+	}
+}
+
+func TestDate_JSONUnmarshaling(t *testing.T) {
+	var d postgres.Date
+	err := json.Unmarshal([]byte(`"2025-05-22"`), &d)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if d.String() != "2025-05-22" {
+		t.Errorf("expected 2025-05-22, got %s", d.String())
+	}
+}
+
+func TestDate_ScanInvalidType(t *testing.T) {
+	var d postgres.Date
+	err := d.Scan(123)
+	if err == nil {
+		t.Fatalf("expected error, got nil")
+	}
+}