f1a0070b71
* feat: add graphql support * fix: use custom editor for viewing graphql schema * fix: graphql parser cargo version * fix: add graphql where missing
295 lines
8.8 KiB
Rust
295 lines
8.8 KiB
Rust
use base64::{engine, Engine as _};
|
|
use core::fmt::Write;
|
|
use jsonwebtoken::{encode, Algorithm, EncodingKey, Header};
|
|
use pem;
|
|
use serde_json::{json, Value};
|
|
use sha2::{Digest, Sha256};
|
|
|
|
use windmill_common::error::Error;
|
|
use windmill_common::jobs::QueuedJob;
|
|
use windmill_parser_sql::parse_snowflake_sig;
|
|
use windmill_queue::HTTP_CLIENT;
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
use crate::{transform_json_value, AuthedClient, JobCompleted};
|
|
|
|
#[derive(Serialize)]
|
|
struct Claims {
|
|
iss: String,
|
|
sub: String,
|
|
iat: i64,
|
|
exp: i64,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct SnowflakeDatabase {
|
|
account_identifier: String,
|
|
public_key: String,
|
|
private_key: String,
|
|
username: String,
|
|
database: Option<String>,
|
|
schema: Option<String>,
|
|
warehouse: Option<String>,
|
|
role: Option<String>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
#[allow(non_snake_case)]
|
|
struct SnowflakeResponse {
|
|
data: Vec<Vec<Value>>,
|
|
resultSetMetaData: SnowflakeResultSetMetaData,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
#[allow(non_snake_case)]
|
|
struct SnowflakeResultSetMetaData {
|
|
numRows: i64,
|
|
rowType: Vec<SnowflakeRowType>,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct SnowflakeRowType {
|
|
name: String,
|
|
r#type: String,
|
|
}
|
|
|
|
#[allow(non_snake_case)]
|
|
#[derive(Deserialize)]
|
|
struct SnowflakeError {
|
|
message: String,
|
|
}
|
|
|
|
pub async fn do_snowflake(
|
|
job: QueuedJob,
|
|
client: &AuthedClient,
|
|
query: &str,
|
|
) -> windmill_common::error::Result<JobCompleted> {
|
|
let args = if let Some(args) = &job.args {
|
|
Some(transform_json_value("args", client, &job.workspace_id, args.clone()).await?)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let snowflake_args: Value = serde_json::from_value(args.unwrap_or_else(|| json!({})))
|
|
.map_err(|e| Error::ExecutionErr(e.to_string()))?;
|
|
|
|
let database = serde_json::from_value::<SnowflakeDatabase>(
|
|
snowflake_args.get("database").unwrap_or(&json!({})).clone(),
|
|
)
|
|
.map_err(|e| Error::ExecutionErr(e.to_string()))?;
|
|
|
|
let qualified_username =
|
|
format!("{}.{}", database.account_identifier, database.username).to_uppercase();
|
|
|
|
let public_key = pem::parse(database.public_key.as_bytes()).map_err(|e| {
|
|
Error::ExecutionErr(format!("Failed to parse public key: {}", e.to_string()))
|
|
})?;
|
|
let mut public_key_hash = Sha256::new();
|
|
public_key_hash.update(public_key.contents());
|
|
|
|
let public_key_fp = engine::general_purpose::STANDARD.encode(public_key_hash.finalize());
|
|
|
|
let iss = format!("{}.SHA256:{}", qualified_username, public_key_fp);
|
|
|
|
let claims = Claims {
|
|
iss: iss,
|
|
sub: qualified_username,
|
|
iat: chrono::Utc::now().timestamp(),
|
|
exp: (chrono::Utc::now() + chrono::Duration::hours(1)).timestamp(),
|
|
};
|
|
|
|
let private_key = EncodingKey::from_rsa_pem(database.private_key.as_bytes()).map_err(|e| {
|
|
Error::ExecutionErr(format!("Failed to parse private key: {}", e.to_string()))
|
|
})?;
|
|
|
|
let token = encode(&Header::new(Algorithm::RS256), &claims, &private_key)
|
|
.map_err(|e| Error::ExecutionErr(e.to_string()))?;
|
|
|
|
let args = &job
|
|
.args
|
|
.clone()
|
|
.unwrap_or_else(|| json!({}))
|
|
.as_object()
|
|
.map(|x| x.to_owned())
|
|
.unwrap_or_else(|| json!({}).as_object().unwrap().to_owned());
|
|
|
|
let mut bindings = serde_json::Map::new();
|
|
let sig = parse_snowflake_sig(&query)
|
|
.map_err(|x| Error::ExecutionErr(x.to_string()))?
|
|
.args;
|
|
|
|
let mut i = 1;
|
|
for arg in &sig {
|
|
let arg_t = arg.otyp.clone().unwrap_or_else(|| "string".to_string());
|
|
let arg_v = args.get(&arg.name).cloned().unwrap_or(json!(""));
|
|
let snowflake_v = convert_typ_val(arg_t, arg_v);
|
|
|
|
bindings.insert(i.to_string(), snowflake_v);
|
|
i += 1;
|
|
}
|
|
|
|
let mut body = serde_json::Map::new();
|
|
if database.schema.is_some() {
|
|
body.insert(
|
|
"schema".to_string(),
|
|
json!(database.schema.unwrap().to_uppercase()),
|
|
);
|
|
}
|
|
if database.warehouse.is_some() {
|
|
body.insert(
|
|
"warehouse".to_string(),
|
|
json!(database.warehouse.unwrap().to_uppercase()),
|
|
);
|
|
}
|
|
if database.role.is_some() {
|
|
body.insert(
|
|
"role".to_string(),
|
|
json!(database.role.unwrap().to_uppercase()),
|
|
);
|
|
}
|
|
if database.database.is_some() {
|
|
body.insert(
|
|
"database".to_string(),
|
|
json!(database.database.unwrap().to_uppercase()),
|
|
);
|
|
}
|
|
body.insert("statement".to_string(), json!(query));
|
|
body.insert("timeout".to_string(), json!(10)); // in seconds
|
|
|
|
if i > 1 {
|
|
body.insert("bindings".to_string(), json!(bindings));
|
|
}
|
|
|
|
let response = HTTP_CLIENT
|
|
.post(format!(
|
|
"https://{}.snowflakecomputing.com/api/v2/statements/",
|
|
database.account_identifier.to_uppercase()
|
|
))
|
|
.bearer_auth(token)
|
|
.header("X-Snowflake-Authorization-Token-Type", "KEYPAIR_JWT")
|
|
.json(&body)
|
|
.send()
|
|
.await
|
|
.map_err(|e| Error::ExecutionErr(e.to_string()))?;
|
|
|
|
match response.error_for_status_ref() {
|
|
Ok(_) => {
|
|
let result = response
|
|
.json::<SnowflakeResponse>()
|
|
.await
|
|
.map_err(|e| Error::ExecutionErr(e.to_string()))?;
|
|
|
|
if result.resultSetMetaData.numRows > 10000 {
|
|
return Err(Error::ExecutionErr(
|
|
"More than 10000 rows were requested, use LIMIT 10000 to limit the number of rows".to_string(),
|
|
));
|
|
}
|
|
|
|
let rows = result
|
|
.data
|
|
.iter()
|
|
.map(|row| {
|
|
let mut row_map = serde_json::Map::new();
|
|
row.iter()
|
|
.zip(result.resultSetMetaData.rowType.iter())
|
|
.for_each(|(val, row_type)| {
|
|
row_map
|
|
.insert(row_type.name.clone(), parse_val(&val, &row_type.r#type));
|
|
});
|
|
Value::from(row_map)
|
|
})
|
|
.collect();
|
|
|
|
Ok(JobCompleted { job: job, result: rows, logs: "".to_string(), success: true })
|
|
}
|
|
Err(e) => {
|
|
let resp = response.text().await.unwrap_or("".to_string());
|
|
match serde_json::from_str::<SnowflakeError>(&resp) {
|
|
Ok(sf_err) => Err(Error::ExecutionErr(sf_err.message)),
|
|
Err(_) => Err(Error::ExecutionErr(e.to_string())),
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn convert_typ_val(arg_t: String, arg_v: Value) -> Value {
|
|
match arg_t.as_str() {
|
|
"date" => {
|
|
let arr = arg_v
|
|
.as_str()
|
|
.unwrap_or("")
|
|
.split("T")
|
|
.collect::<Vec<&str>>();
|
|
json!({
|
|
"type": "TEXT",
|
|
"value": match arr.as_slice() {
|
|
[date, _] => {
|
|
json!(date)
|
|
}
|
|
_ => {
|
|
arg_v
|
|
}
|
|
}
|
|
})
|
|
}
|
|
"time" => {
|
|
let arr = arg_v
|
|
.as_str()
|
|
.unwrap_or("")
|
|
.split("T")
|
|
.collect::<Vec<&str>>();
|
|
json!({
|
|
"type": "TEXT",
|
|
"value": match arr.as_slice() {
|
|
[_, time] => {
|
|
json!(time)
|
|
}
|
|
_ => {
|
|
arg_v
|
|
}
|
|
}
|
|
})
|
|
}
|
|
"binary" => {
|
|
// convert base64 to hex as expected by snowflake
|
|
let bytes = engine::general_purpose::STANDARD
|
|
.decode(arg_v.as_str().unwrap_or(""))
|
|
.unwrap_or(vec![]);
|
|
let mut hex = String::with_capacity(bytes.len() * 2);
|
|
for byte in bytes {
|
|
write!(hex, "{:02X}", byte).unwrap_or(());
|
|
}
|
|
json!({
|
|
"type": "TEXT",
|
|
"value": hex
|
|
})
|
|
}
|
|
_ => {
|
|
let mut v = arg_v;
|
|
|
|
if !v.is_string() {
|
|
// if not string, convert to string for api request
|
|
v = json!(v.to_string());
|
|
}
|
|
|
|
json!({
|
|
"type": "TEXT", // snowflake infer type from schema
|
|
"value": v
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
fn parse_val(value: &Value, typ: &str) -> Value {
|
|
let str_value = value.as_str().unwrap_or("").to_string();
|
|
match typ.to_lowercase().as_str() {
|
|
"boolean" => json!(str_value.parse::<bool>().ok().unwrap_or(false)),
|
|
"real" | "time" | "timestamp_ltz" | "timestamp_ntz" => {
|
|
json!(str_value.parse::<f64>().ok().unwrap_or(0.0))
|
|
}
|
|
"fixed" | "date" | "number" => json!(str_value.parse::<i64>().ok().unwrap_or(0)),
|
|
_ => value.clone(),
|
|
}
|
|
}
|