From f37c9aa4aa96ae4472a986fa3bf8bf7b4400980c Mon Sep 17 00:00:00 2001
From: Pieter <pieter@chesedo.me>
Date: Tue, 30 Aug 2022 07:57:59 +0200
Subject: [PATCH] feat: respect RUST_LOG (#316)

* feat: respect RUST_LOG

* tests: get info logs for sqlx_pool
---
 service/src/logger.rs               | 45 +++++++++++++++++++++++++++--
 service/tests/integration/loader.rs |  3 ++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/service/src/logger.rs b/service/src/logger.rs
index 40009e62e..5e208b16e 100644
--- a/service/src/logger.rs
+++ b/service/src/logger.rs
@@ -1,5 +1,7 @@
+use std::{collections::HashMap, env, str::FromStr};
+
 use chrono::{DateTime, Utc};
-use log::{Level, Metadata, Record};
+use log::{Level, Metadata, ParseLevelError, Record};
 use shuttle_common::{DeploymentId, LogItem};
 use tokio::sync::mpsc::UnboundedSender;
 
@@ -13,17 +15,54 @@ pub struct Log {
 pub struct Logger {
     deployment_id: DeploymentId,
     tx: UnboundedSender<Log>,
+    filter: HashMap<String, Level>,
 }
 
 impl Logger {
     pub fn new(tx: UnboundedSender<Log>, deployment_id: DeploymentId) -> Self {
-        Self { tx, deployment_id }
+        let filter = if let Ok(rust_log) = env::var("RUST_LOG") {
+            let rust_log = rust_log
+                .split(',')
+                .map(|item| {
+                    // Try to get target and level if both are set
+                    if let Some((target, level)) = item.split_once('=') {
+                        Result::<(String, Level), ParseLevelError>::Ok((
+                            target.to_string(),
+                            Level::from_str(level)?,
+                        ))
+                    } else {
+                        // Ok only target or level is set, but which is it
+                        if let Ok(level) = Level::from_str(item) {
+                            Ok((String::new(), level))
+                        } else {
+                            Ok((item.to_string(), Level::Trace))
+                        }
+                    }
+                })
+                .filter_map(Result::ok);
+
+            HashMap::from_iter(rust_log)
+        } else {
+            HashMap::from([(String::new(), Level::Error)])
+        };
+
+        Self {
+            tx,
+            deployment_id,
+            filter,
+        }
     }
 }
 
 impl log::Log for Logger {
     fn enabled(&self, metadata: &Metadata) -> bool {
-        metadata.level() <= Level::Info
+        for (target, level) in self.filter.iter() {
+            if metadata.target().starts_with(target) && &metadata.level() <= level {
+                return true;
+            }
+        }
+
+        false
     }
 
     fn log(&self, record: &Record) {
diff --git a/service/tests/integration/loader.rs b/service/tests/integration/loader.rs
index a47066a57..34f68f443 100644
--- a/service/tests/integration/loader.rs
+++ b/service/tests/integration/loader.rs
@@ -107,6 +107,9 @@ async fn sleep() {
 async fn sqlx_pool() {
     let loader = build_so_create_loader(RESOURCES_PATH, "sqlx-pool").unwrap();
 
+    // Make sure we'll get a log entry
+    std::env::set_var("RUST_LOG", "info");
+
     // Don't initialize a pre-existing PostgresInstance here because the `PostgresInstance::wait_for_connectable()`
     // code has `awaits` and we want to make sure they do not block inside `Service::build()`.
     // At the same time we also want to test the PgPool is created on the correct runtime (ie does not cause a