From 80d11a1972891cca7c6868a340e48dad1ffaddb2 Mon Sep 17 00:00:00 2001
From: leehuwuj <leehuwuj@gmail.com>
Date: Fri, 29 Mar 2024 15:33:25 +0700
Subject: [PATCH] add support for PostgreSQL data source

---
 helpers/python.ts |  8 ++++++--
 questions.ts      | 11 ++++++++---
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/helpers/python.ts b/helpers/python.ts
index d27df466..65aeff5c 100644
--- a/helpers/python.ts
+++ b/helpers/python.ts
@@ -89,6 +89,10 @@ const getAdditionalDependencies = (
         version: "^1.1.0",
         extras: ["rsa"],
       });
+      dependencies.push({
+        name: "psycopg",
+        version: "^3.1.18",
+      });
       break;
   }
 
@@ -335,8 +339,8 @@ export const installPythonTemplate = async ({
       });
 
       const node = dbLoaderConfig.createNode(configEntries);
-      node.commentBefore = ` The configuration for the database loader, only supports MySQL database for now.
- uri: The URI for the database. E.g.: mysql+pymysql://user:password@localhost:3306/db.
+      node.commentBefore = ` The configuration for the database loader, only supports MySQL and PostgreSQL databases for now.
+ uri: The URI for the database. E.g.: mysql+pymysql://user:password@localhost:3306/db or postgresql+psycopg2://user:password@localhost:5432/db
  query: The query to fetch data from the database. E.g.: SELECT * FROM table`;
       loaderConfig.set("db", node);
     }
diff --git a/questions.ts b/questions.ts
index b9637b4a..231be855 100644
--- a/questions.ts
+++ b/questions.ts
@@ -160,7 +160,7 @@ export const getDataSourceChoices = (
       value: "web",
     });
     choices.push({
-      title: "Use data from a database (Mysql)",
+      title: "Use data from a database (Mysql, PostgreSQL)",
       value: "db",
     });
   }
@@ -696,8 +696,13 @@ export const askQuestions = async (
                 validate: (value: string) => {
                   if (!value) {
                     return "Please provide a valid connection string";
-                  } else if (!value.startsWith("mysql+pymysql://")) {
-                    return "The connection string must start with 'mysql+pymysql://'";
+                  } else if (
+                    !(
+                      value.startsWith("mysql+pymysql://") ||
+                      value.startsWith("postgresql+psycopg://")
+                    )
+                  ) {
+                    return "The connection string must start with 'mysql+pymysql://' for MySQL or 'postgresql+psycopg://' for PostgreSQL";
                   }
                   return true;
                 },