Adding more examples of using SPIs with SRF (Set Returning Functions) (…

…#701)

* Adding more examples of using SPIs with SRF (Set Returning Functions)

* adding spi_srf as a workspace crate
* adding extra line and example to workflows
* Explicit add of TableIterator
* properly adding prelude
* Making test pass

.github/workflows/tests.yml

@@ -174,6 +174,9 @@ jobs:
     - name: Run spi example tests
       run: cargo test --package spi --features "pg$PG_VER" --no-default-features
 
+    - name: Run spi_srf example tests
+      run: cargo test --package spi_srf --features "pg$PG_VER" --no-default-features
+
     - name: Run srf example tests
       run: cargo test --package srf --features "pg$PG_VER" --no-default-features
 

Cargo.toml

@@ -26,6 +26,7 @@ members = [
     "pgx-examples/strings",
     "pgx-examples/triggers",
     "pgx-examples/versioned_so",
+    "pgx-examples/spi_srf",
 ]
 
 [profile.dev.build-override]

pgx-examples/spi_srf/.gitignore

@@ -0,0 +1,7 @@
+.DS_Store
+.idea/
+/target
+*.iml
+**/*.rs.bk
+Cargo.lock
+sql/spi-1.0.sql

pgx-examples/spi_srf/Cargo.toml

@@ -0,0 +1,34 @@
+[package]
+name = "spi_srf"
+version = "0.0.0"
+edition = "2021"
+rust-version = "1.58"
+
+[lib]
+crate-type = ["cdylib"]
+
+[features]
+default = ["pg13"]
+pg10 = ["pgx/pg10", "pgx-tests/pg10" ]
+pg11 = ["pgx/pg11", "pgx-tests/pg11" ]
+pg12 = ["pgx/pg12", "pgx-tests/pg12" ]
+pg13 = ["pgx/pg13", "pgx-tests/pg13" ]
+pg14 = ["pgx/pg14", "pgx-tests/pg14" ]
+pg_test = []
+
+[dependencies]
+pgx = { path = "../../pgx", default-features = false }
+
+[dev-dependencies]
+pgx-tests = { path = "../../pgx-tests" }
+
+# uncomment these if compiling outside of 'pgx'
+# [profile.dev]
+# panic = "unwind"
+# lto = "thin"
+
+# [profile.release]
+# panic = "unwind"
+# opt-level = 3
+# lto = "fat"
+# codegen-units = 1

pgx-examples/spi_srf/README.md

@@ -0,0 +1 @@
+Some examples of using SPI (Server Programming Interface) and SRF with pgx.

pgx-examples/spi_srf/spi_srf.control

@@ -0,0 +1,6 @@
+comment = 'spi_srf:  Created by pgx'
+default_version = '@CARGO_VERSION@'
+module_pathname = '$libdir/spi_srf'
+relocatable = false
+superuser = false
+schema = 'spi_srf'

pgx-examples/spi_srf/src/lib.rs

@@ -0,0 +1,144 @@
+/*
+Portions Copyright 2019-2021 ZomboDB, LLC.
+Portions Copyright 2021-2022 Technology Concepts & Design, Inc. <[email protected]>
+
+All rights reserved.
+
+Use of this source code is governed by the MIT license that can be found in the LICENSE file.
+*/
+use pgx::{prelude::*, IntoDatum, SpiTupleTable};
+
+pgx::pg_module_magic!();
+
+extension_sql!(
+    r#"   
+
+CREATE TABLE dog_daycare (
+    dog_name varchar(256),
+    dog_age int,
+    dog_breed varchar(256)
+);
+
+INSERT INTO dog_daycare(dog_name, dog_age, dog_breed) VALUES ('Fido', 3, 'Labrador');
+INSERT INTO dog_daycare(dog_name, dog_age, dog_breed) VALUES ('Spot', 5, 'Poodle');
+INSERT INTO dog_daycare(dog_name, dog_age, dog_breed) VALUES ('Rover', 7, 'Golden Retriever');
+INSERT INTO dog_daycare(dog_name, dog_age, dog_breed) VALUES ('Snoopy', 9, 'Beagle');
+INSERT INTO dog_daycare(dog_name, dog_age, dog_breed) VALUES ('Lassie', 11, 'Collie');
+INSERT INTO dog_daycare(dog_name, dog_age, dog_breed) VALUES ('Scooby', 13, 'Great Dane');
+INSERT INTO dog_daycare(dog_name, dog_age, dog_breed) VALUES ('Moomba', 15, 'Labrador');
+
+
+"#,
+    name = "create_dog_daycare_example_table",
+);
+
+#[pg_extern]
+fn calculate_human_years() -> TableIterator<
+    'static,
+    (
+        name!(dog_name, String),
+        name!(dog_age, i32),
+        name!(dog_breed, String),
+        name!(human_age, i32),
+    ),
+> {
+    /*
+        This function is a simple example of using SPI to return a set of rows
+        from a query. This query will return the same rows as the table, but
+        with an additional column that is the dog's age in human years.
+    */
+    let query = "SELECT * FROM spi_srf.dog_daycare;";
+
+    let mut results = Vec::new();
+
+    Spi::connect(|client| {
+        let mut tup_table: SpiTupleTable = client.select(query, None, None);
+
+        while let Some(row) = tup_table.next() {
+            let dog_name: String = row["dog_name"].value().unwrap();
+            let dog_age: i32 = row["dog_age"].value().unwrap();
+            let dog_breed: String = row["dog_breed"].value().unwrap();
+            let human_age: i32 = dog_age * 7;
+            results.push((dog_name, dog_age, dog_breed, human_age));
+        }
+
+        Ok(Some(()))
+    });
+
+    TableIterator::new(results.into_iter())
+}
+
+#[pg_extern]
+fn filter_by_breed(
+    breed: &str,
+) -> TableIterator<
+    'static,
+    (
+        name!(dog_name, Option<String>),
+        name!(dog_age, Option<i32>),
+        name!(dog_breed, Option<String>),
+    ),
+> {
+    /*
+        This function is a simple example of using SPI to return a set of rows
+        from a query. This query will return the records for the given breed.
+    */
+
+    let query = "SELECT * FROM spi_srf.dog_daycare WHERE dog_breed = $1;";
+    let args = vec![(PgBuiltInOids::TEXTOID.oid(), breed.into_datum())];
+
+    let mut results = Vec::new();
+
+    Spi::connect(|client| {
+        let mut tup_table: SpiTupleTable = client.select(query, None, Some(args));
+
+        while let Some(row) = tup_table.next() {
+            results.push((
+                row["dog_name"].value(),
+                row["dog_age"].value(),
+                row["dog_breed"].value(),
+            ));
+        }
+
+        Ok(Some(()))
+    });
+
+    TableIterator::new(results.into_iter())
+}
+
+#[cfg(any(test, feature = "pg_test"))]
+#[pg_schema]
+mod tests {
+    use crate::calculate_human_years;
+    use pgx::*;
+
+    #[pg_test]
+    fn test_calculate_human_years() {
+        let mut results: Vec<(String, i32, String, i32)> = Vec::new();
+
+        results.push(("Fido".to_string(), 3, "Labrador".to_string(), 21));
+        results.push(("Spot".to_string(), 5, "Poodle".to_string(), 35));
+        results.push(("Rover".to_string(), 7, "Golden Retriever".to_string(), 49));
+        results.push(("Snoopy".to_string(), 9, "Beagle".to_string(), 63));
+        results.push(("Lassie".to_string(), 11, "Collie".to_string(), 77));
+        results.push(("Scooby".to_string(), 13, "Great Dane".to_string(), 91));
+        results.push(("Moomba".to_string(), 15, "Labrador".to_string(), 105));
+        let func_results = calculate_human_years();
+
+        for (expected, actual) in results.iter().zip(func_results) {
+            assert_eq!(expected, &actual);
+        }
+    }
+}
+
+#[cfg(test)]
+pub mod pg_test {
+    pub fn setup(_options: Vec<&str>) {
+        // perform one-off initialization when the pg_test framework starts
+    }
+
+    pub fn postgresql_conf_options() -> Vec<&'static str> {
+        // return any postgresql.conf settings that are required for your tests
+        vec![]
+    }
+}