Create migrations for new db

airtai · Jun 20, 2024 · 9abf4cc · 9abf4cc
1 parent 583a2a5
commit 9abf4cc
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 5 deletions.
diff --git a/google_sheets/app.py b/google_sheets/app.py
@@ -11,7 +11,7 @@
 from prisma.errors import RecordNotFoundError
 
 from . import __version__
-from .db_helpers import get_db_connection, get_wasp_db_url
+from .db_helpers import get_db_connection, get_db_url
 
 __all__ = ["app"]
 
@@ -45,7 +45,7 @@
 
 
 async def get_user(user_id: Union[int, str]) -> Any:
-    wasp_db_url = await get_wasp_db_url()
+    wasp_db_url = await get_db_url(db_name="waspdb")
     async with get_db_connection(db_url=wasp_db_url) as db:
         user = await db.query_first(
             f'SELECT * from "User" where id={user_id}'  # nosec: [B608]

diff --git a/google_sheets/db_helpers.py b/google_sheets/db_helpers.py
@@ -10,7 +10,7 @@ async def get_db_connection(
     db_url: Optional[str] = None,
 ) -> AsyncGenerator[Prisma, None]:
     if not db_url:
-        db_url = environ.get("DATABASE_URL", None)
+        db_url = await get_db_url(db_name="gsheets")
         if not db_url:
             raise ValueError(
                 "No database URL provided nor set as environment variable 'DATABASE_URL'"
@@ -25,9 +25,9 @@ async def get_db_connection(
         await db.disconnect()
 
 
-async def get_wasp_db_url() -> str:
+async def get_db_url(db_name: str) -> str:
     curr_db_url = environ.get("DATABASE_URL")
-    wasp_db_name = environ.get("WASP_DB_NAME", "waspdb")
+    wasp_db_name = environ.get("WASP_DB_NAME", db_name)
     wasp_db_url = curr_db_url.replace(curr_db_url.split("/")[-1], wasp_db_name)  # type: ignore[union-attr]
     if "connect_timeout" not in wasp_db_url:
         wasp_db_url += "?connect_timeout=60"

diff --git a/migrations/20240620124134_initial/migration.sql b/migrations/20240620124134_initial/migration.sql
@@ -0,0 +1,14 @@
+-- CreateTable
+CREATE TABLE "GAuth" (
+    "id" TEXT NOT NULL,
+    "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "updated_at" TIMESTAMP(3) NOT NULL,
+    "user_id" INTEGER NOT NULL,
+    "creds" JSONB NOT NULL,
+    "info" JSONB NOT NULL,
+
+    CONSTRAINT "GAuth_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateIndex
+CREATE UNIQUE INDEX "GAuth_user_id_key" ON "GAuth"("user_id");
diff --git a/migrations/migration_lock.toml b/migrations/migration_lock.toml
@@ -0,0 +1,3 @@
+# Please do not edit this file manually
+# It should be added in your version-control system (i.e. Git)
+provider = "postgresql"