From c235dd8b2a462fc6b12d31c55e4a357a23481271 Mon Sep 17 00:00:00 2001
From: Vadim Voitenko <vadim.voitenko@exness.com>
Date: Wed, 7 Feb 2024 16:11:09 +0200
Subject: [PATCH] Json transformer and connection string fixes

* Fixed Json transformer: previously it did not parse operation fields correctly
* Fixed database connection string builder in pg_restore and pg_dump wrapper
---
 internal/db/postgres/pgdump/pgdump.go         | 16 ++++++++-
 internal/db/postgres/pgrestore/pgrestore.go   | 19 ++++++++--
 internal/db/postgres/transformers/json.go     | 12 +++----
 .../db/postgres/transformers/json_test.go     | 35 +++++++++++++++++++
 4 files changed, 72 insertions(+), 10 deletions(-)

diff --git a/internal/db/postgres/pgdump/pgdump.go b/internal/db/postgres/pgdump/pgdump.go
index d511c3e5..2e3bac0d 100644
--- a/internal/db/postgres/pgdump/pgdump.go
+++ b/internal/db/postgres/pgdump/pgdump.go
@@ -109,7 +109,21 @@ func (o *Options) GetPgDSN() (string, error) {
 		return o.DbName, nil
 	}
 
-	return fmt.Sprintf("host=%s port=%d user=%s dbname=%s", o.Host, o.Port, o.UserName, o.DbName), nil
+	var parts []string
+	if o.Host != "" {
+		parts = append(parts, fmt.Sprintf("host=%s", o.Host))
+	}
+	if o.Port != 5432 {
+		parts = append(parts, fmt.Sprintf("port=%d", o.Port))
+	}
+	if o.UserName != "" {
+		parts = append(parts, fmt.Sprintf("port=%d", o.Port))
+	}
+	if o.DbName != "" {
+		parts = append(parts, fmt.Sprintf("dbname=%d", o.Port))
+	}
+
+	return strings.Join(parts, " "), nil
 }
 
 func (o *Options) GetParams() []string {
diff --git a/internal/db/postgres/pgrestore/pgrestore.go b/internal/db/postgres/pgrestore/pgrestore.go
index 3eadea6b..578a3f85 100644
--- a/internal/db/postgres/pgrestore/pgrestore.go
+++ b/internal/db/postgres/pgrestore/pgrestore.go
@@ -98,11 +98,24 @@ type Options struct {
 }
 
 func (o *Options) GetPgDSN() (string, error) {
-	//return "host=localhost port=5432 user=postgres dbname=postgres", nil
-	if strings.Contains(o.DbName, "=") {
+	if strings.HasPrefix(o.DbName, "postgresql://") || strings.Contains(o.DbName, "=") {
 		return o.DbName, nil
 	}
-	return fmt.Sprintf("host=%s port=%d user=%s dbname=%s", o.Host, o.Port, o.UserName, o.DbName), nil
+
+	var parts []string
+	if o.Host != "" {
+		parts = append(parts, fmt.Sprintf("host=%s", o.Host))
+	}
+	if o.Port != 5432 {
+		parts = append(parts, fmt.Sprintf("port=%d", o.Port))
+	}
+	if o.UserName != "" {
+		parts = append(parts, fmt.Sprintf("port=%d", o.Port))
+	}
+	if o.DbName != "" {
+		parts = append(parts, fmt.Sprintf("dbname=%d", o.Port))
+	}
+	return strings.Join(parts, " "), nil
 }
 
 func (o *Options) GetParams() []string {
diff --git a/internal/db/postgres/transformers/json.go b/internal/db/postgres/transformers/json.go
index 08b76062..77f61f2f 100644
--- a/internal/db/postgres/transformers/json.go
+++ b/internal/db/postgres/transformers/json.go
@@ -74,12 +74,12 @@ var jsonSetOpt = &sjson.Options{
 }
 
 type Operation struct {
-	Operation     string      `mapstructure:"operation"`
-	Value         interface{} `mapstructure:"value,omitempty"`
-	ValueTemplate string      `mapstructure:"value_template,omitempty"`
-	Path          string      `mapstructure:"path"`
-	ErrorNotExist bool        `mapstructure:"error_not_exist"`
-	tmpl          *template.Template
+	Operation     string             `mapstructure:"operation" json:"operation"`
+	Value         interface{}        `mapstructure:"value,omitempty" json:"value"`
+	ValueTemplate string             `mapstructure:"value_template,omitempty" json:"value_template"`
+	Path          string             `mapstructure:"path" json:"path"`
+	ErrorNotExist bool               `mapstructure:"error_not_exist" json:"error_not_exist"`
+	tmpl          *template.Template `json:"tmpl,omitempty"`
 }
 
 func (o *Operation) Apply(inp []byte, tctx *JsonContext, buf *bytes.Buffer) ([]byte, error) {
diff --git a/internal/db/postgres/transformers/json_test.go b/internal/db/postgres/transformers/json_test.go
index 7144b423..0600d2e7 100644
--- a/internal/db/postgres/transformers/json_test.go
+++ b/internal/db/postgres/transformers/json_test.go
@@ -145,3 +145,38 @@ func TestJsonTransformer_Transform_null(t *testing.T) {
 	resValue := string(res.Data)
 	require.JSONEq(t, expected, resValue)
 }
+
+func TestJsonTransformer_structure_tags_encoding_regression(t *testing.T) {
+	// The unexpected behavior described in issue https://github.com/GreenmaskIO/greenmask/issues/4
+	// The problem was that the Operation object did not have an appropriate json tag on the fields
+	rawData := []byte(`
+		[
+			{
+				"operation": "set", 
+				"path": "name", 
+				"value_template": "template_tes", 
+				"value": "value_test", 
+				"error_not_exist": true
+			}
+		]
+	`)
+
+	expected := &Operation{
+		Operation:     "set",
+		Path:          "name",
+		ValueTemplate: "template_tes",
+		Value:         "value_test",
+		ErrorNotExist: true,
+	}
+
+	var ops []*Operation
+	err := json.Unmarshal(rawData, &ops)
+	require.NoError(t, err)
+	require.Len(t, ops, 1)
+	op := ops[0]
+	assert.Equal(t, expected.Operation, op.Operation)
+	assert.Equal(t, expected.Path, op.Path)
+	assert.Equal(t, expected.ValueTemplate, op.ValueTemplate)
+	assert.Equal(t, expected.Value, op.Value)
+	assert.Equal(t, expected.ErrorNotExist, op.ErrorNotExist)
+}