Skip to content

Commit 54325f0

Browse files
burmeciaburmecia
authored
Merge pull request #323 from remilapeyre/import-foreign-schema
Give the full server struct to the FDW when instanciating it
2 parents 410cbe6 + 256d3ee commit 54325f0

File tree

16 files changed

+110
-66
lines changed

16 files changed

+110
-66
lines changed

‎supabase-wrappers/src/instance.rs

+33-2
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,48 @@
1+
use std::collections::HashMap;
2+
use std::ffi::CStr;
3+
14
use crate::prelude::*;
25
use pgrx::pg_sys::panic::ErrorReport;
36
use pgrx::prelude::*;
47

8+
pub struct ForeignServer {
9+
pub server_name: String,
10+
pub server_type: Option<String>,
11+
pub server_version: Option<String>,
12+
pub options: HashMap<String, String>,
13+
}
14+
515
// create a fdw instance from its id
616
pub(super) unsafe fn create_fdw_instance_from_server_id<
717
E: Into<ErrorReport>,
818
W: ForeignDataWrapper<E>,
919
>(
1020
fserver_id: pg_sys::Oid,
1121
) -> W {
22+
let to_string = |raw: *mut std::ffi::c_char| -> Option<String> {
23+
if raw.is_null() {
24+
return None;
25+
}
26+
let c_str = CStr::from_ptr(raw);
27+
let value = c_str
28+
.to_str()
29+
.map_err(|_| {
30+
OptionsError::OptionValueIsInvalidUtf8(
31+
String::from_utf8_lossy(c_str.to_bytes()).to_string(),
32+
)
33+
})
34+
.report_unwrap()
35+
.to_string();
36+
Some(value)
37+
};
1238
let fserver = pg_sys::GetForeignServer(fserver_id);
13-
let fserver_opts = options_to_hashmap((*fserver).options).report_unwrap();
14-
let wrapper = W::new(&fserver_opts);
39+
let server = ForeignServer {
40+
server_name: to_string((*fserver).servername).unwrap(),
41+
server_type: to_string((*fserver).servertype),
42+
server_version: to_string((*fserver).serverversion),
43+
options: options_to_hashmap((*fserver).options).report_unwrap(),
44+
};
45+
let wrapper = W::new(server);
1546
wrapper.report_unwrap()
1647
}
1748

‎supabase-wrappers/src/interface.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
//! Provides interface types and trait to develop Postgres foreign data wrapper
22
//!
33
4+
use crate::instance::ForeignServer;
45
use crate::FdwRoutine;
56
use pgrx::pg_sys::panic::ErrorReport;
67
use pgrx::prelude::{Date, Timestamp, TimestampWithTimeZone};
@@ -526,7 +527,7 @@ pub trait ForeignDataWrapper<E: Into<ErrorReport>> {
526527
/// You can do any initalization in this function, like saving connection
527528
/// info or API url in an variable, but don't do heavy works like database
528529
/// connection or API call.
529-
fn new(options: &HashMap<String, String>) -> Result<Self, E>
530+
fn new(server: ForeignServer) -> Result<Self, E>
530531
where
531532
Self: Sized;
532533

‎supabase-wrappers/src/lib.rs

+4-3
Original file line numberDiff line numberDiff line change
@@ -90,8 +90,8 @@
9090
//! type HelloWorldFdwResult<T> = Result<T, HelloWorldFdwError>;
9191
//!
9292
//! impl ForeignDataWrapper<HelloWorldFdwError> for HelloWorldFdw {
93-
//! fn new(options: &HashMap<String, String>) -> HelloWorldFdwResult<Self> {
94-
//! // 'options' is the key-value pairs defined in `CREATE SERVER` SQL, for example,
93+
//! fn new(server: ForeignServer) -> HelloWorldFdwResult<Self> {
94+
//! // 'server.options' is the key-value pairs defined in `CREATE SERVER` SQL, for example,
9595
//! //
9696
//! // create server my_helloworld_server
9797
//! // foreign data wrapper wrappers_helloworld
@@ -172,7 +172,7 @@
172172
//! }
173173
//!
174174
//! impl ForeignDataWrapper<HelloWorldFdwError> for HelloWorldFdw {
175-
//! fn new(options: &HashMap<String, String>) -> Result<Self, HelloWorldFdwError> {
175+
//! fn new(server: ForeignServer) -> Result<Self, HelloWorldFdwError> {
176176
//! Ok(Self {
177177
//! row_cnt: 0,
178178
//! tgt_cols: Vec::new(),
@@ -299,6 +299,7 @@ pub mod utils;
299299
/// The prelude includes all necessary imports to make Wrappers work
300300
pub mod prelude {
301301
pub use crate::import_foreign_schema::*;
302+
pub use crate::instance::ForeignServer;
302303
pub use crate::interface::*;
303304
pub use crate::options::*;
304305
pub use crate::utils::*;

‎wrappers/src/fdw/airtable_fdw/airtable_fdw.rs

+5-4
Original file line numberDiff line numberDiff line change
@@ -90,16 +90,17 @@ impl AirtableFdw {
9090

9191
// TODO Add support for INSERT, UPDATE, DELETE
9292
impl ForeignDataWrapper<AirtableFdwError> for AirtableFdw {
93-
fn new(options: &HashMap<String, String>) -> AirtableFdwResult<Self> {
94-
let base_url = options
93+
fn new(server: ForeignServer) -> AirtableFdwResult<Self> {
94+
let base_url = server
95+
.options
9596
.get("api_url")
9697
.map(|t| t.to_owned())
9798
.unwrap_or_else(|| "https://api.airtable.com/v0".to_string());
9899

99-
let client = match options.get("api_key") {
100+
let client = match server.options.get("api_key") {
100101
Some(api_key) => Some(create_client(api_key)?),
101102
None => {
102-
let key_id = require_option("api_key_id", options)?;
103+
let key_id = require_option("api_key_id", &server.options)?;
103104
if let Some(api_key) = get_vault_secret(key_id) {
104105
Some(create_client(&api_key)?)
105106
} else {

‎wrappers/src/fdw/auth0_fdw/auth0_fdw.rs

+5-4
Original file line numberDiff line numberDiff line change
@@ -94,12 +94,13 @@ impl ForeignDataWrapper<Auth0FdwError> for Auth0Fdw {
9494
// info or API url in an variable, but don't do any heavy works like making a
9595
// database connection or API call.
9696

97-
fn new(options: &HashMap<String, String>) -> Result<Self, Auth0FdwError> {
98-
let url = require_option("url", options)?.to_string();
99-
let api_key = if let Some(api_key) = options.get("api_key") {
97+
fn new(server: ForeignServer) -> Result<Self, Auth0FdwError> {
98+
let url = require_option("url", &server.options)?.to_string();
99+
let api_key = if let Some(api_key) = server.options.get("api_key") {
100100
api_key.clone()
101101
} else {
102-
let api_key_id = options
102+
let api_key_id = server
103+
.options
103104
.get("api_key_id")
104105
.expect("`api_key_id` must be set if `api_key` is not");
105106
get_vault_secret(api_key_id).ok_or(Auth0FdwError::SecretNotFound(api_key_id.clone()))?

‎wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs

+9-7
Original file line numberDiff line numberDiff line change
@@ -146,12 +146,12 @@ impl BigQueryFdw {
146146
}
147147

148148
impl ForeignDataWrapper<BigQueryFdwError> for BigQueryFdw {
149-
fn new(options: &HashMap<String, String>) -> Result<Self, BigQueryFdwError> {
149+
fn new(server: ForeignServer) -> Result<Self, BigQueryFdwError> {
150150
let mut ret = BigQueryFdw {
151151
rt: create_async_runtime()?,
152152
client: None,
153-
project_id: require_option("project_id", options)?.to_string(),
154-
dataset_id: require_option("dataset_id", options)?.to_string(),
153+
project_id: require_option("project_id", &server.options)?.to_string(),
154+
dataset_id: require_option("dataset_id", &server.options)?.to_string(),
155155
table: "".to_string(),
156156
rowid_col: "".to_string(),
157157
tgt_cols: Vec::new(),
@@ -160,13 +160,15 @@ impl ForeignDataWrapper<BigQueryFdwError> for BigQueryFdw {
160160
};
161161

162162
// Is authentication mocked
163-
let mock_auth: bool = options
163+
let mock_auth: bool = server
164+
.options
164165
.get("mock_auth")
165166
.map(|t| t.to_owned())
166167
.unwrap_or_else(|| "false".to_string())
167168
== *"true";
168169

169-
let api_endpoint = options
170+
let api_endpoint = server
171+
.options
170172
.get("api_endpoint")
171173
.map(|t| t.to_owned())
172174
.unwrap_or_else(|| "https://bigquery.googleapis.com/bigquery/v2".to_string());
@@ -182,10 +184,10 @@ impl ForeignDataWrapper<BigQueryFdwError> for BigQueryFdw {
182184
serde_json::to_string_pretty(&dummy_auth_config)
183185
.expect("dummy auth config should not fail to serialize")
184186
}
185-
false => match options.get("sa_key") {
187+
false => match server.options.get("sa_key") {
186188
Some(sa_key) => sa_key.to_owned(),
187189
None => {
188-
let sa_key_id = require_option("sa_key_id", options)?;
190+
let sa_key_id = require_option("sa_key_id", &server.options)?;
189191
match get_vault_secret(sa_key_id) {
190192
Some(sa_key) => sa_key,
191193
None => return Ok(ret),

‎wrappers/src/fdw/clickhouse_fdw/clickhouse_fdw.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -201,12 +201,12 @@ impl ClickHouseFdw {
201201
}
202202

203203
impl ForeignDataWrapper<ClickHouseFdwError> for ClickHouseFdw {
204-
fn new(options: &HashMap<String, String>) -> ClickHouseFdwResult<Self> {
204+
fn new(server: ForeignServer) -> ClickHouseFdwResult<Self> {
205205
let rt = create_async_runtime()?;
206-
let conn_str = match options.get("conn_string") {
206+
let conn_str = match server.options.get("conn_string") {
207207
Some(conn_str) => conn_str.to_owned(),
208208
None => {
209-
let conn_str_id = require_option("conn_string_id", options)?;
209+
let conn_str_id = require_option("conn_string_id", &server.options)?;
210210
get_vault_secret(conn_str_id).unwrap_or_default()
211211
}
212212
};

‎wrappers/src/fdw/cognito_fdw/cognito_fdw.rs

+8-7
Original file line numberDiff line numberDiff line change
@@ -97,16 +97,17 @@ impl ForeignDataWrapper<CognitoFdwError> for CognitoFdw {
9797
// info or API url in an variable, but don't do any heavy works like making a
9898
// database connection or API call.
9999

100-
fn new(options: &HashMap<String, String>) -> Result<Self, CognitoFdwError> {
101-
let user_pool_id = require_option("user_pool_id", options)?.to_string();
102-
let aws_region = require_option("region", options)?.to_string();
100+
fn new(server: ForeignServer) -> Result<Self, CognitoFdwError> {
101+
let user_pool_id = require_option("user_pool_id", &server.options)?.to_string();
102+
let aws_region = require_option("region", &server.options)?.to_string();
103103

104-
let aws_access_key_id = require_option("aws_access_key_id", options)?.to_string();
104+
let aws_access_key_id = require_option("aws_access_key_id", &server.options)?.to_string();
105105
let aws_secret_access_key =
106-
if let Some(aws_secret_access_key) = options.get("aws_secret_access_key") {
106+
if let Some(aws_secret_access_key) = server.options.get("aws_secret_access_key") {
107107
aws_secret_access_key.clone()
108108
} else {
109-
let aws_secret_access_key = options
109+
let aws_secret_access_key = server
110+
.options
110111
.get("api_key_id")
111112
.expect("`api_key_id` must be set if `aws_secret_access_key` is not");
112113
get_vault_secret(aws_secret_access_key).ok_or(CognitoFdwError::SecretNotFound(
@@ -122,7 +123,7 @@ impl ForeignDataWrapper<CognitoFdwError> for CognitoFdw {
122123
let config = aws_config::load_defaults(BehaviorVersion::latest()).await;
123124

124125
let mut builder = config.to_builder();
125-
if let Some(endpoint_url) = options.get("endpoint_url") {
126+
if let Some(endpoint_url) = server.options.get("endpoint_url") {
126127
if !endpoint_url.is_empty() {
127128
builder.set_endpoint_url(Some(endpoint_url.clone()));
128129
}

‎wrappers/src/fdw/firebase_fdw/firebase_fdw.rs

+5-5
Original file line numberDiff line numberDiff line change
@@ -211,23 +211,23 @@ impl FirebaseFdw {
211211
}
212212

213213
impl ForeignDataWrapper<FirebaseFdwError> for FirebaseFdw {
214-
fn new(options: &HashMap<String, String>) -> FirebaseFdwResult<Self> {
214+
fn new(server: ForeignServer) -> FirebaseFdwResult<Self> {
215215
let mut ret = Self {
216216
rt: create_async_runtime()?,
217-
project_id: require_option("project_id", options)?.to_string(),
217+
project_id: require_option("project_id", &server.options)?.to_string(),
218218
client: None,
219219
scan_result: Vec::default(),
220220
};
221221

222222
// get oauth2 access token if it is directly defined in options
223-
let token = if let Some(access_token) = options.get("access_token") {
223+
let token = if let Some(access_token) = server.options.get("access_token") {
224224
access_token.to_owned()
225225
} else {
226226
// otherwise, get it from the options or Vault
227-
let sa_key = match options.get("sa_key") {
227+
let sa_key = match server.options.get("sa_key") {
228228
Some(sa_key) => sa_key.to_owned(),
229229
None => {
230-
let sa_key_id = require_option("sa_key_id", options)?;
230+
let sa_key_id = require_option("sa_key_id", &server.options)?;
231231
match get_vault_secret(sa_key_id) {
232232
Some(sa_key) => sa_key,
233233
None => return Ok(ret),

‎wrappers/src/fdw/helloworld_fdw/helloworld_fdw.rs

+1-1
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ impl ForeignDataWrapper<HelloWorldFdwError> for HelloWorldFdw {
4242
// You can do any initalization in this new() function, like saving connection
4343
// info or API url in an variable, but don't do any heavy works like making a
4444
// database connection or API call.
45-
fn new(_options: &HashMap<String, String>) -> HelloWorldFdwResult<Self> {
45+
fn new(_server: ForeignServer) -> HelloWorldFdwResult<Self> {
4646
Ok(Self {
4747
row_cnt: 0,
4848
tgt_cols: Vec::new(),

‎wrappers/src/fdw/logflare_fdw/logflare_fdw.rs

+5-4
Original file line numberDiff line numberDiff line change
@@ -193,8 +193,9 @@ impl LogflareFdw {
193193
}
194194

195195
impl ForeignDataWrapper<LogflareFdwError> for LogflareFdw {
196-
fn new(options: &HashMap<String, String>) -> LogflareFdwResult<Self> {
197-
let base_url = options
196+
fn new(server: ForeignServer) -> LogflareFdwResult<Self> {
197+
let base_url = server
198+
.options
198199
.get("api_url")
199200
.map(|t| t.to_owned())
200201
.map(|s| {
@@ -205,10 +206,10 @@ impl ForeignDataWrapper<LogflareFdwError> for LogflareFdw {
205206
}
206207
})
207208
.unwrap_or_else(|| LogflareFdw::BASE_URL.to_string());
208-
let client = match options.get("api_key") {
209+
let client = match server.options.get("api_key") {
209210
Some(api_key) => Some(create_client(api_key)),
210211
None => {
211-
let key_id = require_option("api_key_id", options)?;
212+
let key_id = require_option("api_key_id", &server.options)?;
212213
get_vault_secret(key_id).map(|api_key| create_client(&api_key))
213214
}
214215
}

‎wrappers/src/fdw/mssql_fdw/mssql_fdw.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -170,12 +170,12 @@ impl MssqlFdw {
170170
}
171171

172172
impl ForeignDataWrapper<MssqlFdwError> for MssqlFdw {
173-
fn new(options: &HashMap<String, String>) -> MssqlFdwResult<Self> {
173+
fn new(server: ForeignServer) -> MssqlFdwResult<Self> {
174174
let rt = create_async_runtime()?;
175-
let conn_str = match options.get("conn_string") {
175+
let conn_str = match server.options.get("conn_string") {
176176
Some(conn_str) => conn_str.to_owned(),
177177
None => {
178-
let conn_str_id = require_option("conn_string_id", options)?;
178+
let conn_str_id = require_option("conn_string_id", &server.options)?;
179179
get_vault_secret(conn_str_id).unwrap_or_default()
180180
}
181181
};

‎wrappers/src/fdw/redis_fdw/redis_fdw.rs

+3-3
Original file line numberDiff line numberDiff line change
@@ -240,11 +240,11 @@ impl RedisFdw {
240240
}
241241

242242
impl ForeignDataWrapper<RedisFdwError> for RedisFdw {
243-
fn new(options: &HashMap<String, String>) -> RedisFdwResult<Self> {
244-
let conn_url = match options.get("conn_url") {
243+
fn new(server: ForeignServer) -> RedisFdwResult<Self> {
244+
let conn_url = match server.options.get("conn_url") {
245245
Some(url) => url.to_owned(),
246246
None => {
247-
let conn_url_id = require_option("conn_url_id", options)?;
247+
let conn_url_id = require_option("conn_url_id", &server.options)?;
248248
get_vault_secret(conn_url_id).unwrap_or_default()
249249
}
250250
};

‎wrappers/src/fdw/s3_fdw/s3_fdw.rs

+9-8
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ impl S3Fdw {
112112
}
113113

114114
impl ForeignDataWrapper<S3FdwError> for S3Fdw {
115-
fn new(options: &HashMap<String, String>) -> S3FdwResult<Self> {
115+
fn new(server: ForeignServer) -> S3FdwResult<Self> {
116116
// cannot use create_async_runtime() as the runtime needs to be created
117117
// for multiple threads
118118
let rt = tokio::runtime::Runtime::new()
@@ -128,27 +128,27 @@ impl ForeignDataWrapper<S3FdwError> for S3Fdw {
128128
};
129129

130130
// get is_mock flag
131-
let is_mock: bool = options.get("is_mock") == Some(&"true".to_string());
131+
let is_mock: bool = server.options.get("is_mock") == Some(&"true".to_string());
132132

133133
// get credentials
134134
let creds = if is_mock {
135135
// LocalStack uses hardcoded credentials
136136
Some(("test".to_string(), "test".to_string()))
137137
} else {
138-
match options.get("vault_access_key_id") {
138+
match server.options.get("vault_access_key_id") {
139139
Some(vault_access_key_id) => {
140140
// if using credentials stored in Vault
141141
let vault_secret_access_key =
142-
require_option("vault_secret_access_key", options)?;
142+
require_option("vault_secret_access_key", &server.options)?;
143143
get_vault_secret(vault_access_key_id)
144144
.zip(get_vault_secret(vault_secret_access_key))
145145
}
146146
None => {
147147
// if using credentials directly specified
148148
let aws_access_key_id =
149-
require_option("aws_access_key_id", options)?.to_string();
149+
require_option("aws_access_key_id", &server.options)?.to_string();
150150
let aws_secret_access_key =
151-
require_option("aws_secret_access_key", options)?.to_string();
151+
require_option("aws_secret_access_key", &server.options)?.to_string();
152152
Some((aws_access_key_id, aws_secret_access_key))
153153
}
154154
}
@@ -163,7 +163,8 @@ impl ForeignDataWrapper<S3FdwError> for S3Fdw {
163163
let region = if is_mock {
164164
default_region
165165
} else {
166-
options
166+
server
167+
.options
167168
.get("aws_region")
168169
.map(|t| t.to_owned())
169170
.unwrap_or(default_region)
@@ -177,7 +178,7 @@ impl ForeignDataWrapper<S3FdwError> for S3Fdw {
177178
let mut config_loader = aws_config::defaults(BehaviorVersion::latest());
178179

179180
// endpoint_url not supported as env var in rust https://github.com/awslabs/aws-sdk-rust/issues/932
180-
if let Some(endpoint_url) = options.get("endpoint_url") {
181+
if let Some(endpoint_url) = server.options.get("endpoint_url") {
181182
config_loader = config_loader.endpoint_url(endpoint_url);
182183
}
183184

0 commit comments

Comments
 (0)