feat(qe): enable queries with returning for sqlite (#4640)

prisma · Jan 17, 2024 · ffecfca · ffecfca
1 parent 7f4fc14
commit ffecfca
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 5 deletions.
diff --git a/psl/psl-core/src/builtin_connectors/sqlite_datamodel_connector.rs b/psl/psl-core/src/builtin_connectors/sqlite_datamodel_connector.rs
@@ -25,10 +25,9 @@ const CAPABILITIES: ConnectorCapabilities = enumflags2::make_bitflags!(Connector
     NativeUpsert |
     FilteredInlineChildNestedToOneDisconnect |
     RowIn |
-    // InsertReturning, DeleteReturning, UpdateReturning - While SQLite does support RETURNING, it does not return column information on the
-    // way back from the database. This column type information is necessary in order to preserve consistency for some data types such as int,
-    // where values could overflow.
-    // Since we care to stay consistent with reads, it is not enabled.
+    InsertReturning |
+    DeleteReturning |
+    UpdateReturning |
     SupportsFiltersOnRelationsWithoutJoins
 });
 

diff --git a/quaint/src/visitor/sqlite.rs b/quaint/src/visitor/sqlite.rs
@@ -419,6 +419,41 @@ impl<'a> Visitor<'a> for Sqlite<'a> {
 
         Ok(())
     }
+
+    fn visit_update(&mut self, update: Update<'a>) -> visitor::Result {
+        self.write("UPDATE ")?;
+        self.visit_table(update.table, true)?;
+
+        {
+            self.write(" SET ")?;
+            let pairs = update.columns.into_iter().zip(update.values);
+            let len = pairs.len();
+
+            for (i, (key, value)) in pairs.enumerate() {
+                self.visit_column(key)?;
+                self.write(" = ")?;
+                self.visit_expression(value)?;
+
+                if i < (len - 1) {
+                    self.write(", ")?;
+                }
+            }
+        }
+
+        if let Some(conditions) = update.conditions {
+            self.write(" WHERE ")?;
+            self.visit_conditions(conditions)?;
+        }
+
+        self.returning(update.returning)?;
+
+        if let Some(comment) = update.comment {
+            self.write(" ")?;
+            self.visit_comment(comment)?;
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]

diff --git a/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/metrics.rs b/query-engine/connector-test-kit-rs/query-engine-tests/tests/new/metrics.rs
@@ -34,7 +34,7 @@ mod metrics {
         let total_operations = get_counter(&json, PRISMA_CLIENT_QUERIES_TOTAL);
 
         match runner.connector_version() {
-            Sqlite(_) => assert_eq!(total_queries, 9),
+            Sqlite(_) => assert_eq!(total_queries, 2),
             SqlServer(_) => assert_eq!(total_queries, 17),
             MongoDb(_) => assert_eq!(total_queries, 5),
             CockroachDb(_) => (), // not deterministic