support user input prometheus base url

tembo-io · Oct 6, 2023 · e3aa30b · e3aa30b
1 parent 1bc8b3d
commit e3aa30b
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 14 deletions.
diff --git a/README.md b/README.md
@@ -1,7 +1,7 @@
 ## Prometheus_fdw
-NOTE: ask the user to enter the link to prometheus
 NOTE: update step value
 NOTE: add a cron job to make this automatic
+NOTE: write tests
 
 ### Pre-requisistes
 
@@ -23,7 +23,9 @@ Create the server:
 
 ```
 create server my_prometheus_server
-  foreign data wrapper prometheus_wrapper;
+  foreign data wrapper prometheus_wrapper
+  options (
+    base_url '<base prometheus url>');
 ```
 
 Create Foreign Table:
@@ -118,11 +120,4 @@ INNER JOIN
 ON
     ml.metric_name = mlab.name AND ml.metric_labels = mlab.labels
 ON CONFLICT (id, time) DO NOTHING;
-```
-
-
-
-
-
-
-
+```
diff --git a/src/lib.rs b/src/lib.rs
@@ -64,14 +64,13 @@ fn resp_to_rows(obj: &str, resp: &JsonValue) -> Vec<Row> {
 
 pub(crate) struct PrometheusFdw {
     rt: Runtime,
+    base_url: Option<String>,
     client: Option<Client>,
     scan_result: Option<Vec<Row>>,
     tgt_cols: Vec<Column>,
 }
 
 impl PrometheusFdw {
-    const DEFAULT_BASE_URL: &'static str = "https://prometheus-data-1.use1.prod.plat.cdb-svc.com";
-
     fn value_to_promql_string(value: &supabase_wrappers::interface::Value) -> String {
         match value {
             supabase_wrappers::interface::Value::Cell(cell) => match cell {
@@ -111,7 +110,7 @@ impl PrometheusFdw {
                     let upper_timestamp = Self::value_to_promql_string(&upper_timestamp.value);
                     let ret = format!(
                         "{}/api/v1/query_range?query={}&start={}&end={}&step=10m",
-                        Self::DEFAULT_BASE_URL,
+                        self.base_url.as_ref().unwrap(),
                         metric_name,
                         lower_timestamp,
                         upper_timestamp
@@ -131,13 +130,24 @@ impl PrometheusFdw {
 }
 
 impl ForeignDataWrapper for PrometheusFdw {
-    fn new(_options: &HashMap<String, String>) -> Self {
+    fn new(options: &HashMap<String, String>) -> Self {
         let mut ret = Self {
             rt: create_async_runtime(),
+            base_url: None,
             client: None,
             tgt_cols: Vec::new(),
             scan_result: None,
         };
+
+        let base_url = if let Some(prom_url) = options.get("base_url") {
+            prom_url.to_owned()
+        } else {
+            warning!("Cannot find prometheus base url in options");
+            let prom_url = env::var("PROMETHEUS_BASE_URL").unwrap();
+            prom_url
+        };
+
+        ret.base_url = Some(base_url);
         ret.client = Some(reqwest::Client::new());
 
         ret