also add dsn for postgres

rilldata · Dec 28, 2024 · 8b052d5 · 8b052d5
1 parent df1006a
commit 8b052d5
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/runtime/drivers/duckdb/model_executor_sqlstore_self.go b/runtime/drivers/duckdb/model_executor_sqlstore_self.go
@@ -92,7 +92,7 @@ func (e *sqlStoreToSelfExecutor) modelInputProperties(modelName, inputConnector
 			if err := mapstructure.Decode(inputHandle.Config(), &config); err != nil {
 				return nil, err
 			}
-			dsn = config.DatabaseURL
+			dsn = config.ResolveDSN()
 		}
 		if dsn == "" {
 			return nil, fmt.Errorf("must set `dsn` for models that transfer data from `postgres` to `duckdb`")

diff --git a/runtime/drivers/postgres/postgres.go b/runtime/drivers/postgres/postgres.go
@@ -58,6 +58,14 @@ type driver struct{}
 
 type ConfigProperties struct {
 	DatabaseURL string `mapstructure:"database_url"`
+	DSN         string `mapstructure:"dsn"`
+}
+
+func (c *ConfigProperties) ResolveDSN() string {
+	if c.DSN != "" {
+		return c.DSN
+	}
+	return c.DatabaseURL
 }
 
 func (d driver) Open(instanceID string, config map[string]any, st *storage.Client, ac *activity.Client, logger *zap.Logger) (drivers.Handle, error) {