From f484436c4385a67d14e627e5c039d4c6ea6d6a0d Mon Sep 17 00:00:00 2001
From: Sarah Wooders <sarahwooders@gmail.com>
Date: Thu, 26 Oct 2023 12:02:29 -0700
Subject: [PATCH] add db connector

---
 memgpt/connectors/connector.py | 31 ++++++++++++++++++++++++++++++-
 memgpt/utils.py                |  7 +++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/memgpt/connectors/connector.py b/memgpt/connectors/connector.py
index 2b3c22a2..8e2b828c 100644
--- a/memgpt/connectors/connector.py
+++ b/memgpt/connectors/connector.py
@@ -9,7 +9,6 @@ memgpt load <data-connector-type> --name <dataset-name> [ADDITIONAL ARGS]
 """
 
 from llama_index import download_loader
-from llama_index.embeddings import OpenAIEmbedding
 from typing import List
 import os
 import typer
@@ -64,3 +63,33 @@ def load_webpage(
     # save connector information into .memgpt metadata file
     save_index(index, name)
 
+
+@app.command("database")
+def load_database(
+    name: str = typer.Option(help="Name of dataset to load."),
+    scheme: str = typer.Option(help="Database scheme."),
+    host: str = typer.Option(help="Database host."),
+    port: int = typer.Option(help="Database port."),
+    user: str = typer.Option(help="Database user."),
+    password: str = typer.Option(help="Database password."),
+    dbname: str = typer.Option(help="Database name."),
+    query: str = typer.Option(None, help="Database query."),
+):
+    from llama_index.readers.database import DatabaseReader
+    
+    db = DatabaseReader(
+        scheme="postgresql",  # Database Scheme
+        host="localhost",  # Database Host
+        port="5432",  # Database Port
+        user="postgres",  # Database User
+        password="FakeExamplePassword",  # Database Password
+        dbname="postgres",  # Database Name
+    )
+
+    # load data
+    docs = db.load_data(query=query)
+    
+    index = index_docs(docs)
+    save_index(index, name)
+
+
diff --git a/memgpt/utils.py b/memgpt/utils.py
index 92c82978..a8b5a07a 100644
--- a/memgpt/utils.py
+++ b/memgpt/utils.py
@@ -415,6 +415,13 @@ def save_index(index, name):
     # TODO: load directory from config 
     # TODO: save to vectordb/local depending on config
     dir = f"{MEMGPT_DIR}/archival/{name}"
+
+    # check if directory exists
+    if os.path.exists(dir):
+        confirm = typer.confirm(typer.style(f"Index with name {name} already exists -- overwrite?", fg="red"), default=False)
+        if not confirm:
+            typer.secho("Aborting.", fg="red")
+            exit()
     # create directory, even if it already exists
     os.makedirs(dir, exist_ok=True)
     index.storage_context.persist(dir)