Change studio ID type from UUID to integer (#938)

All uses of UUIDs to manage the `studios` table have been switched to a plain `INTEGER` as a primary key. **Warning**: This will drop the current studio tables. Given that a core ID was changed and the feature is not yet released, we just clear them instead of a more complex migration.
BloopAI · Sep 11, 2023 · 23156dd · 23156dd
1 parent b56b4e9
commit 23156dd
Show file tree

Hide file tree

Showing 4 changed files with 119 additions and 89 deletions.
diff --git a/server/bleep/migrations/20230907154037_studio_int_id.sql b/server/bleep/migrations/20230907154037_studio_int_id.sql
@@ -0,0 +1,21 @@
+DROP TABLE studios;
+DROP TABLE studio_snapshots;
+
+-- Identical to before, except we change the primary key to be of type `INTEGER`.
+CREATE TABLE studios (
+    id INTEGER PRIMARY KEY,
+    name TEXT NOT NULL,
+    user_id TEXT NOT NULL
+);
+
+-- Identical to before, except we change the `studio_id` column to be of type `INTEGER`.
+CREATE TABLE studio_snapshots (
+    id INTEGER PRIMARY KEY,
+    studio_id INTEGER NOT NULL REFERENCES studios(id) ON DELETE CASCADE,
+
+    modified_at DATETIME NOT NULL DEFAULT (datetime('now')),
+
+    -- JSON serialized fields
+    context TEXT NOT NULL,
+    messages TEXT NOT NULL
+);
diff --git a/server/bleep/sqlx-data.json b/server/bleep/sqlx-data.json
@@ -1,5 +1,41 @@
 {
   "db": "SQLite",
+  "02ca4d99b13160cb4c78a793f32bec20760e112f4045d8c31541932b4459d7fe": {
+    "describe": {
+      "columns": [
+        {
+          "name": "id",
+          "ordinal": 0,
+          "type_info": "Int64"
+        },
+        {
+          "name": "name",
+          "ordinal": 1,
+          "type_info": "Text"
+        },
+        {
+          "name": "modified_at!",
+          "ordinal": 2,
+          "type_info": "Datetime"
+        },
+        {
+          "name": "context",
+          "ordinal": 3,
+          "type_info": "Text"
+        }
+      ],
+      "nullable": [
+        false,
+        false,
+        false,
+        false
+      ],
+      "parameters": {
+        "Right": 1
+      }
+    },
+    "query": "SELECT\n            s.id,\n            s.name,\n            ss.modified_at as \"modified_at!\",\n            ss.context\n        FROM studios s\n        INNER JOIN studio_snapshots ss ON s.id = ss.studio_id\n        WHERE s.user_id = ? AND (ss.studio_id, ss.modified_at) IN (\n            SELECT studio_id, MAX(modified_at)\n            FROM studio_snapshots\n            GROUP BY studio_id\n        )"
+  },
   "0c72c51f4a5b726f6f524fcb269f074550b1a3a25ed050fb2701f8bec02678d7": {
     "describe": {
       "columns": [],
@@ -92,16 +128,6 @@
     },
     "query": "SELECT id FROM templates WHERE id = ?"
   },
-  "32257914a77039f1317e864d1a5b2c2a309610ccf199c60d5ace2b5e1a2253f1": {
-    "describe": {
-      "columns": [],
-      "nullable": [],
-      "parameters": {
-        "Right": 3
-      }
-    },
-    "query": "INSERT INTO studios(id, name, user_id) VALUES (?, ?, ?)"
-  },
   "359b4d0fa1fcb081767303103b23f0650568cf4e79787c7ddcd21af5bad6761b": {
     "describe": {
       "columns": [],
@@ -146,7 +172,7 @@
         {
           "name": "id",
           "ordinal": 0,
-          "type_info": "Blob"
+          "type_info": "Int64"
         }
       ],
       "nullable": [
@@ -348,6 +374,24 @@
     },
     "query": "SELECT messages, context FROM studio_snapshots WHERE id = ?"
   },
+  "69d67f761ba7fb01ffc310f12cd181acb744b76de90e33c039ce7030c27642bc": {
+    "describe": {
+      "columns": [
+        {
+          "name": "id",
+          "ordinal": 0,
+          "type_info": "Int64"
+        }
+      ],
+      "nullable": [
+        false
+      ],
+      "parameters": {
+        "Right": 2
+      }
+    },
+    "query": "INSERT INTO studios(name, user_id) VALUES (?, ?) RETURNING id"
+  },
   "69e7fcbe274e645ac8dece40ddad39b5f594731e4647a683f4537ede7e20b3d4": {
     "describe": {
       "columns": [],
@@ -462,16 +506,6 @@
     },
     "query": "DELETE FROM templates WHERE id = ? RETURNING id"
   },
-  "9b8ecb6b98f7008b114c31c9f56f68f2fa2c52c48e1c6f3e428f5f65c6a89092": {
-    "describe": {
-      "columns": [],
-      "nullable": [],
-      "parameters": {
-        "Right": 3
-      }
-    },
-    "query": "INSERT INTO studios (id, user_id, name) VALUES (?, ?, ?)"
-  },
   "9f862a56e79cc9ae6e9b896064a0057335b40225be0a8c8d29d9227de12ae364": {
     "describe": {
       "columns": [],
@@ -498,7 +532,7 @@
         {
           "name": "id",
           "ordinal": 0,
-          "type_info": "Blob"
+          "type_info": "Int64"
         }
       ],
       "nullable": [
@@ -516,7 +550,7 @@
         {
           "name": "id",
           "ordinal": 0,
-          "type_info": "Blob"
+          "type_info": "Int64"
         },
         {
           "name": "name",
@@ -638,6 +672,24 @@
     },
     "query": "SELECT messages FROM studio_snapshots WHERE id = ?"
   },
+  "d477bfff9f1880a91e700aad60fdd92141c3c5a6494e37ba965289aff8e9a956": {
+    "describe": {
+      "columns": [
+        {
+          "name": "id",
+          "ordinal": 0,
+          "type_info": "Int64"
+        }
+      ],
+      "nullable": [
+        false
+      ],
+      "parameters": {
+        "Right": 2
+      }
+    },
+    "query": "INSERT INTO studios (user_id, name) VALUES (?, ?) RETURNING id"
+  },
   "d5ee5becde7005920d7094fca5b7974bbf19713b3625fbf6d1a3e198e7cf4de4": {
     "describe": {
       "columns": [
@@ -712,42 +764,6 @@
     },
     "query": "SELECT repo_ref, exchanges FROM conversations WHERE user_id = ? AND thread_id = ?"
   },
-  "e57e6b31c3d2087ac0eb0eee5403156c383f159b0db0359d2c4560ee4c479c1c": {
-    "describe": {
-      "columns": [
-        {
-          "name": "id: Uuid",
-          "ordinal": 0,
-          "type_info": "Blob"
-        },
-        {
-          "name": "name",
-          "ordinal": 1,
-          "type_info": "Text"
-        },
-        {
-          "name": "modified_at!",
-          "ordinal": 2,
-          "type_info": "Datetime"
-        },
-        {
-          "name": "context",
-          "ordinal": 3,
-          "type_info": "Text"
-        }
-      ],
-      "nullable": [
-        false,
-        false,
-        false,
-        false
-      ],
-      "parameters": {
-        "Right": 1
-      }
-    },
-    "query": "SELECT\n            s.id as \"id: Uuid\",\n            s.name,\n            ss.modified_at as \"modified_at!\",\n            ss.context\n        FROM studios s\n        INNER JOIN studio_snapshots ss ON s.id = ss.studio_id\n        WHERE s.user_id = ? AND (ss.studio_id, ss.modified_at) IN (\n            SELECT studio_id, MAX(modified_at)\n            FROM studio_snapshots\n            GROUP BY studio_id\n        )"
-  },
   "ec193a038eb7fc3aaca3c3adebcc4dbde01b47ae34ac2df2227c5e5459617182": {
     "describe": {
       "columns": [],

diff --git a/server/bleep/src/analytics.rs b/server/bleep/src/analytics.rs
@@ -27,7 +27,7 @@ pub struct QueryEvent {
 
 #[derive(Debug, Clone)]
 pub struct StudioEvent {
-    pub studio_id: Uuid,
+    pub studio_id: i64,
 
     // This is not a `Map<K, V>`, to prevent RudderStack from collapsing these fields into columns
     // in the analytics DB.
@@ -36,7 +36,7 @@ pub struct StudioEvent {
 }
 
 impl StudioEvent {
-    pub fn new(studio_id: Uuid, type_: &str) -> Self {
+    pub fn new(studio_id: i64, type_: &str) -> Self {
         Self {
             studio_id,
             type_: type_.to_owned(),

diff --git a/server/bleep/src/webserver/studio.rs b/server/bleep/src/webserver/studio.rs
@@ -41,23 +41,21 @@ pub async fn create(
     user: Extension<User>,
     params: Json<Create>,
 ) -> webserver::Result<String> {
-    let studio_id = Uuid::new_v4();
-
     let mut transaction = app.sql.begin().await?;
 
     let user_id = user
         .login()
         .ok_or_else(|| super::Error::user("didn't have user ID"))?
         .to_string();
 
-    sqlx::query! {
-        "INSERT INTO studios (id, user_id, name) VALUES (?, ?, ?)",
-        studio_id,
+    let studio_id: i64 = sqlx::query! {
+        "INSERT INTO studios (user_id, name) VALUES (?, ?) RETURNING id",
         user_id,
         params.name,
     }
-    .execute(&mut transaction)
-    .await?;
+    .fetch_one(&mut transaction)
+    .await?
+    .id;
 
     sqlx::query! {
         "INSERT INTO studio_snapshots (studio_id, context, messages)
@@ -67,7 +65,8 @@ pub async fn create(
         "[]",
     }
     .execute(&mut transaction)
-    .await?;
+    .await
+    .unwrap();
 
     transaction.commit().await?;
 
@@ -76,7 +75,7 @@ pub async fn create(
 
 #[derive(serde::Serialize)]
 pub struct ListItem {
-    id: Uuid,
+    id: i64,
     name: String,
     modified_at: NaiveDateTime,
     repos: Vec<String>,
@@ -94,7 +93,7 @@ pub async fn list(
 
     let studios = sqlx::query!(
         "SELECT
-            s.id as \"id: Uuid\",
+            s.id,
             s.name,
             ss.modified_at as \"modified_at!\",
             ss.context
@@ -196,11 +195,7 @@ impl From<&Message> for llm_gateway::api::Message {
     }
 }
 
-async fn latest_snapshot_id<'a, E>(
-    studio_id: Uuid,
-    exec: E,
-    user_id: &str,
-) -> webserver::Result<i64>
+async fn latest_snapshot_id<'a, E>(studio_id: i64, exec: E, user_id: &str) -> webserver::Result<i64>
 where
     E: sqlx::Executor<'a, Database = sqlx::Sqlite>,
 {
@@ -228,7 +223,7 @@ pub struct Get {
 pub async fn get(
     app: Extension<Application>,
     user: Extension<User>,
-    Path(id): Path<Uuid>,
+    Path(id): Path<i64>,
     Query(params): Query<Get>,
 ) -> webserver::Result<Json<Studio>> {
     let user_id = user
@@ -280,7 +275,7 @@ pub struct Patch {
 pub async fn patch(
     app: Extension<Application>,
     user: Extension<User>,
-    Path(studio_id): Path<Uuid>,
+    Path(studio_id): Path<i64>,
     Json(patch): Json<Patch>,
 ) -> webserver::Result<Json<TokenCounts>> {
     let user_id = user
@@ -377,7 +372,7 @@ pub async fn patch(
 pub async fn delete(
     app: Extension<Application>,
     user: Extension<User>,
-    Path(id): Path<Uuid>,
+    Path(id): Path<i64>,
 ) -> webserver::Result<()> {
     let user_id = user
         .login()
@@ -546,7 +541,7 @@ pub async fn get_file_token_count(
 pub async fn generate(
     app: Extension<Application>,
     Extension(user): Extension<User>,
-    Path(studio_id): Path<Uuid>,
+    Path(studio_id): Path<i64>,
 ) -> webserver::Result<Sse<Pin<Box<dyn tokio_stream::Stream<Item = Result<sse::Event>> + Send>>>> {
     let user_id = user
         .login()
@@ -706,7 +701,7 @@ async fn generate_llm_context(app: Application, context: Vec<ContextFile>) -> Re
 pub struct Import {
     pub thread_id: Uuid,
     /// An optional studio ID to import into.
-    pub studio_id: Option<Uuid>,
+    pub studio_id: Option<i64>,
 }
 
 async fn extract_relevant_chunks(
@@ -845,16 +840,14 @@ pub async fn import(
     let studio_id = match params.studio_id {
         Some(id) => id,
         None => {
-            let id = Uuid::new_v4();
             sqlx::query!(
-                "INSERT INTO studios(id, name, user_id) VALUES (?, ?, ?)",
-                id,
+                "INSERT INTO studios(name, user_id) VALUES (?, ?) RETURNING id",
                 conversation.title,
                 user_id,
             )
-            .execute(&mut transaction)
-            .await?;
-            id
+            .fetch_one(&mut transaction)
+            .await?
+            .id
         }
     };
 
@@ -952,7 +945,7 @@ pub struct Snapshot {
 pub async fn list_snapshots(
     app: Extension<Application>,
     user: Extension<User>,
-    Path(studio_id): Path<Uuid>,
+    Path(studio_id): Path<i64>,
 ) -> webserver::Result<Json<Vec<Snapshot>>> {
     let user_id = user
         .login()
@@ -987,7 +980,7 @@ pub async fn list_snapshots(
 pub async fn delete_snapshot(
     app: Extension<Application>,
     user: Extension<User>,
-    Path((studio_id, snapshot_id)): Path<(Uuid, i64)>,
+    Path((studio_id, snapshot_id)): Path<(i64, i64)>,
 ) -> webserver::Result<()> {
     let user_id = user
         .login()