From 1b87b138c03fe2351bf5b9817c7370079312fdf2 Mon Sep 17 00:00:00 2001
From: Eddie A Tejeda <669988+eddietejeda@users.noreply.github.com>
Date: Wed, 23 Sep 2026 23:05:21 -0700
Subject: [PATCH] fix: recognise quoted function names; cap instruction and
 label sizes

`"prompt_jev"(...)` was compared through the name's Display, which keeps
the quotes, so a quoted call fell through to DataFusion as an unknown
function. Match the bare identifier case-insensitively instead.

Instructions over a few thousand characters are rejected by the provider
with a bare HTTP status. Cap instructions at 4,000 characters, labels at
256, and descriptions at 1,024 during validation, so the caller sees the
actual problem at plan time and no request is made.

Found by fuzzing the SQL surface with 2,000 generated queries.
---
 README.md    |  3 ++-
 src/sql.rs   | 34 +++++++++++++++++++++++++++++++++-
 tests/sql.rs | 41 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 434d642..eea6321 100644
--- a/README.md
+++ b/README.md
@@ -154,7 +154,8 @@ each row must be judged in complete isolation.
   applied before inference, so `WHERE region = 'eu' AND prompt_jev(...) > 0.8`
   only asks about the `eu` rows.
 - **Limits.** Each row's text may be up to 64 KiB after JSON encoding. Longer text
-  fails the query. Requests are capped at 256 KiB.
+  fails the query. Requests are capped at 256 KiB. Questions may be up to 4,000
+  characters, labels 256, and descriptions 1,024.
 - **Outages.** Each request is tried three times with a 30-second timeout. If all
   fail, the affected rows return NULL and the query completes. Bad credentials,
   invalid options, or a malformed reply fail the query instead.
diff --git a/src/sql.rs b/src/sql.rs
index f8a2913..54d32c4 100644
--- a/src/sql.rs
+++ b/src/sql.rs
@@ -21,12 +21,25 @@ pub struct Question {
     pub batch_size: usize,
 }
 impl Question {
+    /// Longest instructions accepted, in characters. The provider rejects
+    /// requests well before the crate's request-size cap would, and a plan-time
+    /// error names the actual problem instead of a provider HTTP status.
+    pub const MAX_INSTRUCTION_CHARS: usize = 4_000;
+    pub const MAX_LABEL_CHARS: usize = 256;
+    pub const MAX_DESCRIPTION_CHARS: usize = 1_024;
+
     pub fn validate(&self) -> Result<()> {
         if self.instructions.trim().is_empty() {
             return Err(plan_datafusion_err!(
                 "prompt_jev instructions must not be empty"
             ));
         }
+        if self.instructions.chars().count() > Self::MAX_INSTRUCTION_CHARS {
+            return Err(plan_datafusion_err!(
+                "prompt_jev instructions exceed {} characters",
+                Self::MAX_INSTRUCTION_CHARS
+            ));
+        }
         if !(1..=64).contains(&self.batch_size) {
             return Err(plan_datafusion_err!(
                 "prompt_jev batch_size must be between 1 and 64"
@@ -42,6 +55,17 @@ impl Question {
                     "prompt_jev criteria require unique non-empty labels and non-empty descriptions"
                 ));
             }
+            if c.label.chars().count() > Self::MAX_LABEL_CHARS
+                || c.description
+                    .as_ref()
+                    .is_some_and(|d| d.chars().count() > Self::MAX_DESCRIPTION_CHARS)
+            {
+                return Err(plan_datafusion_err!(
+                    "prompt_jev labels are limited to {} characters and descriptions to {}",
+                    Self::MAX_LABEL_CHARS,
+                    Self::MAX_DESCRIPTION_CHARS
+                ));
+            }
         }
         let valid = match self.kind.as_str() {
             "choice" => (2..=255).contains(&self.criteria.len()),
@@ -122,7 +146,15 @@ pub fn rewrite_statement(statement: &mut ast::Statement) -> Result<()> {
         let Expr::Function(f) = expr else {
             return ControlFlow::Continue(());
         };
-        if f.name.to_string().to_lowercase() != "prompt_jev" {
+        // Match the bare identifier, so `"prompt_jev"(...)` and `PROMPT_JEV(...)`
+        // are the same call; a schema-qualified name is left to DataFusion.
+        let is_prompt_jev = match f.name.0.as_slice() {
+            [part] => part
+                .as_ident()
+                .is_some_and(|i| i.value.eq_ignore_ascii_case("prompt_jev")),
+            _ => false,
+        };
+        if !is_prompt_jev {
             return ControlFlow::Continue(());
         }
         let result = (|| {
diff --git a/tests/sql.rs b/tests/sql.rs
index b55b2a0..9b84ce8 100644
--- a/tests/sql.rs
+++ b/tests/sql.rs
@@ -850,3 +850,44 @@ async fn a_folded_limit_survives_the_filter_split() {
         "rows failing `id > 1` must not be asked about: {asked}"
     );
 }
+
+#[tokio::test]
+async fn quoted_and_uppercase_names_are_the_same_function() {
+    let mock = Arc::new(Mock::default());
+    let result = run(
+        mock.clone(),
+        "SELECT \"prompt_jev\"('a', 'q') AS a, PROMPT_JEV('a', 'q') AS b, \"PROMPT_JEV\"('a', 'q') AS c",
+    )
+    .await
+    .unwrap();
+    assert_eq!(result.iter().map(|b| b.num_rows()).sum::<usize>(), 1);
+    assert!(display(&result).contains("0.9"));
+}
+
+#[tokio::test]
+async fn oversized_instructions_and_labels_fail_at_plan_time() {
+    let long_q = "q".repeat(5_000);
+    let long_label = "L".repeat(300);
+    let long_desc = "d".repeat(2_000);
+    for (sql, needle) in [
+        (
+            format!("SELECT prompt_jev('a', '{long_q}')"),
+            "4000 characters",
+        ),
+        (
+            format!("SELECT prompt_jev('a', 'q', choice := ['{long_label}', 'b'])"),
+            "labels are limited",
+        ),
+        (
+            format!(
+                "SELECT prompt_jev('a', 'q', choice := [{{label: 'a', description: '{long_desc}'}}, 'b'])"
+            ),
+            "descriptions to 1024",
+        ),
+    ] {
+        let mock = Arc::new(Mock::default());
+        let err = run(mock.clone(), &sql).await.err().unwrap().to_string();
+        assert!(err.contains(needle), "{needle}: {err}");
+        assert!(mock.calls.lock().unwrap().is_empty());
+    }
+}