From ac4562c8693eca5fb772d0804d688dc5ab3b842f Mon Sep 17 00:00:00 2001
From: Anshul Khandelwal <12948312+k-anshul@users.noreply.github.com>
Date: Sat, 28 Dec 2024 07:22:00 +0530
Subject: [PATCH] also add dsn for postgres

---
 .../duckdb/model_executor_sqlstore_self.go    | 27 ++++++++++++-------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/runtime/drivers/duckdb/model_executor_sqlstore_self.go b/runtime/drivers/duckdb/model_executor_sqlstore_self.go
index 7fe33e68a43..08e60b28d45 100644
--- a/runtime/drivers/duckdb/model_executor_sqlstore_self.go
+++ b/runtime/drivers/duckdb/model_executor_sqlstore_self.go
@@ -12,14 +12,25 @@ import (
 )
 
 type sqlStoreToSelfInputProps struct {
-	SQL string `mapstructure:"sql"`
-	DSN string `mapstructure:"dsn"`
+	SQL         string `mapstructure:"sql"`
+	DSN         string `mapstructure:"dsn"`
+	DatabaseURL string `mapstructure:"database_url"`
+}
+
+func (p *sqlStoreToSelfInputProps) resolveDSN() string {
+	if p.DSN != "" {
+		return p.DSN
+	}
+	return p.DatabaseURL
 }
 
 func (p *sqlStoreToSelfInputProps) Validate() error {
 	if p.SQL == "" {
 		return fmt.Errorf("missing property 'sql'")
 	}
+	if p.DSN != "" && p.DatabaseURL != "" {
+		return fmt.Errorf("cannot set both 'dsn' and 'database_url'")
+	}
 	return nil
 }
 
@@ -66,10 +77,8 @@ func (e *sqlStoreToSelfExecutor) modelInputProperties(modelName, inputConnector
 	userQuery, _ := strings.CutSuffix(inputProps.SQL, ";") // trim trailing semi colon
 	switch inputHandle.Driver() {
 	case "mysql":
-		var dsn string
-		if inputProps.DSN != "" {
-			dsn = inputProps.DSN
-		} else {
+		dsn := inputProps.resolveDSN()
+		if dsn == "" {
 			// may be configured via a connector
 			var config *mysql.ConfigProperties
 			if err := mapstructure.Decode(inputHandle.Config(), &config); err != nil {
@@ -83,10 +92,8 @@ func (e *sqlStoreToSelfExecutor) modelInputProperties(modelName, inputConnector
 		m.PreExec = fmt.Sprintf("INSTALL 'MYSQL'; LOAD 'MYSQL'; ATTACH %s AS %s (TYPE mysql, READ_ONLY)", safeSQLString(dsn), safeDBName)
 		m.SQL = fmt.Sprintf("SELECT * FROM mysql_query(%s, %s)", safeSQLString(dbName), safeSQLString(userQuery))
 	case "postgres":
-		var dsn string
-		if inputProps.DSN != "" {
-			dsn = inputProps.DSN
-		} else {
+		dsn := inputProps.resolveDSN()
+		if dsn == "" {
 			// may be configured via a connector
 			var config *postgres.ConfigProperties
 			if err := mapstructure.Decode(inputHandle.Config(), &config); err != nil {