1use crate::{
2 constants::{CX_REWRITER_PATH, J4RS_BASE_PATH},
3 prelude::*,
4};
5use fehler::throws;
6use j4rs::{ClasspathEntry, Instance, InvocationArg, Jvm, JvmBuilder};
7use log::debug;
8use std::collections::HashMap;
9use std::convert::TryFrom;
10use std::{env, fs};
11
12pub struct Plan {
13 pub db_name: String,
14 pub db_alias: String,
15 pub sql: String,
16 pub cardinality: usize,
17}
18
19pub struct FederatedDataSourceInfo<'a> {
20 pub conn_str_info: Option<SourceConn>,
21 pub manual_info: Option<HashMap<String, Vec<String>>>,
22 pub is_local: bool,
23 pub jdbc_url: &'a str,
24 pub jdbc_driver: &'a str,
25}
26
27impl<'a> FederatedDataSourceInfo<'a> {
28 pub fn new_from_conn_str(
29 source_conn: SourceConn,
30 is_local: bool,
31 jdbc_url: &'a str,
32 jdbc_driver: &'a str,
33 ) -> Self {
34 Self {
35 conn_str_info: Some(source_conn),
36 manual_info: None,
37 is_local,
38 jdbc_url,
39 jdbc_driver,
40 }
41 }
42 pub fn new_from_manual_schema(
43 manual_schema: HashMap<String, Vec<String>>,
44 is_local: bool,
45 ) -> Self {
46 Self {
47 conn_str_info: None,
48 manual_info: Some(manual_schema),
49 is_local,
50 jdbc_url: "",
51 jdbc_driver: "",
52 }
53 }
54}
55
56#[throws(ConnectorXOutError)]
57fn init_jvm(j4rs_base: Option<&str>) -> Jvm {
58 let base = match j4rs_base {
59 Some(path) => fs::canonicalize(path)
60 .map_err(|_| ConnectorXOutError::FileNotFoundError(path.to_string()))?,
61 None => fs::canonicalize(J4RS_BASE_PATH)
62 .map_err(|_| ConnectorXOutError::FileNotFoundError(J4RS_BASE_PATH.to_string()))?,
63 };
64 debug!("j4rs base path: {:?}", base);
65
66 let rewriter_path = env::var("CX_REWRITER_PATH").unwrap_or(CX_REWRITER_PATH.to_string());
67 let path = fs::canonicalize(rewriter_path.as_str())
68 .map_err(|_| ConnectorXOutError::FileNotFoundError(rewriter_path))?;
69
70 debug!("rewriter path: {:?}", path);
71
72 let entry = ClasspathEntry::new(path.to_str().unwrap());
73 JvmBuilder::new()
74 .skip_setting_native_lib()
75 .classpath_entry(entry)
76 .with_base_path(base.to_str().unwrap())
77 .build()?
78}
79
80#[allow(dead_code)]
81#[throws(ConnectorXOutError)]
82fn create_sources(
83 jvm: &Jvm,
84 db_map: &HashMap<String, FederatedDataSourceInfo>,
85) -> (Instance, Instance) {
86 debug!("Could not find environment variable `FED_CONFIG_PATH`, use manual configuration (c++ API only)!");
87 let mut db_config = vec![];
88 let db_manual = jvm.create_instance("java.util.HashMap", InvocationArg::empty())?;
89
90 for (db_name, db_info) in db_map.iter() {
91 if let Some(manual_info) = &db_info.manual_info {
92 let schema_info = jvm.create_instance("java.util.HashMap", InvocationArg::empty())?;
93 for (name, columns) in manual_info {
94 let arr_instance = jvm.java_list("java.lang.String", columns.to_vec())?;
95 jvm.invoke(
96 &schema_info,
97 "put",
98 &[
99 InvocationArg::try_from(name).unwrap(),
100 InvocationArg::from(arr_instance),
101 ],
102 )?;
103 }
104 let fed_ds = jvm.create_instance(
105 "ai.dataprep.federated.FederatedDataSource",
106 &[
107 InvocationArg::try_from(db_info.is_local).unwrap(),
108 InvocationArg::from(schema_info),
109 ],
110 )?;
111 jvm.invoke(
112 &db_manual,
113 "put",
114 &[
115 InvocationArg::try_from(db_name).unwrap(),
116 InvocationArg::from(fed_ds),
117 ],
118 )?;
119 } else {
120 db_config.push(String::from(db_name));
121 }
122 }
123 let db_config = jvm.java_list("java.lang.String", db_config)?;
124 (db_config, db_manual)
125}
126
127#[allow(dead_code)]
128#[throws(ConnectorXOutError)]
129fn create_sources2(
130 jvm: &Jvm,
131 db_map: &HashMap<String, FederatedDataSourceInfo>,
132) -> (Instance, Instance) {
133 debug!("Found environment variable `FED_CONFIG_PATH`, use configurations!");
134 let mut dbs = vec![];
135 let db_manual = jvm.create_instance("java.util.HashMap", InvocationArg::empty())?;
136 for db in db_map.keys() {
137 dbs.push(String::from(db));
138 }
139 (jvm.java_list("java.lang.String", dbs)?, db_manual)
140}
141
142#[throws(ConnectorXOutError)]
143pub fn rewrite_sql(
144 sql: &str,
145 db_map: &HashMap<String, FederatedDataSourceInfo>,
146 j4rs_base: Option<&str>,
147 strategy: &str,
148) -> Vec<Plan> {
149 let jvm = init_jvm(j4rs_base)?;
150 debug!("init jvm successfully!");
151
152 let sql = InvocationArg::try_from(sql).unwrap();
153 let strategy = InvocationArg::try_from(strategy).unwrap();
154
155 let (db_config, db_manual) = match env::var("FED_CONFIG_PATH") {
156 Ok(_) => create_sources2(&jvm, db_map)?,
157 _ => create_sources(&jvm, db_map)?,
158 };
159 let rewriter = jvm.create_instance(
160 "ai.dataprep.accio.FederatedQueryRewriter",
161 InvocationArg::empty(),
162 )?;
163 let db_config = InvocationArg::from(db_config);
164 let db_manual = InvocationArg::from(db_manual);
165 let plan = jvm.invoke(&rewriter, "rewrite", &[sql, db_config, db_manual, strategy])?;
166
167 let count = jvm.invoke(&plan, "getCount", InvocationArg::empty())?;
168 let count: i32 = jvm.to_rust(count)?;
169 debug!("rewrite finished, got {} queries", count);
170
171 let mut fed_plan = vec![];
172 for i in 0..count {
173 let idx = [InvocationArg::try_from(i).unwrap().into_primitive()?];
174
175 let db = jvm.invoke(&plan, "getDBName", &idx)?;
176 let db: String = jvm.to_rust(db)?;
177
178 let alias_db = jvm.invoke(&plan, "getAliasDBName", &idx)?;
179 let alias_db: String = jvm.to_rust(alias_db)?;
180
181 let rewrite_sql = jvm.invoke(&plan, "getSql", &idx)?;
182 let rewrite_sql: String = jvm.to_rust(rewrite_sql)?;
183
184 let cardinality = jvm.invoke(&plan, "getCardinality", &idx)?;
185 let cardinality: usize = jvm.to_rust(cardinality)?;
186
187 debug!(
188 "{} - db: {}, alias: {}, cardinality: {}, rewrite sql: {}",
189 i, db, alias_db, cardinality, rewrite_sql
190 );
191 fed_plan.push(Plan {
192 db_name: db,
193 db_alias: alias_db,
194 sql: rewrite_sql,
195 cardinality,
196 });
197 }
198 fed_plan
199}