From 7c03af09076e9f816daaec816b889b627f636b3e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philip=20Dub=C3=A9?= <philip@peerdb.io>
Date: Wed, 18 Sep 2024 15:22:13 +0000
Subject: [PATCH] Query parser: encode ?? as ?

Also borrow from source template instead of copying its parts into separate strings
---
 src/lib.rs     |  2 +-
 src/query.rs   |  8 ++++----
 src/sql/mod.rs | 56 +++++++++++++++++++++++++++-----------------------
 3 files changed, 35 insertions(+), 31 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 692211fd..492bb486 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -299,7 +299,7 @@ impl Client {
     }
 
     /// Starts a new SELECT/DDL query.
-    pub fn query(&self, query: &str) -> query::Query {
+    pub fn query<'a>(&self, query: &'a str) -> query::Query<'a> {
         query::Query::new(self, query)
     }
 
diff --git a/src/query.rs b/src/query.rs
index 5a03d178..14b35327 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -17,13 +17,13 @@ const MAX_QUERY_LEN_TO_USE_GET: usize = 8192;
 
 #[must_use]
 #[derive(Clone)]
-pub struct Query {
+pub struct Query<'a> {
     client: Client,
-    sql: SqlBuilder,
+    sql: SqlBuilder<'a>,
 }
 
-impl Query {
-    pub(crate) fn new(client: &Client, template: &str) -> Self {
+impl<'parts> Query<'parts> {
+    pub(crate) fn new(client: &Client, template: &'parts str) -> Self {
         Self {
             client: client.clone(),
             sql: SqlBuilder::new(template),
diff --git a/src/sql/mod.rs b/src/sql/mod.rs
index edb42e71..8ed0290d 100644
--- a/src/sql/mod.rs
+++ b/src/sql/mod.rs
@@ -12,34 +12,41 @@ pub(crate) mod escape;
 mod ser;
 
 #[derive(Clone)]
-pub(crate) enum SqlBuilder {
-    InProgress(Vec<Part>),
+pub(crate) enum SqlBuilder<'a> {
+    InProgress(Vec<Part<'a>>),
     Failed(String),
 }
 
 #[derive(Clone)]
-pub(crate) enum Part {
+pub(crate) enum Part<'a> {
     Arg,
     Fields,
-    Text(String),
+    Str(&'a str),
+    String(String),
 }
 
-impl SqlBuilder {
-    pub(crate) fn new(template: &str) -> Self {
-        let mut iter = template.split('?');
-        let prefix = String::from(iter.next().unwrap());
-        let mut parts = vec![Part::Text(prefix)];
-
-        for s in iter {
-            let text = if let Some(text) = s.strip_prefix("fields") {
+impl<'parts> SqlBuilder<'parts> {
+    pub(crate) fn new(template: &'parts str) -> Self {
+        let mut parts = Vec::new();
+        let mut rest = template;
+        while let Some(idx) = rest.find('?') {
+            if idx != 0 {
+                parts.push(Part::Str(&rest[..idx]));
+            }
+            rest = &rest[idx+1..];
+            if let Some(restqq) = rest.strip_prefix('?') {
+                parts.push(Part::Str("?"));
+                rest = restqq;
+            } else if let Some(restfields) = rest.strip_prefix("fields") {
                 parts.push(Part::Fields);
-                text
+                rest = restfields;
             } else {
                 parts.push(Part::Arg);
-                s
-            };
-
-            parts.push(Part::Text(text.into()));
+            }
+        }
+        
+        if !rest.is_empty() {
+            parts.push(Part::Str(rest));
         }
 
         SqlBuilder::InProgress(parts)
@@ -57,7 +64,7 @@ impl SqlBuilder {
                 return self.error(format_args!("invalid argument: {err}"));
             }
 
-            *part = Part::Text(s);
+            *part = Part::String(s);
         } else {
             self.error("unexpected bind(), all arguments are already bound");
         }
@@ -70,23 +77,19 @@ impl SqlBuilder {
 
         if let Some(fields) = row::join_column_names::<T>() {
             for part in parts.iter_mut().filter(|p| matches!(p, Part::Fields)) {
-                *part = Part::Text(fields.clone());
+                *part = Part::String(fields.clone());
             }
         } else if parts.iter().any(|p| matches!(p, Part::Fields)) {
             self.error("argument ?fields cannot be used with non-struct row types");
         }
     }
 
-    pub(crate) fn append(&mut self, suffix: &str) {
+    pub(crate) fn append(&mut self, part: &'static str) {
         let Self::InProgress(parts) = self else {
             return;
         };
 
-        if let Some(Part::Text(text)) = parts.last_mut() {
-            text.push_str(suffix);
-        } else {
-            // Do nothing, it will fail in `finish()`.
-        }
+        parts.push(Part::Str(part));
     }
 
     pub(crate) fn finish(mut self) -> Result<String> {
@@ -95,7 +98,8 @@ impl SqlBuilder {
         if let Self::InProgress(parts) = &self {
             for part in parts {
                 match part {
-                    Part::Text(text) => sql.push_str(text),
+                    Part::Str(text) => sql.push_str(text),
+                    Part::String(text) => sql.push_str(&text[..]),
                     Part::Arg => {
                         self.error("unbound query argument");
                         break;