feat(snowflake): add timeout_secs foreign table option (supabase#534)

* feat(snowflake): add timeout_secs table option

* fix: change to use server option and add docs

Author: burmecia

docs/catalog/snowflake.md

@@ -17,6 +17,7 @@ The Snowflake Wrapper is a WebAssembly(Wasm) foreign data wrapper which allows y
 
 | Version | Wasm Package URL                                                                                      | Checksum                                                           | Required Wrappers Version |
 | ------- | ----------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------ | ------------------------- |
+| 0.2.1   | `https://github.com/supabase/wrappers/releases/download/wasm_snowflake_fdw_v0.2.1/snowflake_fdw.wasm` | `tbd` | >=0.5.0                   |
 | 0.2.0   | `https://github.com/supabase/wrappers/releases/download/wasm_snowflake_fdw_v0.2.0/snowflake_fdw.wasm` | `921b18a1e9c20c4ef5a09af17b5d76fd6ebe56d41bcfa565b74a530420532437` | >=0.5.0                   |
 | 0.1.1   | `https://github.com/supabase/wrappers/releases/download/wasm_snowflake_fdw_v0.1.1/snowflake_fdw.wasm` | `7aaafc7edc1726bc93ddc04452d41bda9e1a264a1df2ea9bf1b00b267543b860` | >=0.4.0                   |
 | 0.1.0   | `https://github.com/supabase/wrappers/releases/download/wasm_snowflake_fdw_v0.1.0/snowflake_fdw.wasm` | `2fb46fd8afa63f3975dadf772338106b609b131861849356e0c09dde032d1af8` | >=0.4.0                   |
@@ -75,7 +76,8 @@ We need to provide Postgres with the credentials to connect to Snowflake, and an
         account_identifier 'MYORGANIZATION-MYACCOUNT',
         user 'MYUSER',
         public_key_fingerprint 'SizgPofeFX0jwC8IhbOfGFyOggFgo8oTOS1uPLZhzUQ=',
-        private_key_id '<key_ID>' -- The Key ID from above.
+        private_key_id '<key_ID>', -- The Key ID from above.
+        timeout_secs '60'  -- Timeout in seconds for Snowflake statement execution, default is 60, value range is [0, 6048]
       );
     ```
 
@@ -92,7 +94,8 @@ We need to provide Postgres with the credentials to connect to Snowflake, and an
         account_identifier 'MYORGANIZATION-MYACCOUNT',
         user 'MYUSER',
         public_key_fingerprint 'SizgPofeFX0jwC8IhbOfGFyOggFgo8oTOS1uPLZhzUQ=',
-        private_key E'-----BEGIN PRIVATE KEY-----\n...\n-----END PRIVATE KEY-----'
+        private_key E'-----BEGIN PRIVATE KEY-----\n...\n-----END PRIVATE KEY-----',
+        timeout_secs '60'  -- Timeout in seconds for Snowflake statement execution, default is 60, value range is [0, 6048]
       );
     ```
 

wasm-wrappers/fdw/Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "snowflake_fdw"
-version = "0.2.0"
+version = "0.2.1"
 edition = { workspace = true }
 homepage = { workspace = true }
 rust-version = { workspace = true }

wasm-wrappers/fdw/snowflake_fdw/Cargo.toml

@@ -25,6 +25,7 @@ struct SnowflakeFdw {
     src_rows: Vec<JsonValue>,
     src_idx: usize,
     rowid_col: String,
+    timeout_secs: i32,
 }
 
 static mut INSTANCE: *mut SnowflakeFdw = std::ptr::null_mut::<SnowflakeFdw>();
@@ -204,7 +205,10 @@ impl SnowflakeFdw {
     // make the first SQL query request
     fn make_init_request(&mut self, sql: &str) -> FdwResult {
         let url = format!("{}?async=false", self.base_url);
-        let body = format!(r#"{{ "statement": "{sql}", "timeout": 60 }}"#);
+        let body = format!(
+            r#"{{ "statement": "{sql}", "timeout": {} }}"#,
+            self.timeout_secs
+        );
         let (mut resp, mut resp_json) = self.make_post_request(&url, &body)?;
 
         // polling query result
@@ -292,6 +296,14 @@ impl Guest for SnowflakeFdw {
             }
         };
 
+        // parse query timeout from server option
+        // ref: https://docs.snowflake.com/en/developer-guide/sql-api/reference#label-sql-api-reference-post-statements-request-body
+        let timeout_secs = opts
+            .require_or("timeout_secs", "60")
+            .parse::<i32>()
+            .map_err(|e| format!("invalid timeout_secs value: {e}"))?
+            .clamp(0, 6048);
+
         // Snowflake api authentication ref:
         // https://docs.snowflake.com/en/developer-guide/sql-api/authenticating#using-key-pair-authentication
         let claims = vec![
@@ -322,6 +334,8 @@ impl Guest for SnowflakeFdw {
             "KEYPAIR_JWT".to_owned(),
         ));
 
+        this.timeout_secs = timeout_secs;
+
         stats::inc_stats(FDW_NAME, stats::Metric::CreateTimes, 1);
 
         Ok(())

wasm-wrappers/fdw/snowflake_fdw/src/lib.rs

@@ -25,6 +25,7 @@ mod tests {
                        api_url 'http://localhost:8096/snowflake/{}',
                        account_identifier 'TEST-ABC1234',
                        user 'TESTUSER',
+                       timeout_secs '30',
                        public_key_fingerprint 'WVggEofeFX0jwCdImbOfGFyOggF2o8DT7S1uTLZhCJQ=',
                        private_key E'-----BEGIN PRIVATE KEY-----\nMIIEvwIBADANBgkqhkiG9w0BAQEFAASCBKkwggSlAgEAAoIBAQCOQv0mMe1yElR8\nhiQgduWu7OrMR3iW8xbu+i04LkMDKB/JtdMl1Mq3Gs/XWUB07BOIytcK4L7k1Z/3\nnkTGvib85S5+VLsngBzWltltvCOfnM2uCLWHmVpfMR1WVFrCU1r7NP92U2APpwvN\ncJps39VNyQWTm+TvQN25enqAC2uR6xdvItb86dn4Tab3KcAXj0f0qHompa8SSwBO\nGZFpgAjt7QFPWTniNQyBU7wXTntwV+a2WC/bf1i9MRAf8bqSr+yqijLHVGrmBPSb\njmvo/8bj69DzeY2//gZXAMe1m6cEjXPm2MY/POd66Xg1YZDzYsv95GZ54kJ+EgEF\n9rjsi/J5AgMBAAECggEAAsG9kh3pkgpU5Mzcqlxjew5QRoEkDxjK2vqyIaKT3d3L\nL+d8HgGPpBi66ltqalmgz0fO/wD38gtJvEyu3IMW0lPGoOAXeF59MJNfx0acEh3B\nxpuYmPYZ0DptbRzZXWasHq4aPTrEY8lC60pBU9bKlWVN3FxrBU/mfA+pjA2smflC\n54brJTsSb1/1xAExxsvB2Leb6VcNWKRCaN6Z4gdWd1Qofi080LVWxE3MXhxHpTMj\nVf3KHhKI5DHJXlZPU/w4KXOlp99UH0vnx3EJzD07kI+nR2k7tfh8PxwFzn8g6hEC\nK9Q+HmzUUTxFh1M8eXi7IMRjLRJThVSl/Kbqr6cpeQKBgQC+410aRuHUteQIgJBQ\nbceAOjEh5MetByIEFdXLEgYspl1rSjN/JoIMUguyJ8KZGj5G3NaRaJklNOofYekI\nhIL3SBWZ9U58MJVMnSUVeBdazCu0k9HnOfOFrFJIDoRPBfjnP0UJmC+9ggoVayOX\nVW5psrxGiQXWiG7mho1bshSlNwKBgQC+yX1wFSF8gGsAb41iw60K9W+O5PLVWxAu\nYst8CQTY64RVctvAypWtNTb4nmIBe9aX9k5loe+uv8Zse/t1hgCVGR7n70EyT9Y+\nGNrGqYtVjtZQ+L+dAivrlUKsTDGqzWldUTg7gpOqkFaQbV0O11ytyKJKYXCpTrL2\nwib6V4X9zwKBgQCllTJAxfXFfxZUbblBm0iwKUpPXVX7+LEAHDS9F2B1wMZOeCod\nhLjQmSb+HlFGX6Zf79bMgZA+3xyrplHviorUmBns2AaB4d7Qe4wciHSx1WOgG427\n5uAgNy+Uw8rvhX24koB/Zx0aZT/7/lj8QCYr19hL0zZWNzkEDPl37gzMlwKBgQC3\noOsww8XVNSzH4JZupvOYhp53JHltTRaH7uL3YR7fQd++9qv4JYRmj793D8o4r17e\nKF1QiMpOoZpzs+lVNkK9Ps52YduYdys33WhEqc7H7JDuolya3Ao11xWzDCsJwGdX\nP+MltAo4sm/+1qQosrQrN96sRJjQ/ERYKIqnjTIUFQKBgQC6xaC5SB1UMepFbTVa\n2tuRuYwcrU/NxjeW85SAyeyv2dMg7S+Ot8mnA1Js2O0NHlczUZgRZalYkCuSUE78\nb6rIbezIW2azrw3tqAAPLsB+rhXvaUpICoybu+j6aCiqtZYsDx7zIj/FTD27Tpwx\nYfLx1Erqd3vM/LzOIaIOqlfETw==\n-----END PRIVATE KEY-----'
                        )"#,