From db1508e27c813dc24e9c846e7d730e4cf0d46db8 Mon Sep 17 00:00:00 2001
From: Damien Elmes <gpg@ankiweb.net>
Date: Fri, 6 Mar 2020 10:11:05 +1000
Subject: [PATCH] support first_row_only

---
 pylib/anki/dbproxy.py        |  3 +--
 pylib/anki/rsbackend.py      |  8 ++++++--
 rslib/src/backend/dbproxy.rs | 38 +++++++++++++++++++++++++++++++++++-
 3 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/pylib/anki/dbproxy.py b/pylib/anki/dbproxy.py
index 80cb3795d..197e5ef21 100644
--- a/pylib/anki/dbproxy.py
+++ b/pylib/anki/dbproxy.py
@@ -67,8 +67,7 @@ class DBProxy:
                 self.mod = True
         assert ":" not in sql
         # fetch rows
-        # fixme: first_row_only
-        return self._backend.db_query(sql, args)
+        return self._backend.db_query(sql, args, first_row_only)
 
     # Query shortcuts
     ###################
diff --git a/pylib/anki/rsbackend.py b/pylib/anki/rsbackend.py
index db427b4a2..b098aaa10 100644
--- a/pylib/anki/rsbackend.py
+++ b/pylib/anki/rsbackend.py
@@ -386,8 +386,12 @@ class RustBackend:
     def restore_trash(self):
         self._run_command(pb.BackendInput(restore_trash=pb.Empty()))
 
-    def db_query(self, sql: str, args: List[ValueForDB]) -> List[DBRow]:
-        return self._db_command(dict(kind="query", sql=sql, args=args))
+    def db_query(
+        self, sql: str, args: List[ValueForDB], first_row_only: bool
+    ) -> List[DBRow]:
+        return self._db_command(
+            dict(kind="query", sql=sql, args=args, first_row_only=first_row_only)
+        )
 
     def db_execute_many(self, sql: str, args: List[List[ValueForDB]]) -> List[DBRow]:
         return self._db_command(dict(kind="executemany", sql=sql, args=args))
diff --git a/rslib/src/backend/dbproxy.rs b/rslib/src/backend/dbproxy.rs
index f4ebaf8b6..aa4cb173c 100644
--- a/rslib/src/backend/dbproxy.rs
+++ b/rslib/src/backend/dbproxy.rs
@@ -4,6 +4,7 @@
 use crate::err::Result;
 use crate::storage::StorageContext;
 use rusqlite::types::{FromSql, FromSqlError, ToSql, ToSqlOutput, ValueRef};
+use rusqlite::OptionalExtension;
 use serde_derive::{Deserialize, Serialize};
 
 #[derive(Deserialize)]
@@ -12,6 +13,7 @@ pub(super) enum DBRequest {
     Query {
         sql: String,
         args: Vec<SqlValue>,
+        first_row_only: bool,
     },
     Begin,
     Commit,
@@ -68,7 +70,17 @@ impl FromSql for SqlValue {
 pub(super) fn db_command_bytes(ctx: &StorageContext, input: &[u8]) -> Result<String> {
     let req: DBRequest = serde_json::from_slice(input)?;
     let resp = match req {
-        DBRequest::Query { sql, args } => db_query(ctx, &sql, &args)?,
+        DBRequest::Query {
+            sql,
+            args,
+            first_row_only,
+        } => {
+            if first_row_only {
+                db_query_row(ctx, &sql, &args)?
+            } else {
+                db_query(ctx, &sql, &args)?
+            }
+        }
         DBRequest::Begin => {
             ctx.begin_trx()?;
             DBResult::None
@@ -86,6 +98,30 @@ pub(super) fn db_command_bytes(ctx: &StorageContext, input: &[u8]) -> Result<Str
     Ok(serde_json::to_string(&resp)?)
 }
 
+pub(super) fn db_query_row(ctx: &StorageContext, sql: &str, args: &[SqlValue]) -> Result<DBResult> {
+    let mut stmt = ctx.db.prepare_cached(sql)?;
+    let columns = stmt.column_count();
+
+    let row = stmt
+        .query_row(args, |row| {
+            let mut orow = Vec::with_capacity(columns);
+            for i in 0..columns {
+                let v: SqlValue = row.get(i)?;
+                orow.push(v);
+            }
+            Ok(orow)
+        })
+        .optional()?;
+
+    let rows = if let Some(row) = row {
+        vec![row]
+    } else {
+        vec![]
+    };
+
+    Ok(DBResult::Rows(rows))
+}
+
 pub(super) fn db_query(ctx: &StorageContext, sql: &str, args: &[SqlValue]) -> Result<DBResult> {
     let mut stmt = ctx.db.prepare_cached(sql)?;
     let columns = stmt.column_count();