use anyhow::{Context as _, Result};use flate2::read::GzDecoder;use gpui::BackgroundExecutor;use http_client::{AsyncBody, HttpClient, Method, Request};use indoc::indoc;use serde::Deserialize;use serde_json::{Value as JsonValue, json};use std::fmt::Write as _;use std::io::Read;use std::sync::Arc;use std::time::Duration;use telemetry_events::EditPredictionRating; use zeta_prompt::{ZetaFormat, ZetaPromptInput, excerpt_range_for_format}; use crate::example::Example;use crate::progress::{InfoStyle, Progress, Step};const EDIT_PREDICTION_DEPLOYMENT_EVENT: &str = "Edit Prediction Deployment";use edit_prediction::example_spec::{ExampleSpec, TelemetrySource}; pub(crate) const SNOWFLAKE_SUCCESS_CODE: &str = "090001";pub(crate) const SNOWFLAKE_ASYNC_IN_PROGRESS_CODE: &str = "333334";const PREDICTIVE_EDIT_REQUESTED_EVENT: &str = "Predictive Edit Requested";const PREDICTIVE_EDIT_REJECTED_EVENT: &str = "Predictive Edit Rejected";const EDIT_PREDICTION_RATED_EVENT: &str = "Edit Prediction Rated";const EDIT_PREDICTION_SETTLED_EVENT: &str = "Edit Prediction Settled"; /// Minimum Zed version for filtering captured examples./// For example, `MinCaptureVersion { minor: 224, patch: 1 }` means only pull examples/// where `zed_version >= 0.224.1`.#[derive(Clone, Copy, Debug)]pub struct MinCaptureVersion { pub minor: u32, pub patch: u32,} const DEFAULT_STATEMENT_TIMEOUT_SECONDS: u64 = 240;const SETTLED_STATEMENT_TIMEOUT_SECONDS: u64 = 240;pub(crate) const POLL_INTERVAL: Duration = Duration::from_secs(2);pub(crate) const MAX_POLL_ATTEMPTS: usize = 120; /// Parse an input token of the form `captured-after:{timestamp}`.pub fn parse_captured_after_input(input: &str) -> Option<&str> { input.strip_prefix("captured-after:")} /// Parse an input token of the form `rejected-after:{timestamp}`.pub fn parse_rejected_after_input(input: &str) -> Option<&str> { input.strip_prefix("rejected-after:")} /// Parse an input token of the form `requested-after:{timestamp}`.pub fn parse_requested_after_input(input: &str) -> Option<&str> { input.strip_prefix("requested-after:")} /// Parse an input token of the form `settled-after:{timestamp}`.pub fn parse_settled_after_input(input: &str) -> Option<&str> { input.strip_prefix("settled-after:")} /// Parse an input token of the form `rated-after:{timestamp}`, `rated-positive-after:{timestamp}`,/// or `rated-negative-after:{timestamp}`./// Returns `(timestamp, Option<EditPredictionRating>)` where `None` means all ratings.pub fn parse_rated_after_input(input: &str) -> Option<(&str, Option<EditPredictionRating>)> { if let Some(timestamp) = input.strip_prefix("rated-positive-after:") { Some((timestamp, Some(EditPredictionRating::Positive))) } else if let Some(timestamp) = input.strip_prefix("rated-negative-after:") { Some((timestamp, Some(EditPredictionRating::Negative))) } else if let Some(timestamp) = input.strip_prefix("rated-after:") { Some((timestamp, None)) } else { None }} #[derive(Debug, Clone, Deserialize)]#[serde(rename_all = "camelCase")]pub(crate) struct SnowflakeStatementResponse { #[serde(default)] pub(crate) data: Vec<Vec<JsonValue>>, #[serde(default)] pub(crate) result_set_meta_data: Option<SnowflakeResultSetMetaData>, #[serde(default)] pub(crate) code: Option<String>, #[serde(default)] pub(crate) message: Option<String>, #[serde(default)] pub(crate) statement_handle: Option<String>,} #[derive(Debug, Clone, Deserialize)]#[serde(rename_all = "camelCase")]pub(crate) struct SnowflakeResultSetMetaData { #[serde(default, rename = "rowType")] row_type: Vec<SnowflakeColumnMeta>, #[serde(default)] num_rows: Option<i64>, #[serde(default)] partition_info: Vec<SnowflakePartitionInfo>,} #[derive(Debug, Clone, Deserialize)]#[serde(rename_all = "camelCase")]struct SnowflakePartitionInfo {} #[derive(Debug, Clone, Deserialize)]struct SnowflakeColumnMeta { #[serde(default)] name: String,} async fn run_sql_with_polling( http_client: Arc<dyn HttpClient>, base_url: &str, token: &str, request: &serde_json::Value, step_progress: &crate::progress::StepProgress, background_executor: BackgroundExecutor,) -> Result<SnowflakeStatementResponse> { let mut response = run_sql(http_client.clone(), base_url, token, request).await?; if response.code.as_deref() == Some(SNOWFLAKE_ASYNC_IN_PROGRESS_CODE) { let statement_handle = response .statement_handle .as_ref() .context("async query response missing statementHandle")? .clone(); for attempt in 1..=MAX_POLL_ATTEMPTS { step_progress.set_substatus(format!("polling ({attempt})")); background_executor.timer(POLL_INTERVAL).await; response = fetch_partition(http_client.clone(), base_url, token, &statement_handle, 0).await?; if response.code.as_deref() != Some(SNOWFLAKE_ASYNC_IN_PROGRESS_CODE) { break; } } if response.code.as_deref() == Some(SNOWFLAKE_ASYNC_IN_PROGRESS_CODE) { anyhow::bail!( "query still running after {} poll attempts ({} seconds)", MAX_POLL_ATTEMPTS, MAX_POLL_ATTEMPTS as u64 * POLL_INTERVAL.as_secs() ); } } Ok(response)} struct SnowflakeConfig { token: String, base_url: String, role: Option<String>,} async fn fetch_examples_with_query( http_client: Arc<dyn HttpClient>, step_progress: &crate::progress::StepProgress, background_executor: BackgroundExecutor, statement: &str, bindings: JsonValue, timeout_seconds: u64, required_columns: &[&str], parse_response: for<'a> fn( &'a SnowflakeStatementResponse, &'a std::collections::HashMap<String, usize>, ) -> Result<Box<dyn Iterator<Item = Example> + 'a>>,) -> Result<Vec<Example>> { let snowflake = SnowflakeConfig { token: std::env::var("EP_SNOWFLAKE_API_KEY") .context("missing required environment variable EP_SNOWFLAKE_API_KEY")?, base_url: std::env::var("EP_SNOWFLAKE_BASE_URL").context( "missing required environment variable EP_SNOWFLAKE_BASE_URL (e.g. https://<account>.snowflakecomputing.com)", )?, role: std::env::var("EP_SNOWFLAKE_ROLE").ok(), }; let request = json!({ "statement": statement, "timeout": timeout_seconds, "database": "EVENTS", "schema": "PUBLIC", "warehouse": "DBT", "role": snowflake.role.as_deref(), "bindings": bindings }); let response = run_sql_with_polling( http_client.clone(), &snowflake.base_url, &snowflake.token, &request, step_progress, background_executor, ) .await?; let total_rows = response .result_set_meta_data .as_ref() .and_then(|meta| meta.num_rows) .unwrap_or(response.data.len() as i64); let partition_count = response .result_set_meta_data .as_ref() .map(|meta| meta.partition_info.len()) .unwrap_or(1) .max(1); step_progress.set_info(format!("{} rows", total_rows), InfoStyle::Normal); step_progress.set_substatus("parsing"); let column_indices = get_column_indices(&response.result_set_meta_data, required_columns); let mut parsed_examples = Vec::with_capacity(total_rows as usize); parsed_examples.extend(parse_response(&response, &column_indices)?); if partition_count > 1 { let statement_handle = response .statement_handle .as_ref() .context("response has multiple partitions but no statementHandle")?; for partition in 1..partition_count { step_progress.set_substatus(format!( "fetching partition {}/{}", partition + 1, partition_count )); let partition_response = fetch_partition( http_client.clone(), &snowflake.base_url, &snowflake.token, statement_handle, partition, ) .await?; parsed_examples.extend(parse_response(&partition_response, &column_indices)?); } } step_progress.set_substatus("done"); Ok(parsed_examples)} pub(crate) async fn fetch_partition( http_client: Arc<dyn HttpClient>, base_url: &str, token: &str, statement_handle: &str, partition: usize,) -> Result<SnowflakeStatementResponse> { let url = format!( "{}/api/v2/statements/{}?partition={}", base_url.trim_end_matches('/'), statement_handle, partition ); let http_request = Request::builder() .method(Method::GET) .uri(url.as_str()) .header("Authorization", format!("Bearer {token}")) .header( "X-Snowflake-Authorization-Token-Type", "PROGRAMMATIC_ACCESS_TOKEN", ) .header("Accept", "application/json") .header("Accept-Encoding", "gzip") .header("User-Agent", "edit_prediction_cli") .body(AsyncBody::empty())?; let response = http_client .send(http_request) .await .context("failed to send partition request to Snowflake SQL API")?; let status = response.status(); let content_encoding = response .headers() .get("content-encoding") .and_then(|v| v.to_str().ok()) .map(|s| s.to_lowercase()); let body_bytes = { use futures::AsyncReadExt as _; let mut body = response.into_body(); let mut bytes = Vec::new(); body.read_to_end(&mut bytes) .await .context("failed to read Snowflake SQL API partition response body")?; bytes }; let body_bytes = if content_encoding.as_deref() == Some("gzip") { let mut decoder = GzDecoder::new(&body_bytes[..]); let mut decompressed = Vec::new(); decoder .read_to_end(&mut decompressed) .context("failed to decompress gzip response")?; decompressed } else { body_bytes }; if !status.is_success() && status.as_u16() != 202 { let body_text = String::from_utf8_lossy(&body_bytes); anyhow::bail!( "snowflake sql api partition request http {}: {}", status.as_u16(), body_text ); } if body_bytes.is_empty() { anyhow::bail!( "snowflake sql api partition {} returned empty response body (http {})", partition, status.as_u16() ); } serde_json::from_slice::<SnowflakeStatementResponse>(&body_bytes).with_context(|| { let body_preview = String::from_utf8_lossy(&body_bytes[..body_bytes.len().min(500)]); format!( "failed to parse Snowflake SQL API partition {} response JSON (http {}): {}", partition, status.as_u16(), body_preview ) })} pub(crate) async fn run_sql( http_client: Arc<dyn HttpClient>, base_url: &str, token: &str, request: &serde_json::Value,) -> Result<SnowflakeStatementResponse> { let url = format!("{}/api/v2/statements", base_url.trim_end_matches('/')); let request_body = serde_json::to_vec(request).context("failed to serialize Snowflake SQL API request")?; let http_request = Request::builder() .method(Method::POST) .uri(url.as_str()) .header("Authorization", format!("Bearer {token}")) .header( "X-Snowflake-Authorization-Token-Type", "PROGRAMMATIC_ACCESS_TOKEN", ) .header("Content-Type", "application/json") .header("Accept", "application/json") .header("User-Agent", "edit_prediction_cli") .body(AsyncBody::from(request_body.clone()))?; let response = http_client .send(http_request) .await .context("failed to send request to Snowflake SQL API")?; let status = response.status(); let body_bytes = { use futures::AsyncReadExt as _; let mut body = response.into_body(); let mut bytes = Vec::new(); body.read_to_end(&mut bytes) .await .context("failed to read Snowflake SQL API response body")?; bytes }; if !status.is_success() && status.as_u16() != 202 { let body_text = String::from_utf8_lossy(&body_bytes); anyhow::bail!("snowflake sql api http {}: {}", status.as_u16(), body_text); } serde_json::from_slice::<SnowflakeStatementResponse>(&body_bytes) .context("failed to parse Snowflake SQL API response JSON")} pub async fn fetch_rejected_examples_after( http_client: Arc<dyn HttpClient>, after_timestamps: &[String], max_rows_per_timestamp: usize, offset: usize, background_executor: BackgroundExecutor, min_capture_version: Option<MinCaptureVersion>,) -> Result<Vec<Example>> { if after_timestamps.is_empty() { return Ok(Vec::new()); } let progress = Progress::global(); let mut all_examples = Vec::new(); for after_date in after_timestamps.iter() { let step_progress_name = format!("rejected>{after_date}"); let step_progress = progress.start(Step::PullExamples, &step_progress_name); step_progress.set_substatus("querying"); let min_minor_str = min_capture_version.map(|version| version.minor.to_string()); let min_patch_str = min_capture_version.map(|version| version.patch.to_string()); let min_minor_str_ref = min_minor_str.as_deref(); let min_patch_str_ref = min_patch_str.as_deref(); let statement = indoc! {r#" SELECT req.event_properties:request_id::string AS request_id, req.device_id::string AS device_id, req.time::string AS time, req.event_properties:input AS input, req.event_properties:prompt::string AS prompt, req.event_properties:output::string AS output, rej.event_properties:was_shown::boolean AS was_shown, rej.event_properties:reason::string AS reason, req.event_properties:zed_version::string AS zed_version FROM events req INNER JOIN events rej ON req.event_properties:request_id = rej.event_properties:request_id WHERE req.event_type = ? AND rej.event_type = ? AND req.event_properties:version = 'V3' AND rej.event_properties:was_shown = true AND req.event_properties:input:can_collect_data = true AND req.time > TRY_TO_TIMESTAMP_NTZ(?) AND (? IS NULL OR ( TRY_CAST(SPLIT_PART(req.event_properties:zed_version::string, '.', 2) AS INTEGER) > ? OR ( TRY_CAST(SPLIT_PART(req.event_properties:zed_version::string, '.', 2) AS INTEGER) = ? AND TRY_CAST(SPLIT_PART(SPLIT_PART(req.event_properties:zed_version::string, '.', 3), '+', 1) AS INTEGER) >= ? ) )) ORDER BY req.time ASC LIMIT ? OFFSET ? "#}; let bindings = json!({ "1": { "type": "TEXT", "value": PREDICTIVE_EDIT_REQUESTED_EVENT }, "2": { "type": "TEXT", "value": PREDICTIVE_EDIT_REJECTED_EVENT }, "3": { "type": "TEXT", "value": after_date }, "4": { "type": "FIXED", "value": min_minor_str_ref }, "5": { "type": "FIXED", "value": min_minor_str_ref }, "6": { "type": "FIXED", "value": min_minor_str_ref }, "7": { "type": "FIXED", "value": min_patch_str_ref }, "8": { "type": "FIXED", "value": max_rows_per_timestamp.to_string() }, "9": { "type": "FIXED", "value": offset.to_string() } }); let examples = fetch_examples_with_query( http_client.clone(), &step_progress, background_executor.clone(), statement, bindings, DEFAULT_STATEMENT_TIMEOUT_SECONDS, &[ "request_id", "device_id", "time", "input", "prompt", "output", "was_shown", "reason", "zed_version", ], rejected_examples_from_response, ) .await?; all_examples.extend(examples); } Ok(all_examples)} pub async fn fetch_requested_examples_after( http_client: Arc<dyn HttpClient>, after_timestamps: &[String], max_rows_per_timestamp: usize, offset: usize, background_executor: BackgroundExecutor, min_capture_version: Option<MinCaptureVersion>,) -> Result<Vec<Example>> { if after_timestamps.is_empty() { return Ok(Vec::new()); } Lines 1–500 of 1653. Load more
use anyhow::{Context as _, Result}; use flate2::read::GzDecoder; use gpui::BackgroundExecutor; use http_client::{AsyncBody, HttpClient, Method, Request}; use indoc::indoc; use serde::Deserialize; use serde_json::{Value as JsonValue, json}; use std::fmt::Write as _; use std::io::Read; use std::sync::Arc; use std::time::Duration; use telemetry_events::EditPredictionRating; use zeta_prompt::{ZetaFormat, ZetaPromptInput, excerpt_range_for_format}; use crate::example::Example; use crate::progress::{InfoStyle, Progress, Step}; const EDIT_PREDICTION_DEPLOYMENT_EVENT: &str = "Edit Prediction Deployment"; use edit_prediction::example_spec::{ExampleSpec, TelemetrySource}; pub(crate) const SNOWFLAKE_SUCCESS_CODE: &str = "090001"; pub(crate) const SNOWFLAKE_ASYNC_IN_PROGRESS_CODE: &str = "333334"; const PREDICTIVE_EDIT_REQUESTED_EVENT: &str = "Predictive Edit Requested"; const PREDICTIVE_EDIT_REJECTED_EVENT: &str = "Predictive Edit Rejected"; const EDIT_PREDICTION_RATED_EVENT: &str = "Edit Prediction Rated"; const EDIT_PREDICTION_SETTLED_EVENT: &str = "Edit Prediction Settled"; /// Minimum Zed version for filtering captured examples. /// For example, `MinCaptureVersion { minor: 224, patch: 1 }` means only pull examples /// where `zed_version >= 0.224.1`. #[derive(Clone, Copy, Debug)] pub struct MinCaptureVersion { pub minor: u32, pub patch: u32, } const DEFAULT_STATEMENT_TIMEOUT_SECONDS: u64 = 240; const SETTLED_STATEMENT_TIMEOUT_SECONDS: u64 = 240; pub(crate) const POLL_INTERVAL: Duration = Duration::from_secs(2); pub(crate) const MAX_POLL_ATTEMPTS: usize = 120; /// Parse an input token of the form `captured-after:{timestamp}`. pub fn parse_captured_after_input(input: &str) -> Option<&str> { input.strip_prefix("captured-after:") } /// Parse an input token of the form `rejected-after:{timestamp}`. pub fn parse_rejected_after_input(input: &str) -> Option<&str> { input.strip_prefix("rejected-after:") } /// Parse an input token of the form `requested-after:{timestamp}`. pub fn parse_requested_after_input(input: &str) -> Option<&str> { input.strip_prefix("requested-after:") } /// Parse an input token of the form `settled-after:{timestamp}`. pub fn parse_settled_after_input(input: &str) -> Option<&str> { input.strip_prefix("settled-after:") } /// Parse an input token of the form `rated-after:{timestamp}`, `rated-positive-after:{timestamp}`, /// or `rated-negative-after:{timestamp}`. /// Returns `(timestamp, Option<EditPredictionRating>)` where `None` means all ratings. pub fn parse_rated_after_input(input: &str) -> Option<(&str, Option<EditPredictionRating>)> { if let Some(timestamp) = input.strip_prefix("rated-positive-after:") { Some((timestamp, Some(EditPredictionRating::Positive))) } else if let Some(timestamp) = input.strip_prefix("rated-negative-after:") { Some((timestamp, Some(EditPredictionRating::Negative))) } else if let Some(timestamp) = input.strip_prefix("rated-after:") { Some((timestamp, None)) } else { None } } #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase")] pub(crate) struct SnowflakeStatementResponse { #[serde(default)] pub(crate) data: Vec<Vec<JsonValue>>, #[serde(default)] pub(crate) result_set_meta_data: Option<SnowflakeResultSetMetaData>, #[serde(default)] pub(crate) code: Option<String>, #[serde(default)] pub(crate) message: Option<String>, #[serde(default)] pub(crate) statement_handle: Option<String>, } #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase")] pub(crate) struct SnowflakeResultSetMetaData { #[serde(default, rename = "rowType")] row_type: Vec<SnowflakeColumnMeta>, #[serde(default)] num_rows: Option<i64>, #[serde(default)] partition_info: Vec<SnowflakePartitionInfo>, } #[derive(Debug, Clone, Deserialize)] #[serde(rename_all = "camelCase")] struct SnowflakePartitionInfo {} #[derive(Debug, Clone, Deserialize)] struct SnowflakeColumnMeta { #[serde(default)] name: String, } async fn run_sql_with_polling( http_client: Arc<dyn HttpClient>, base_url: &str, token: &str, request: &serde_json::Value, step_progress: &crate::progress::StepProgress, background_executor: BackgroundExecutor, ) -> Result<SnowflakeStatementResponse> { let mut response = run_sql(http_client.clone(), base_url, token, request).await?; if response.code.as_deref() == Some(SNOWFLAKE_ASYNC_IN_PROGRESS_CODE) { let statement_handle = response .statement_handle .as_ref() .context("async query response missing statementHandle")? .clone(); for attempt in 1..=MAX_POLL_ATTEMPTS { step_progress.set_substatus(format!("polling ({attempt})")); background_executor.timer(POLL_INTERVAL).await; response = fetch_partition(http_client.clone(), base_url, token, &statement_handle, 0).await?; if response.code.as_deref() != Some(SNOWFLAKE_ASYNC_IN_PROGRESS_CODE) { break; } } if response.code.as_deref() == Some(SNOWFLAKE_ASYNC_IN_PROGRESS_CODE) { anyhow::bail!( "query still running after {} poll attempts ({} seconds)", MAX_POLL_ATTEMPTS, MAX_POLL_ATTEMPTS as u64 * POLL_INTERVAL.as_secs() ); } } Ok(response) } struct SnowflakeConfig { token: String, base_url: String, role: Option<String>, } async fn fetch_examples_with_query( http_client: Arc<dyn HttpClient>, step_progress: &crate::progress::StepProgress, background_executor: BackgroundExecutor, statement: &str, bindings: JsonValue, timeout_seconds: u64, required_columns: &[&str], parse_response: for<'a> fn( &'a SnowflakeStatementResponse, &'a std::collections::HashMap<String, usize>, ) -> Result<Box<dyn Iterator<Item = Example> + 'a>>, ) -> Result<Vec<Example>> { let snowflake = SnowflakeConfig { token: std::env::var("EP_SNOWFLAKE_API_KEY") .context("missing required environment variable EP_SNOWFLAKE_API_KEY")?, base_url: std::env::var("EP_SNOWFLAKE_BASE_URL").context( "missing required environment variable EP_SNOWFLAKE_BASE_URL (e.g. https://<account>.snowflakecomputing.com)", )?, role: std::env::var("EP_SNOWFLAKE_ROLE").ok(), }; let request = json!({ "statement": statement, "timeout": timeout_seconds, "database": "EVENTS", "schema": "PUBLIC", "warehouse": "DBT", "role": snowflake.role.as_deref(), "bindings": bindings }); let response = run_sql_with_polling( http_client.clone(), &snowflake.base_url, &snowflake.token, &request, step_progress, background_executor, ) .await?; let total_rows = response .result_set_meta_data .as_ref() .and_then(|meta| meta.num_rows) .unwrap_or(response.data.len() as i64); let partition_count = response .result_set_meta_data .as_ref() .map(|meta| meta.partition_info.len()) .unwrap_or(1) .max(1); step_progress.set_info(format!("{} rows", total_rows), InfoStyle::Normal); step_progress.set_substatus("parsing"); let column_indices = get_column_indices(&response.result_set_meta_data, required_columns); let mut parsed_examples = Vec::with_capacity(total_rows as usize); parsed_examples.extend(parse_response(&response, &column_indices)?); if partition_count > 1 { let statement_handle = response .statement_handle .as_ref() .context("response has multiple partitions but no statementHandle")?; for partition in 1..partition_count { step_progress.set_substatus(format!( "fetching partition {}/{}", partition + 1, partition_count )); let partition_response = fetch_partition( http_client.clone(), &snowflake.base_url, &snowflake.token, statement_handle, partition, ) .await?; parsed_examples.extend(parse_response(&partition_response, &column_indices)?); } } step_progress.set_substatus("done"); Ok(parsed_examples) } pub(crate) async fn fetch_partition( http_client: Arc<dyn HttpClient>, base_url: &str, token: &str, statement_handle: &str, partition: usize, ) -> Result<SnowflakeStatementResponse> { let url = format!( "{}/api/v2/statements/{}?partition={}", base_url.trim_end_matches('/'), statement_handle, partition ); let http_request = Request::builder() .method(Method::GET) .uri(url.as_str()) .header("Authorization", format!("Bearer {token}")) .header( "X-Snowflake-Authorization-Token-Type", "PROGRAMMATIC_ACCESS_TOKEN", ) .header("Accept", "application/json") .header("Accept-Encoding", "gzip") .header("User-Agent", "edit_prediction_cli") .body(AsyncBody::empty())?; let response = http_client .send(http_request) .await .context("failed to send partition request to Snowflake SQL API")?; let status = response.status(); let content_encoding = response .headers() .get("content-encoding") .and_then(|v| v.to_str().ok()) .map(|s| s.to_lowercase()); let body_bytes = { use futures::AsyncReadExt as _; let mut body = response.into_body(); let mut bytes = Vec::new(); body.read_to_end(&mut bytes) .await .context("failed to read Snowflake SQL API partition response body")?; bytes }; let body_bytes = if content_encoding.as_deref() == Some("gzip") { let mut decoder = GzDecoder::new(&body_bytes[..]); let mut decompressed = Vec::new(); decoder .read_to_end(&mut decompressed) .context("failed to decompress gzip response")?; decompressed } else { body_bytes }; if !status.is_success() && status.as_u16() != 202 { let body_text = String::from_utf8_lossy(&body_bytes); anyhow::bail!( "snowflake sql api partition request http {}: {}", status.as_u16(), body_text ); } if body_bytes.is_empty() { anyhow::bail!( "snowflake sql api partition {} returned empty response body (http {})", partition, status.as_u16() ); } serde_json::from_slice::<SnowflakeStatementResponse>(&body_bytes).with_context(|| { let body_preview = String::from_utf8_lossy(&body_bytes[..body_bytes.len().min(500)]); format!( "failed to parse Snowflake SQL API partition {} response JSON (http {}): {}", partition, status.as_u16(), body_preview ) }) } pub(crate) async fn run_sql( http_client: Arc<dyn HttpClient>, base_url: &str, token: &str, request: &serde_json::Value, ) -> Result<SnowflakeStatementResponse> { let url = format!("{}/api/v2/statements", base_url.trim_end_matches('/')); let request_body = serde_json::to_vec(request).context("failed to serialize Snowflake SQL API request")?; let http_request = Request::builder() .method(Method::POST) .uri(url.as_str()) .header("Authorization", format!("Bearer {token}")) .header( "X-Snowflake-Authorization-Token-Type", "PROGRAMMATIC_ACCESS_TOKEN", ) .header("Content-Type", "application/json") .header("Accept", "application/json") .header("User-Agent", "edit_prediction_cli") .body(AsyncBody::from(request_body.clone()))?; let response = http_client .send(http_request) .await .context("failed to send request to Snowflake SQL API")?; let status = response.status(); let body_bytes = { use futures::AsyncReadExt as _; let mut body = response.into_body(); let mut bytes = Vec::new(); body.read_to_end(&mut bytes) .await .context("failed to read Snowflake SQL API response body")?; bytes }; if !status.is_success() && status.as_u16() != 202 { let body_text = String::from_utf8_lossy(&body_bytes); anyhow::bail!("snowflake sql api http {}: {}", status.as_u16(), body_text); } serde_json::from_slice::<SnowflakeStatementResponse>(&body_bytes) .context("failed to parse Snowflake SQL API response JSON") } pub async fn fetch_rejected_examples_after( http_client: Arc<dyn HttpClient>, after_timestamps: &[String], max_rows_per_timestamp: usize, offset: usize, background_executor: BackgroundExecutor, min_capture_version: Option<MinCaptureVersion>, ) -> Result<Vec<Example>> { if after_timestamps.is_empty() { return Ok(Vec::new()); } let progress = Progress::global(); let mut all_examples = Vec::new(); for after_date in after_timestamps.iter() { let step_progress_name = format!("rejected>{after_date}"); let step_progress = progress.start(Step::PullExamples, &step_progress_name); step_progress.set_substatus("querying"); let min_minor_str = min_capture_version.map(|version| version.minor.to_string()); let min_patch_str = min_capture_version.map(|version| version.patch.to_string()); let min_minor_str_ref = min_minor_str.as_deref(); let min_patch_str_ref = min_patch_str.as_deref(); let statement = indoc! {r#" SELECT req.event_properties:request_id::string AS request_id, req.device_id::string AS device_id, req.time::string AS time, req.event_properties:input AS input, req.event_properties:prompt::string AS prompt, req.event_properties:output::string AS output, rej.event_properties:was_shown::boolean AS was_shown, rej.event_properties:reason::string AS reason, req.event_properties:zed_version::string AS zed_version FROM events req INNER JOIN events rej ON req.event_properties:request_id = rej.event_properties:request_id WHERE req.event_type = ? AND rej.event_type = ? AND req.event_properties:version = 'V3' AND rej.event_properties:was_shown = true AND req.event_properties:input:can_collect_data = true AND req.time > TRY_TO_TIMESTAMP_NTZ(?) AND (? IS NULL OR ( TRY_CAST(SPLIT_PART(req.event_properties:zed_version::string, '.', 2) AS INTEGER) > ? OR ( TRY_CAST(SPLIT_PART(req.event_properties:zed_version::string, '.', 2) AS INTEGER) = ? AND TRY_CAST(SPLIT_PART(SPLIT_PART(req.event_properties:zed_version::string, '.', 3), '+', 1) AS INTEGER) >= ? ) )) ORDER BY req.time ASC LIMIT ? OFFSET ? "#}; let bindings = json!({ "1": { "type": "TEXT", "value": PREDICTIVE_EDIT_REQUESTED_EVENT }, "2": { "type": "TEXT", "value": PREDICTIVE_EDIT_REJECTED_EVENT }, "3": { "type": "TEXT", "value": after_date }, "4": { "type": "FIXED", "value": min_minor_str_ref }, "5": { "type": "FIXED", "value": min_minor_str_ref }, "6": { "type": "FIXED", "value": min_minor_str_ref }, "7": { "type": "FIXED", "value": min_patch_str_ref }, "8": { "type": "FIXED", "value": max_rows_per_timestamp.to_string() }, "9": { "type": "FIXED", "value": offset.to_string() } }); let examples = fetch_examples_with_query( http_client.clone(), &step_progress, background_executor.clone(), statement, bindings, DEFAULT_STATEMENT_TIMEOUT_SECONDS, &[ "request_id", "device_id", "time", "input", "prompt", "output", "was_shown", "reason", "zed_version", ], rejected_examples_from_response, ) .await?; all_examples.extend(examples); } Ok(all_examples) } pub async fn fetch_requested_examples_after( http_client: Arc<dyn HttpClient>, after_timestamps: &[String], max_rows_per_timestamp: usize, offset: usize, background_executor: BackgroundExecutor, min_capture_version: Option<MinCaptureVersion>, ) -> Result<Vec<Example>> { if after_timestamps.is_empty() { return Ok(Vec::new()); }