fix: Add schema option to connection config (#1112)

cloudquery · Jul 8, 2022 · 4e597e4 · 4e597e4
1 parent 8e6cd35
commit 4e597e4
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 6 deletions.
diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -54,6 +54,7 @@ type Connection struct {
 	Port     int      `yaml:"port,omitempty" json:"port,omitempty"`
 	Database string   `yaml:"database,omitempty" json:"database,omitempty"`
 	SSLMode  string   `yaml:"sslmode,omitempty" json:"sslmode,omitempty"`
+	Schema   string   `yaml:"schema,omitempty" json:"schema,omitempty"`
 	Extras   []string `yaml:"extras,omitempty" json:"extras,omitempty"`
 }
 
@@ -96,7 +97,7 @@ func (c CloudQuery) GetRequiredProvider(name string) (*RequiredProvider, error)
 }
 
 func (c Connection) IsAnyConnParamsSet() bool {
-	return c.Type != "" || c.Username != "" || c.Password != "" || c.Host != "" || c.Port != 0 || c.Database != "" || c.SSLMode != "" || len(c.Extras) > 0
+	return c.Type != "" || c.Username != "" || c.Password != "" || c.Host != "" || c.Port != 0 || c.Database != "" || c.SSLMode != "" || c.Schema != "" || len(c.Extras) > 0
 }
 
 func (c *Connection) BuildFromConnParams() {
@@ -118,11 +119,22 @@ func (c *Connection) BuildFromConnParams() {
 	}
 
 	v := url.Values{}
+
+	if c.Schema != "" {
+		v.Set("search_path", c.Schema)
+	} else {
+		v.Set("search_path", "public") // default
+	}
+
 	if c.Extras != nil {
 		for _, extra := range c.Extras {
 			parts := strings.SplitN(extra, "=", 2)
 			if len(parts) == 1 {
 				v.Add(parts[0], "")
+				continue
+			}
+			if parts[0] == "search_path" {
+				v.Set(parts[0], parts[1])
 			} else {
 				v.Add(parts[0], parts[1])
 			}

diff --git a/pkg/config/config_parser_test.go b/pkg/config/config_parser_test.go
@@ -136,7 +136,7 @@ func Test_ProcessConfig_Connection(t *testing.T) {
 				Host:     `localhost`,
 				Database: `postgres`,
 			},
-			"postgres://user:pass@localhost:5432/postgres",
+			"postgres://user:pass@localhost:5432/postgres?search_path=public",
 			false,
 		},
 		{
@@ -148,7 +148,20 @@ func Test_ProcessConfig_Connection(t *testing.T) {
 				Port:     15432,
 				Database: `postgres`,
 			},
-			"postgres://user@localhost:15432/postgres",
+			"postgres://user@localhost:15432/postgres?search_path=public",
+			false,
+		},
+		{
+			"should use the provided schema if specified",
+			&Connection{
+				Username: `user`,
+				Type:     `postgres`,
+				Host:     `localhost`,
+				Port:     5432,
+				Database: `postgres`,
+				Schema:   `test`,
+			},
+			"postgres://user@localhost:5432/postgres?search_path=test",
 			false,
 		},
 		{
@@ -161,7 +174,20 @@ func Test_ProcessConfig_Connection(t *testing.T) {
 				SSLMode:  `disable`,
 				Extras:   []string{"a=b", "c=d", "e", "sslmode=enable"},
 			},
-			"postgres://user:pass@localhost:5432/postdb?a=b&c=d&e=&sslmode=disable",
+			"postgres://user:pass@localhost:5432/postdb?a=b&c=d&e=&search_path=public&sslmode=disable",
+			false,
+		},
+		{
+			"extras should override search_path",
+			&Connection{
+				Username: `user`,
+				Password: `pass`,
+				Host:     `localhost`,
+				Database: `postdb`,
+				SSLMode:  `disable`,
+				Extras:   []string{"a=b", "c=d", "e", "sslmode=enable", "search_path=test2"},
+			},
+			"postgres://user:pass@localhost:5432/postdb?a=b&c=d&e=&search_path=test2&sslmode=disable",
 			false,
 		},
 		{

diff --git a/pkg/config/loader_test.go b/pkg/config/loader_test.go
@@ -226,7 +226,7 @@ cloudquery:
     - name: aws
       version: latest
 `,
-			"postgres://postgres:pass@localhost:15432/cq?sslmode=disable",
+			"postgres://postgres:pass@localhost:15432/cq?search_path=public&sslmode=disable",
 			false,
 		},
 		{

diff --git a/pkg/config/schema.json b/pkg/config/schema.json
@@ -38,6 +38,10 @@
               "description": "Connection port",
               "type": "integer"
             },
+            "schema": {
+              "description": "Connection schema",
+              "type": "string"
+            },
             "sslmode": {
               "description": "Connection sslmode",
               "type": "string"

diff --git a/pkg/errors/error.go b/pkg/errors/error.go
@@ -104,9 +104,10 @@ func shouldIgnorePgCode(code string) bool {
 		// Class 08 - Connection Exception
 		// Class 28 - Invalid Authorization Specification
 		// Class 3D - Invalid Catalog Name
+		// Class 3F - Invalid Schema Name
 		// Class 53 - Insufficient Resources
 		// Class 57 - Operator Intervention
-		case "08", "28", "3D", "53", "57":
+		case "08", "28", "3D", "3F", "53", "57":
 			return true
 		}
 	}