From 52eebfce77b63d8289dad0c74513d1700ae00bf1 Mon Sep 17 00:00:00 2001
From: yihong <zouzou0208@gmail.com>
Date: Wed, 8 Jan 2025 17:44:54 +0800
Subject: [PATCH] feat: support `select session_user;` (#5313)

* feat: support `select session_user;`

This commit is part of support DBeaver that support function
select session_user like postgres did.

Signed-off-by: yihong0618 <zouzou0208@gmail.com>

* fix: lint problem

Signed-off-by: yihong0618 <zouzou0208@gmail.com>

* fix: address comments add tests

Signed-off-by: yihong0618 <zouzou0208@gmail.com>

---------

Signed-off-by: yihong0618 <zouzou0208@gmail.com>
---
 src/common/function/src/system.rs             |  5 ++--
 src/common/function/src/system/database.rs    | 28 +++++++++++++++++++
 .../common/system/pg_catalog.result           | 20 +++++++++++++
 .../standalone/common/system/pg_catalog.sql   |  8 ++++++
 4 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/src/common/function/src/system.rs b/src/common/function/src/system.rs
index b543fb9523..b2e6c41135 100644
--- a/src/common/function/src/system.rs
+++ b/src/common/function/src/system.rs
@@ -22,7 +22,7 @@ mod version;
 use std::sync::Arc;
 
 use build::BuildFunction;
-use database::{CurrentSchemaFunction, DatabaseFunction};
+use database::{CurrentSchemaFunction, DatabaseFunction, SessionUserFunction};
 use pg_catalog::PGCatalogFunction;
 use procedure_state::ProcedureStateFunction;
 use timezone::TimezoneFunction;
@@ -36,8 +36,9 @@ impl SystemFunction {
     pub fn register(registry: &FunctionRegistry) {
         registry.register(Arc::new(BuildFunction));
         registry.register(Arc::new(VersionFunction));
-        registry.register(Arc::new(DatabaseFunction));
         registry.register(Arc::new(CurrentSchemaFunction));
+        registry.register(Arc::new(DatabaseFunction));
+        registry.register(Arc::new(SessionUserFunction));
         registry.register(Arc::new(TimezoneFunction));
         registry.register_async(Arc::new(ProcedureStateFunction));
         PGCatalogFunction::register(registry);
diff --git a/src/common/function/src/system/database.rs b/src/common/function/src/system/database.rs
index fece862ee0..ae0985a084 100644
--- a/src/common/function/src/system/database.rs
+++ b/src/common/function/src/system/database.rs
@@ -28,9 +28,11 @@ pub struct DatabaseFunction;
 
 #[derive(Clone, Debug, Default)]
 pub struct CurrentSchemaFunction;
+pub struct SessionUserFunction;
 
 const DATABASE_FUNCTION_NAME: &str = "database";
 const CURRENT_SCHEMA_FUNCTION_NAME: &str = "current_schema";
+const SESSION_USER_FUNCTION_NAME: &str = "session_user";
 
 impl Function for DatabaseFunction {
     fn name(&self) -> &str {
@@ -72,6 +74,26 @@ impl Function for CurrentSchemaFunction {
     }
 }
 
+impl Function for SessionUserFunction {
+    fn name(&self) -> &str {
+        SESSION_USER_FUNCTION_NAME
+    }
+
+    fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
+        Ok(ConcreteDataType::string_datatype())
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::uniform(0, vec![], Volatility::Immutable)
+    }
+
+    fn eval(&self, func_ctx: FunctionContext, _columns: &[VectorRef]) -> Result<VectorRef> {
+        let user = func_ctx.query_ctx.current_user();
+
+        Ok(Arc::new(StringVector::from_slice(&[user.username()])) as _)
+    }
+}
+
 impl fmt::Display for DatabaseFunction {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "DATABASE")
@@ -84,6 +106,12 @@ impl fmt::Display for CurrentSchemaFunction {
     }
 }
 
+impl fmt::Display for SessionUserFunction {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "SESSION_USER")
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;
diff --git a/tests/cases/standalone/common/system/pg_catalog.result b/tests/cases/standalone/common/system/pg_catalog.result
index 4e3e217b2e..392ed7662a 100644
--- a/tests/cases/standalone/common/system/pg_catalog.result
+++ b/tests/cases/standalone/common/system/pg_catalog.result
@@ -3,6 +3,26 @@ create database pg_catalog;
 
 Error: 1004(InvalidArguments), Schema pg_catalog already exists
 
+-- session_user because session_user is based on the current user so is not null is for test 
+-- SQLNESS PROTOCOL POSTGRES
+SELECT session_user is not null;
+
++----------------------------+
+| session_user() IS NOT NULL |
++----------------------------+
+| t                          |
++----------------------------+
+
+-- session_user and current_schema
+-- SQLNESS PROTOCOL POSTGRES
+select current_schema();
+
++------------------+
+| current_schema() |
++------------------+
+| public           |
++------------------+
+
 -- make sure all the pg_catalog tables are only visible to postgres
 select * from pg_catalog.pg_class;
 
diff --git a/tests/cases/standalone/common/system/pg_catalog.sql b/tests/cases/standalone/common/system/pg_catalog.sql
index 712302254d..67a720a802 100644
--- a/tests/cases/standalone/common/system/pg_catalog.sql
+++ b/tests/cases/standalone/common/system/pg_catalog.sql
@@ -1,6 +1,14 @@
 -- should not able to create pg_catalog
 create database pg_catalog;
 
+-- session_user because session_user is based on the current user so is not null is for test 
+-- SQLNESS PROTOCOL POSTGRES
+SELECT session_user is not null;
+
+-- session_user and current_schema
+-- SQLNESS PROTOCOL POSTGRES
+select current_schema();
+
 -- make sure all the pg_catalog tables are only visible to postgres
 select * from pg_catalog.pg_class;
 select * from pg_catalog.pg_namespace;