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>
This commit is contained in:
yihong
2025-01-08 17:44:54 +08:00
committed by GitHub
parent e18416a726
commit 52eebfce77
4 changed files with 59 additions and 2 deletions

View File

@@ -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);

View File

@@ -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;

View File

@@ -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;

View File

@@ -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;