feat: add set search_path to 'xxx' for pg (#5342)

* feat: add set search_path to 'xxx' for pg

Signed-off-by: yihong0618 <[email protected]>

* fix: address comments

Signed-off-by: yihong0618 <[email protected]>

---------

Signed-off-by: yihong0618 <[email protected]>

Author: yihong0618

src/operator/src/statement.rs

@@ -72,7 +72,9 @@ use table::table_name::TableName;
 use table::table_reference::TableReference;
 use table::TableRef;
 
-use self::set::{set_bytea_output, set_datestyle, set_timezone, validate_client_encoding};
+use self::set::{
+    set_bytea_output, set_datestyle, set_search_path, set_timezone, validate_client_encoding,
+};
 use crate::error::{
     self, CatalogSnafu, ExecLogicalPlanSnafu, ExternalSnafu, InvalidSqlSnafu, NotSupportedSnafu,
     PlanStatementSnafu, Result, SchemaNotFoundSnafu, StatementTimeoutSnafu,
@@ -408,6 +410,16 @@ impl StatementExecutor {
                     .fail();
                 }
             }
+            "SEARCH_PATH" => {
+                if query_ctx.channel() == Channel::Postgres {
+                    set_search_path(set_var.value, query_ctx)?
+                } else {
+                    return NotSupportedSnafu {
+                        feat: format!("Unsupported set variable {}", var_name),
+                    }
+                    .fail();
+                }
+            }
             _ => {
                 // for postgres, we give unknown SET statements a warning with
                 //  success, this is prevent the SET call becoming a blocker

src/operator/src/statement/set.rs

@@ -81,6 +81,26 @@ pub fn set_bytea_output(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
     Ok(())
 }
 
+pub fn set_search_path(exprs: Vec<Expr>, ctx: QueryContextRef) -> Result<()> {
+    let search_expr = exprs.first().context(NotSupportedSnafu {
+        feat: "No search path find in set variable statement",
+    })?;
+    match search_expr {
+        Expr::Value(Value::SingleQuotedString(search_path))
+        | Expr::Value(Value::DoubleQuotedString(search_path)) => {
+            ctx.set_current_schema(&search_path.clone());
+            Ok(())
+        }
+        expr => NotSupportedSnafu {
+            feat: format!(
+                "Unsupported search path expr {} in set variable statement",
+                expr
+            ),
+        }
+        .fail(),
+    }
+}
+
 pub fn validate_client_encoding(set: SetVariables) -> Result<()> {
     let Some((encoding, [])) = set.value.split_first() else {
         return InvalidSqlSnafu {

tests/cases/standalone/common/system/pg_catalog.result

@@ -33,6 +33,34 @@ show search_path;
 | public      |
 +-------------+
 
+-- set search_path for pg using schema for now FIXME when support real search_path
+create database test;
+
+Affected Rows: 1
+
+-- SQLNESS PROTOCOL POSTGRES
+set search_path to 'test';
+
+Affected Rows: 0
+
+drop database test;
+
+Affected Rows: 0
+
+-- SQLNESS PROTOCOL POSTGRES
+set search_path to 'public';
+
+Affected Rows: 0
+
+-- 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;
 

tests/cases/standalone/common/system/pg_catalog.sql

@@ -13,6 +13,17 @@ select current_schema();
 -- SQLNESS PROTOCOL POSTGRES
 show search_path;
 
+-- set search_path for pg using schema for now FIXME when support real search_path
+create database test;
+-- SQLNESS PROTOCOL POSTGRES
+set search_path to 'test';
+drop database test;
+-- SQLNESS PROTOCOL POSTGRES
+set search_path to 'public';
+
+-- 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;