From 50900bc7cc6bca805502ad46ae73d17d1e763f31 Mon Sep 17 00:00:00 2001
From: Lev Kokotov <lev.kokotov@gmail.com>
Date: Fri, 10 Jan 2025 16:36:42 -0800
Subject: [PATCH] improve example + add CI badge

---
 README.md                          |  1 +
 examples/routing-plugin/src/lib.rs | 39 ++++++++++++++++++++----------
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index 885a4e4..7a64373 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,7 @@
 # pgDog - PostgreSQL Load Balancer
 
 [![Documentation](https://img.shields.io/badge/documentation-blue?style=flat)](https://pgdog.dev)
+![CI](https://github.com/levkk/pgdog/actions/workflows/ci.yml/badge.svg)
 
 pgDog is a PostgreSQL pooler, load balancer and sharding proxy, written in Rust.
 Spiritual successor to [pgcat](https://github.com/levkk/pgcat), pgDog comes with a lot of
diff --git a/examples/routing-plugin/src/lib.rs b/examples/routing-plugin/src/lib.rs
index ac4bc64..ac6c823 100644
--- a/examples/routing-plugin/src/lib.rs
+++ b/examples/routing-plugin/src/lib.rs
@@ -4,19 +4,32 @@ use pgdog_plugin::*;
 /// Route query.
 #[no_mangle]
 pub extern "C" fn pgdog_route_query(input: Input) -> Output {
-    if let Some(query) = input.query() {
-        let _id = if let Some(id) = query.parameter(0) {
-            if let Some(id) = id.as_str() {
-                id.parse::<i64>().ok()
-            } else if let Ok(id) = id.as_bytes().try_into() {
-                Some(i64::from_be_bytes(id))
-            } else {
-                None
+    let is_read = input
+        .query()
+        .map(|query| query.query().to_lowercase().trim().starts_with("select"))
+        .unwrap_or(false);
+
+    // This is just an example of extracing a parameter from
+    // the query. In the future, we'll use this to shard transactions.
+    let _parameter = input.query().map(|query| {
+        query.parameter(0).map(|parameter| {
+            let id = parameter.as_str().map(|str| str.parse::<i64>());
+            match id {
+                Some(Ok(id)) => id,
+                _ => i64::from_be_bytes(
+                    parameter
+                        .as_bytes()
+                        .try_into()
+                        .map(|bytes| bytes)
+                        .unwrap_or([0u8; 8]),
+                ),
             }
-        } else {
-            None
-        };
-    }
+        })
+    });
 
-    Output::skip()
+    if is_read {
+        Output::forward(Route::read_any())
+    } else {
+        Output::forward(Route::write_any())
+    }
 }