diff --git a/Cargo.lock b/Cargo.lock index b383fcf3af688..041e746d8672e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8707,6 +8707,7 @@ dependencies = [ "pretty_assertions", "prometheus", "rand", + "regex", "risingwave_batch", "risingwave_common", "risingwave_common_service", diff --git a/e2e_test/batch/catalog/sysinfo.slt.part b/e2e_test/batch/catalog/sysinfo.slt.part index 521b42326e267..0b9520f2e3158 100644 --- a/e2e_test/batch/catalog/sysinfo.slt.part +++ b/e2e_test/batch/catalog/sysinfo.slt.part @@ -34,4 +34,36 @@ statement ok drop index tab_idx; statement ok -drop table tab; \ No newline at end of file +drop table tab; + +statement ok +create table tab(a int, b int, c int, d int); + +statement ok +CREATE INDEX tab_idx ON tab (a, (b + c + (1 + 1))) include (d); + +query T +select pg_get_indexdef('tab_idx'::regclass), pg_get_indexdef('tab_idx'::regclass, 0, true); +---- +CREATE INDEX tab_idx ON tab(a, (b + c + (1 + 1))) INCLUDE(d) CREATE INDEX tab_idx ON tab(a, (b + c + (1 + 1))) INCLUDE(d) + +query T +select pg_get_indexdef('tab_idx'::regclass, 1, true), pg_get_indexdef('tab_idx'::regclass, 2, true); +---- +a ((b + c) + (1:Int32 + 1:Int32)) + +query T +select pg_get_indexdef('tab_idx'::regclass, 3, true); +---- +d + +query T +select pg_get_indexdef('tab_idx'::regclass, -1, true), pg_get_indexdef('tab_idx'::regclass, 4, true); +---- +(empty) (empty) + +statement ok +drop index tab_idx; + +statement ok +drop table tab; diff --git a/src/frontend/Cargo.toml b/src/frontend/Cargo.toml index 73df1156cd035..7f719590ed9a2 100644 --- a/src/frontend/Cargo.toml +++ b/src/frontend/Cargo.toml @@ -51,6 +51,7 @@ pretty-xmlish = "0.1.13" pretty_assertions = "1" prometheus = { version = "0.13", features = ["process"] } rand = "0.8" +regex = "1" risingwave_batch = { workspace = true } risingwave_common = { workspace = true } risingwave_common_service = { workspace = true } diff --git a/src/frontend/src/catalog/index_catalog.rs b/src/frontend/src/catalog/index_catalog.rs index ca4b4036332d3..255c6af24810c 100644 --- a/src/frontend/src/catalog/index_catalog.rs +++ b/src/frontend/src/catalog/index_catalog.rs @@ -18,6 +18,7 @@ use std::sync::Arc; use educe::Educe; use itertools::Itertools; +use regex::Regex; use risingwave_common::catalog::IndexId; use risingwave_common::util::epoch::Epoch; use risingwave_common::util::sort_util::ColumnOrder; @@ -188,6 +189,24 @@ impl IndexCatalog { } } + pub fn get_column_def(&self, column_no: usize) -> String { + if let Some(col) = self.index_table.columns.get(column_no) { + if col.is_hidden { + return String::new(); + } + } else { + return String::new(); + } + let col = format!("{:?}", self.index_item[column_no]); + let column_names = self + .primary_table + .columns + .iter() + .map(|x| x.name()) + .collect_vec(); + replace_with_column_names(&column_names, &col) + } + pub fn display(&self) -> IndexDisplay { let index_table = self.index_table.clone(); let index_columns_with_ordering = index_table @@ -240,3 +259,49 @@ impl OwnedByUserCatalog for IndexCatalog { self.index_table.owner } } + +fn replace_with_column_names(v: &[&str], s: &str) -> String { + let re = Regex::new(r"\$(\d+)").unwrap(); + let mut result = String::new(); + + let mut last_end = 0; + for cap in re.captures_iter(s) { + let start = cap.get(0).unwrap().start(); + let end = cap.get(0).unwrap().end(); + let index: usize = cap[1].parse().unwrap(); + + result.push_str(&s[last_end..start]); + + match v.get(index) { + Some(value) => result.push_str(value), + None => result.push_str(&s[start..end]), + } + + last_end = end; + } + + result.push_str(&s[last_end..]); + + result +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_replace_with_column_names() { + { + let v = vec!["a", "b", "c"]; + let s = "select * from $0, $4 where $1 = $2"; + let result = replace_with_column_names(&v, s); + assert_eq!(result, "select * from a, $4 where b = c"); + } + { + let v = vec!["a", "b", "c"]; + let s = "$0"; + let result = replace_with_column_names(&v, s); + assert_eq!(result, "a"); + } + } +} diff --git a/src/frontend/src/catalog/root_catalog.rs b/src/frontend/src/catalog/root_catalog.rs index d54497c64e8d7..51779bc6d0043 100644 --- a/src/frontend/src/catalog/root_catalog.rs +++ b/src/frontend/src/catalog/root_catalog.rs @@ -685,15 +685,15 @@ impl Catalog { .ok_or_else(|| CatalogError::NotFound("index", index_name.to_string())) } - pub fn get_table_by_index_id( + pub fn get_index_by_id( &self, db_name: &str, index_id: u32, - ) -> CatalogResult<&Arc> { + ) -> CatalogResult<&Arc> { let index_id = IndexId::from(index_id); for schema in self.get_database_by_name(db_name)?.iter_schemas() { if let Some(index) = schema.get_index_by_id(&index_id) { - return Ok(&index.index_table); + return Ok(index); } } Err(CatalogError::NotFound("index", index_id.to_string())) diff --git a/src/frontend/src/expr/function_impl/pg_get_indexdef.rs b/src/frontend/src/expr/function_impl/pg_get_indexdef.rs index 01e4ea1d64b86..9b87b2cbb2bbc 100644 --- a/src/frontend/src/expr/function_impl/pg_get_indexdef.rs +++ b/src/frontend/src/expr/function_impl/pg_get_indexdef.rs @@ -22,7 +22,17 @@ use crate::catalog::CatalogReader; #[function("pg_get_indexdef(int4) -> varchar")] fn pg_get_indexdef(oid: i32, writer: &mut impl Write) -> Result<()> { - pg_get_indexdef_impl_captured(oid, writer) + pg_get_indexdef_impl_captured(oid, 0, writer) +} + +#[function("pg_get_indexdef(int4, int4, boolean) -> varchar")] +fn pg_get_indexdef_col( + oid: i32, + column_no: i32, + _pretty_bool: bool, + writer: &mut impl Write, +) -> Result<()> { + pg_get_indexdef_impl_captured(oid, column_no, writer) } #[capture_context(CATALOG_READER, DB_NAME)] @@ -30,20 +40,29 @@ fn pg_get_indexdef_impl( catalog: &CatalogReader, db_name: &str, oid: i32, + column_no: i32, writer: &mut impl Write, ) -> Result<()> { - write!( - writer, - "{}", + let ans = if column_no == 0 { catalog .read_guard() - .get_table_by_index_id(db_name, oid as u32) + .get_index_by_id(db_name, oid as u32) .map_err(|e| ExprError::InvalidParam { name: "oid", reason: e.to_report_string().into(), })? + .index_table .create_sql() - ) - .unwrap(); + } else { + catalog + .read_guard() + .get_index_by_id(db_name, oid as u32) + .map_err(|e| ExprError::InvalidParam { + name: "oid", + reason: e.to_report_string().into(), + })? + .get_column_def(column_no as usize - 1) + }; + write!(writer, "{}", ans).unwrap(); Ok(()) }