Skip to content

Commit c206c90

Browse files
MazterQyouMazterQyou
authored
fix(cubesql): Improve DBeaver compatibility (#9769)
1 parent 185db54 commit c206c90

15 files changed

+325
-24
lines changed

‎rust/cubesql/cubesql/src/compile/engine/context_postgresql.rs

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,9 @@ use super::information_schema::postgres::{
2323
PgCatalogIndexProvider, PgCatalogInheritsProvider, PgCatalogMatviewsProvider,
2424
PgCatalogNamespaceProvider, PgCatalogPartitionedTableProvider, PgCatalogProcProvider,
2525
PgCatalogRangeProvider, PgCatalogRolesProvider, PgCatalogSequenceProvider,
26-
PgCatalogSettingsProvider, PgCatalogStatActivityProvider, PgCatalogStatUserTablesProvider,
27-
PgCatalogStatioUserTablesProvider, PgCatalogStatsProvider, PgCatalogTableProvider,
28-
PgCatalogTypeProvider, PgCatalogUserProvider, PgCatalogViewsProvider,
26+
PgCatalogSettingsProvider, PgCatalogShdescriptionProvider, PgCatalogStatActivityProvider,
27+
PgCatalogStatUserTablesProvider, PgCatalogStatioUserTablesProvider, PgCatalogStatsProvider,
28+
PgCatalogTableProvider, PgCatalogTypeProvider, PgCatalogUserProvider, PgCatalogViewsProvider,
2929
PgPreparedStatementsProvider,
3030
};
3131
use crate::{
@@ -136,6 +136,8 @@ impl DatabaseProtocol {
136136
"pg_catalog.pg_views".to_string()
137137
} else if let Some(_) = any.downcast_ref::<PgCatalogStatUserTablesProvider>() {
138138
"pg_catalog.pg_stat_user_tables".to_string()
139+
} else if let Some(_) = any.downcast_ref::<PgCatalogShdescriptionProvider>() {
140+
"pg_catalog.pg_shdescription".to_string()
139141
} else if let Some(_) = any.downcast_ref::<RedshiftPgExternalSchemaProvider>() {
140142
"pg_catalog.pg_external_schema".to_string()
141143
} else if let Some(_) = any.downcast_ref::<RedshiftSvvTablesTableProvider>() {
@@ -401,6 +403,7 @@ impl DatabaseProtocol {
401403
&context.meta.tables,
402404
)))
403405
}
406+
"pg_shdescription" => return Some(Arc::new(PgCatalogShdescriptionProvider::new())),
404407
"pg_external_schema" => {
405408
return Some(Arc::new(RedshiftPgExternalSchemaProvider::new()))
406409
}

‎rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ mod pg_range;
3434
mod pg_roles;
3535
mod pg_sequence;
3636
mod pg_settings;
37+
mod pg_shdescription;
3738
mod pg_stat_activity;
3839
mod pg_stat_user_tables;
3940
mod pg_statio_user_tables;
@@ -69,6 +70,7 @@ pub use pg_range::*;
6970
pub use pg_roles::*;
7071
pub use pg_sequence::*;
7172
pub use pg_settings::*;
73+
pub use pg_shdescription::*;
7274
pub use pg_stat_activity::*;
7375
pub use pg_stat_user_tables::*;
7476
pub use pg_statio_user_tables::*;

‎rust/cubesql/cubesql/src/compile/engine/information_schema/postgres/pg_shdescription.rs

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
use std::{any::Any, sync::Arc};
2+
3+
use async_trait::async_trait;
4+
5+
use datafusion::{
6+
arrow::{
7+
array::{Array, ArrayRef, StringBuilder, UInt32Builder},
8+
datatypes::{DataType, Field, Schema, SchemaRef},
9+
record_batch::RecordBatch,
10+
},
11+
datasource::{datasource::TableProviderFilterPushDown, TableProvider, TableType},
12+
error::DataFusionError,
13+
logical_plan::Expr,
14+
physical_plan::{memory::MemoryExec, ExecutionPlan},
15+
};
16+
17+
struct PgCatalogShdescriptionBuilder {
18+
objoid: UInt32Builder,
19+
classoid: UInt32Builder,
20+
description: StringBuilder,
21+
}
22+
23+
impl PgCatalogShdescriptionBuilder {
24+
fn new(capacity: usize) -> Self {
25+
Self {
26+
objoid: UInt32Builder::new(capacity),
27+
classoid: UInt32Builder::new(capacity),
28+
description: StringBuilder::new(capacity),
29+
}
30+
}
31+
32+
fn finish(mut self) -> Vec<Arc<dyn Array>> {
33+
let columns: Vec<Arc<dyn Array>> = vec![
34+
Arc::new(self.objoid.finish()),
35+
Arc::new(self.classoid.finish()),
36+
Arc::new(self.description.finish()),
37+
];
38+
39+
columns
40+
}
41+
}
42+
43+
pub struct PgCatalogShdescriptionProvider {
44+
data: Arc<Vec<ArrayRef>>,
45+
}
46+
47+
// https://www.postgresql.org/docs/14/catalog-pg-shdescription.html
48+
impl PgCatalogShdescriptionProvider {
49+
pub fn new() -> Self {
50+
let builder = PgCatalogShdescriptionBuilder::new(0);
51+
52+
Self {
53+
data: Arc::new(builder.finish()),
54+
}
55+
}
56+
}
57+
58+
#[async_trait]
59+
impl TableProvider for PgCatalogShdescriptionProvider {
60+
fn as_any(&self) -> &dyn Any {
61+
self
62+
}
63+
64+
fn table_type(&self) -> TableType {
65+
TableType::View
66+
}
67+
68+
fn schema(&self) -> SchemaRef {
69+
Arc::new(Schema::new(vec![
70+
Field::new("objoid", DataType::UInt32, false),
71+
Field::new("classoid", DataType::UInt32, false),
72+
Field::new("description", DataType::Utf8, false),
73+
]))
74+
}
75+
76+
async fn scan(
77+
&self,
78+
projection: &Option<Vec<usize>>,
79+
_filters: &[Expr],
80+
_limit: Option<usize>,
81+
) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
82+
let batch = RecordBatch::try_new(self.schema(), self.data.to_vec())?;
83+
84+
Ok(Arc::new(MemoryExec::try_new(
85+
&[vec![batch]],
86+
self.schema(),
87+
projection.clone(),
88+
)?))
89+
}
90+
91+
fn supports_filter_pushdown(
92+
&self,
93+
_filter: &Expr,
94+
) -> Result<TableProviderFilterPushDown, DataFusionError> {
95+
Ok(TableProviderFilterPushDown::Unsupported)
96+
}
97+
}

‎rust/cubesql/cubesql/src/compile/engine/udf/common.rs

Lines changed: 53 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3999,6 +3999,59 @@ pub fn create_inet_server_addr_udf() -> ScalarUDF {
39993999
)
40004000
}
40014001

4002+
pub fn create_pg_get_partkeydef_udf() -> ScalarUDF {
4003+
let fun = make_scalar_function(move |args: &[ArrayRef]| {
4004+
let table_oids = downcast_primitive_arg!(args[0], "table_oid", OidType);
4005+
4006+
let result = table_oids
4007+
.iter()
4008+
.map(|_| None::<String>)
4009+
.collect::<StringArray>();
4010+
4011+
Ok(Arc::new(result))
4012+
});
4013+
4014+
create_udf(
4015+
"pg_get_partkeydef",
4016+
vec![DataType::UInt32],
4017+
Arc::new(DataType::Utf8),
4018+
Volatility::Immutable,
4019+
fun,
4020+
)
4021+
}
4022+
4023+
pub fn create_pg_relation_size_udf() -> ScalarUDF {
4024+
let fun = make_scalar_function(move |args: &[ArrayRef]| {
4025+
assert!(args.len() == 1);
4026+
4027+
let relids = downcast_primitive_arg!(args[0], "relid", OidType);
4028+
4029+
// 8192 is the lowest size for a table that has at least one column
4030+
// TODO: check if the requested table actually exists
4031+
let result = relids
4032+
.iter()
4033+
.map(|relid| relid.map(|_| 8192))
4034+
.collect::<PrimitiveArray<Int64Type>>();
4035+
4036+
Ok(Arc::new(result))
4037+
});
4038+
4039+
let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(DataType::Int64)));
4040+
4041+
ScalarUDF::new(
4042+
"pg_relation_size",
4043+
&Signature::one_of(
4044+
vec![
4045+
TypeSignature::Exact(vec![DataType::UInt32]),
4046+
TypeSignature::Exact(vec![DataType::UInt32, DataType::Utf8]),
4047+
],
4048+
Volatility::Immutable,
4049+
),
4050+
&return_type,
4051+
&fun,
4052+
)
4053+
}
4054+
40024055
pub fn register_fun_stubs(mut ctx: SessionContext) -> SessionContext {
40034056
macro_rules! register_fun_stub {
40044057
($FTYP:ident, $NAME:expr, argc=$ARGC:expr $(, rettyp=$RETTYP:ident)? $(, vol=$VOL:ident)?) => {
@@ -4863,13 +4916,6 @@ pub fn register_fun_stubs(mut ctx: SessionContext) -> SessionContext {
48634916
rettyp = Utf8,
48644917
vol = Volatile
48654918
);
4866-
register_fun_stub!(
4867-
udf,
4868-
"pg_relation_size",
4869-
tsigs = [[Regclass], [Regclass, Utf8],],
4870-
rettyp = Int64,
4871-
vol = Volatile
4872-
);
48734919
register_fun_stub!(
48744920
udf,
48754921
"pg_reload_conf",

‎rust/cubesql/cubesql/src/compile/mod.rs

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5722,6 +5722,20 @@ ORDER BY
57225722
Ok(())
57235723
}
57245724

5725+
#[tokio::test]
5726+
async fn test_pgcatalog_pgshdescription_postgres() -> Result<(), CubeError> {
5727+
insta::assert_snapshot!(
5728+
"pgcatalog_pgshdescription_postgres",
5729+
execute_query(
5730+
"SELECT * FROM pg_catalog.pg_shdescription".to_string(),
5731+
DatabaseProtocol::PostgreSQL
5732+
)
5733+
.await?
5734+
);
5735+
5736+
Ok(())
5737+
}
5738+
57255739
#[tokio::test]
57265740
async fn test_constraint_column_usage_postgres() -> Result<(), CubeError> {
57275741
insta::assert_snapshot!(

‎rust/cubesql/cubesql/src/compile/parser.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,12 @@ pub fn parse_sql_to_statements(
7373
"SELECT n.oid,n.*,d.description FROM",
7474
"SELECT n.oid as _oid,n.*,d.description FROM",
7575
);
76+
let query = query.replace("SELECT c.oid,c.*,", "SELECT c.oid as _oid,c.*,");
77+
let query = query.replace("SELECT a.oid,a.*,", "SELECT a.oid as _oid,a.*,");
78+
let query = query.replace(
79+
"LEFT OUTER JOIN pg_depend dep on dep.refobjid = a.attrelid AND dep.deptype = 'i' and dep.refobjsubid = a.attnum and dep.classid = dep.refclassid",
80+
"LEFT OUTER JOIN pg_depend dep on dep.refobjid = a.attrelid AND dep.deptype = 'i' and dep.refobjsubid = a.attnum",
81+
);
7682

7783
// TODO Superset introspection: LEFT JOIN by ANY() is not supported
7884
let query = query.replace(

‎rust/cubesql/cubesql/src/compile/query_engine.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -516,6 +516,8 @@ impl QueryEngine for SqlQueryEngine {
516516
ctx.register_udf(create_pg_get_indexdef_udf());
517517
ctx.register_udf(create_inet_server_addr_udf());
518518
ctx.register_udf(create_age_udf());
519+
ctx.register_udf(create_pg_get_partkeydef_udf());
520+
ctx.register_udf(create_pg_relation_size_udf());
519521

520522
// udaf
521523
ctx.register_udaf(create_measure_udaf());

‎rust/cubesql/cubesql/src/compile/snapshots/cubesql__compile__tests__pgcatalog_pgshdescription_postgres.snap

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
---
2+
source: cubesql/src/compile/mod.rs
3+
expression: "execute_query(\"SELECT * FROM pg_catalog.pg_shdescription\".to_string(),\nDatabaseProtocol::PostgreSQL).await?"
4+
---
5+
+--------+----------+-------------+
6+
| objoid | classoid | description |
7+
+--------+----------+-------------+
8+
+--------+----------+-------------+

‎rust/cubesql/cubesql/src/compile/test/snapshots/cubesql__compile__test__test_introspection__dbeaver_introspection_columns.snap

10 KB
Binary file not shown.

‎rust/cubesql/cubesql/src/compile/test/snapshots/cubesql__compile__test__test_introspection__dbeaver_introspection_tables.snap

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
---
2+
source: cubesql/src/compile/test/test_introspection.rs
3+
expression: "execute_query(r#\"\n select c.oid,pg_catalog.pg_total_relation_size(c.oid) as total_rel_size,pg_catalog.pg_relation_size(c.oid) as rel_size\n FROM pg_class c\n WHERE c.relnamespace=2200\n ORDER BY c.oid\n \"#.to_string(),\nDatabaseProtocol::PostgreSQL).await?"
4+
---
5+
+-------+----------------+----------+
6+
| oid | total_rel_size | rel_size |
7+
+-------+----------------+----------+
8+
| 18000 | 8192 | 8192 |
9+
| 18020 | 8192 | 8192 |
10+
| 18030 | 8192 | 8192 |
11+
| 18036 | 8192 | 8192 |
12+
| 18246 | 8192 | 8192 |
13+
+-------+----------------+----------+

0 commit comments

Comments
 (0)