use std::{
collections::{HashMap, HashSet},
path::PathBuf,
};
use crate::{
collection::Collection,
request::{RedactionRules, RequestTemplate, VariableStore, expand_templates},
schema::SchemaRegistry,
};
use super::{
Rule, context,
declarations::{
DeclarationSpanIndex, parse_request_collection, register_declaration,
remember_declaration_span, validate_schema_registry,
},
legacy_header::normalize_legacy_collection_header,
preprocessor,
protocol::SessionRequestTarget,
redaction::apply_redaction_directive,
request::{
parse_request_template,
},
spans::{
prompt_error_to_span, template_error_to_collection_error,
template_error_to_span,
},
variables::{
assign_global_variable, assign_global_variable_for_selected_environment_override,
SecretValueCache,
},
};
pub fn parse_collection(
input: &str,
working_dir: PathBuf,
) -> Result<Collection, pest::error::Error<Rule>> {
parse_collection_with_environment(input, working_dir, None)
}
pub fn parse_collection_with_environment(
input: &str,
working_dir: PathBuf,
selected_environment: Option<&str>,
) -> Result<Collection, pest::error::Error<Rule>> {
let preprocessed = preprocessor::preprocess(input, &working_dir).map_err(|e| {
pest::error::Error::new_from_span(
pest::error::ErrorVariant::CustomError {
message: e.to_string(),
},
pest::Span::new(input, 0, input.len()).unwrap(),
)
})?;
let preprocessed = normalize_legacy_collection_header(preprocessed.as_str());
log::debug!("PREPROCESSED COMPLETE:\n{}", preprocessed);
let collection = parse_request_collection(preprocessed.as_str())?;
let selected_environment_override_keys = selected_environment_override_keys(
preprocessed.as_str(),
selected_environment,
)?;
let mut name = String::new();
let mut description = String::new();
let mut request_templates: Vec<RequestTemplate> = Vec::new();
let mut global_store = VariableStore::new();
let mut global_headers: HashMap<String, String> = HashMap::new();
let mut global_queries: HashMap<String, String> = HashMap::new();
let mut global_callbacks: Vec<String> = vec![];
let mut schema_registry = SchemaRegistry::default();
let mut declaration_spans = DeclarationSpanIndex::new();
let mut session_targets: HashMap<String, SessionRequestTarget> = HashMap::new();
let mut environments: HashMap<String, HashMap<String, String>> = HashMap::new();
let mut redaction_rules = RedactionRules::default();
let mut secret_cache = SecretValueCache::default();
for pair in collection.into_inner() {
match pair.as_rule() {
Rule::collection_name => {
name = pair.as_str().trim().to_string();
}
Rule::collection_description => {
description.push_str(pair.as_str().trim());
}
Rule::variable => {
let span = pair.as_span();
let mut inner_pairs = pair.into_inner();
let key = inner_pairs.next().unwrap().as_str().trim().to_string();
let value = inner_pairs.next().unwrap().as_str().to_string();
if selected_environment_override_keys.contains(&key) {
assign_global_variable_for_selected_environment_override(
&mut global_store,
key,
value.as_str(),
);
} else {
assign_global_variable(
&mut global_store,
key,
value.as_str(),
&working_dir,
&mut secret_cache,
)
.map_err(|err| template_error_to_span(err, span.clone()))?;
}
}
Rule::environment_block => {
let span = pair.as_span();
let (name, assignments) = parse_environment_block(
pair,
&global_store,
&working_dir,
&mut secret_cache,
)?;
if environments.insert(name.clone(), assignments).is_some() {
return Err(pest::error::Error::new_from_span(
pest::error::ErrorVariant::CustomError {
message: format!("Duplicate environment '{}' declared.", name),
},
span,
));
}
}
Rule::header => {
let span = pair.as_span();
let mut inner_pairs = pair.into_inner();
let key = inner_pairs.next().unwrap().as_str().trim().to_string();
let value = inner_pairs.next().unwrap().as_str().trim().to_string();
let value = context::try_inject_from_prompt(&value)
.map_err(|err| prompt_error_to_span(err, span.clone()))?;
global_headers.insert(key, value);
}
Rule::query => {
let span = pair.as_span();
let mut inner_pairs = pair.into_inner();
let key = inner_pairs.next().unwrap().as_str().trim().to_string();
let value = inner_pairs.next().unwrap().as_str().trim().to_string();
let value = context::try_inject_from_prompt(&value)
.map_err(|err| prompt_error_to_span(err, span.clone()))?;
global_queries.insert(key, value);
}
Rule::callback => {
global_callbacks.push(pair.as_str().strip_prefix('!').unwrap().to_string());
}
Rule::redact_header_directive
| Rule::redact_capture_directive
| Rule::redact_body_directive => {
apply_redaction_directive(&mut redaction_rules, pair)?;
}
Rule::declaration => {
remember_declaration_span(&mut declaration_spans, &pair);
register_declaration(&mut schema_registry, pair)?;
}
Rule::requests => {
validate_schema_registry(
&schema_registry,
preprocessed.as_str(),
&declaration_spans,
)?;
let mut effective_global_store = global_store.clone();
if let Some(selected_environment) = selected_environment {
let assignments = environments.get(selected_environment).ok_or_else(|| {
pest::error::Error::new_from_span(
pest::error::ErrorVariant::CustomError {
message: format!(
"Unknown environment '{}'. Available environments: {}.",
selected_environment,
environment_names(&environments)
),
},
pest::Span::new(preprocessed.as_str(), 0, preprocessed.len())
.expect("source span should exist"),
)
})?;
for (key, value) in assignments {
effective_global_store.set_scalar(key.clone(), value.clone());
}
}
for request_pair in pair.into_inner() {
let template = parse_request_template(
request_pair,
&effective_global_store,
&global_headers,
&global_queries,
&global_callbacks,
&schema_registry,
&redaction_rules,
&working_dir,
&mut secret_cache,
&mut session_targets,
)?;
request_templates.push(template);
}
}
Rule::EOI => {}
_ => {
unreachable!("unexpected rule: {:?}", pair.as_rule());
}
}
}
let requests = expand_templates(request_templates)
.map_err(|err| template_error_to_collection_error(err, preprocessed.as_str()))?;
let mut available_environments = environments.keys().cloned().collect::<Vec<_>>();
available_environments.sort();
Ok(Collection {
name,
description,
available_environments,
selected_environment: selected_environment.map(str::to_string),
schema_registry,
requests,
})
}
fn parse_environment_block(
pair: pest::iterators::Pair<Rule>,
global_store: &VariableStore,
working_dir: &PathBuf,
secret_cache: &mut SecretValueCache,
) -> Result<(String, HashMap<String, String>), pest::error::Error<Rule>> {
let span = pair.as_span();
let mut inner_pairs = pair.into_inner();
let name = inner_pairs
.next()
.expect("environment block should include a name")
.as_str()
.trim()
.to_string();
let mut env_store = global_store.clone();
let mut assignments = HashMap::new();
for env_variable in inner_pairs {
let env_span = env_variable.as_span();
let mut variable_parts = env_variable.into_inner();
let key = variable_parts
.next()
.expect("environment variable should include a key")
.as_str()
.trim()
.to_string();
let value = variable_parts
.next()
.expect("environment variable should include a value")
.as_str()
.to_string();
if !env_store.scalars().contains_key(&key) {
return Err(template_error_to_span(
crate::request::TemplateError::UnknownEnvironmentVariable {
environment: name.clone(),
variable: key,
},
env_span,
));
}
super::variables::assign_environment_variable(
&mut env_store,
&name,
key.clone(),
&value,
working_dir,
secret_cache,
)
.map_err(|err| template_error_to_span(err, env_span.clone()))?;
let resolved = env_store
.scalars()
.get(&key)
.expect("environment assignment should set the target key")
.clone();
assignments.insert(key, resolved);
}
if assignments.is_empty() {
return Err(pest::error::Error::new_from_span(
pest::error::ErrorVariant::CustomError {
message: "Environment blocks must declare at least one variable override.".to_string(),
},
span,
));
}
let _ = working_dir;
Ok((name, assignments))
}
fn environment_names(environments: &HashMap<String, HashMap<String, String>>) -> String {
let mut names = environments.keys().cloned().collect::<Vec<_>>();
names.sort();
if names.is_empty() {
"<none>".to_string()
} else {
names.join(", ")
}
}
fn selected_environment_override_keys(
source: &str,
selected_environment: Option<&str>,
) -> Result<HashSet<String>, pest::error::Error<Rule>> {
let Some(selected_environment) = selected_environment else {
return Ok(HashSet::new());
};
let collection = parse_request_collection(source)?;
let mut keys = HashSet::new();
for pair in collection.into_inner() {
if pair.as_rule() != Rule::environment_block {
continue;
}
let mut inner_pairs = pair.into_inner();
let name = inner_pairs
.next()
.expect("environment block should include a name")
.as_str()
.trim()
.to_string();
if name != selected_environment {
continue;
}
for env_variable in inner_pairs {
let mut variable_parts = env_variable.into_inner();
let key = variable_parts
.next()
.expect("environment variable should include a key")
.as_str()
.trim()
.to_string();
keys.insert(key);
}
}
Ok(keys)
}