add support for PostgreSQL data source

run-llama · Mar 29, 2024 · 80d11a1 · 80d11a1
1 parent 52bf693
commit 80d11a1
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git 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
@@ -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;
                 },