Skip to content

Commit b6e6400

Browse files
committed
feat(cubesql): Allow to bind float64 (support in pg-srv)
1 parent 195402f commit b6e6400

File tree

4 files changed

+248
-8
lines changed

4 files changed

+248
-8
lines changed

rust/cubesql/CLAUDE.md

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
# CLAUDE.md
2+
3+
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
4+
5+
## Repository Overview
6+
7+
CubeSQL is a SQL proxy server that enables SQL-based access to Cube.js semantic layer. It emulates the PostgreSQL wire protocol, allowing standard SQL clients and BI tools to query Cube.js deployments as if they were traditional databases. Note: MySQL protocol support has been deprecated and is no longer available.
8+
9+
This is a Rust workspace containing three crates:
10+
- **cubesql**: Main SQL proxy server with query compilation and protocol emulation
11+
- **cubeclient**: Rust client library for Cube.js API communication
12+
- **pg-srv**: PostgreSQL wire protocol server implementation
13+
14+
## Development Commands
15+
16+
### Prerequisites
17+
```bash
18+
# Install required Rust toolchain (1.84.1)
19+
rustup update
20+
21+
# Install snapshot testing tool
22+
cargo install cargo-insta
23+
```
24+
25+
### Core Build Commands
26+
```bash
27+
# Build all workspace members
28+
cargo build
29+
30+
# Build release version
31+
cargo build --release
32+
33+
# Format code
34+
cargo fmt
35+
36+
# Run linting (note: many clippy rules are disabled)
37+
cargo clippy
38+
```
39+
40+
### Running CubeSQL Server
41+
```bash
42+
# Run with required environment variables
43+
CUBESQL_CUBE_URL=$CUBE_URL/cubejs-api \
44+
CUBESQL_CUBE_TOKEN=$CUBE_TOKEN \
45+
CUBESQL_LOG_LEVEL=debug \
46+
CUBESQL_BIND_ADDR=0.0.0.0:4444 \
47+
cargo run --bin cubesqld
48+
49+
# Connect via PostgreSQL client
50+
psql -h 127.0.0.1 -p 4444 -U root
51+
```
52+
53+
### Testing Commands
54+
```bash
55+
# Run all unit tests
56+
cargo test
57+
58+
# Run specific test module
59+
cargo test test_introspection
60+
cargo test test_udfs
61+
62+
# Run integration tests (requires Cube.js instance)
63+
cargo test --test e2e
64+
65+
# Review snapshot test changes
66+
cargo insta review
67+
68+
# Run benchmarks
69+
cargo bench
70+
```
71+
72+
## Architecture Overview
73+
74+
### Query Processing Pipeline
75+
1. **Protocol Layer**: Accepts PostgreSQL wire protocol connections
76+
2. **SQL Parser**: Modified sqlparser-rs parses incoming SQL queries
77+
3. **Query Rewriter**: egg-based rewrite engine transforms SQL to Cube.js queries
78+
4. **Compilation**: Generates Cube.js REST API calls or DataFusion execution plans
79+
5. **Execution**: DataFusion executes queries or proxies to Cube.js
80+
6. **Result Formatting**: Converts results back to wire protocol format
81+
82+
### Key Components
83+
84+
#### cubesql crate structure:
85+
- **`/compile`**: SQL compilation and query planning
86+
- `/engine`: DataFusion integration and query execution
87+
- `/rewrite`: egg-based query optimization rules
88+
- **`/sql`**: Database protocol implementations
89+
- `/postgres`: PostgreSQL system catalog emulation
90+
- `/database_variables`: Variable system for PostgreSQL protocol
91+
- **`/transport`**: Network transport and session management
92+
- **`/config`**: Configuration and service initialization
93+
94+
#### Testing Approach:
95+
- **Unit Tests**: Inline tests in source files using `#[cfg(test)]`
96+
- **Integration Tests**: End-to-end tests in `/e2e` directory
97+
- **Snapshot Tests**: Extensive use of `insta` for SQL compilation snapshots
98+
- **BI Tool Tests**: Compatibility tests for Metabase, Tableau, PowerBI, etc.
99+
100+
### Important Implementation Details
101+
102+
1. **DataFusion Integration**: Uses forked Apache Arrow DataFusion for query execution
103+
2. **Rewrite Rules**: Complex SQL transformations using egg e-graph library
104+
3. **Protocol Emulation**: Implements enough of PostgreSQL protocol for BI tools
105+
4. **System Catalogs**: Emulates pg_catalog (PostgreSQL)
106+
5. **Variable Handling**: Supports SET/SHOW commands for protocol compatibility
107+
108+
## Common Development Tasks
109+
110+
### Adding New SQL Support
111+
1. Add parsing support in `/compile/parser`
112+
2. Create rewrite rules in `/compile/rewrite/rules`
113+
3. Add tests with snapshot expectations
114+
4. Update protocol-specific handling if needed
115+
116+
### Debugging Query Compilation
117+
```bash
118+
# Enable detailed logging
119+
CUBESQL_LOG_LEVEL=trace cargo run --bin cubesqld
120+
121+
# Check rewrite traces in logs
122+
# Look for "Rewrite" entries showing transformation steps
123+
```
124+
125+
### Working with Snapshots
126+
```bash
127+
# After making changes that affect SQL compilation
128+
cargo test
129+
cargo insta review # Review and accept/reject changes
130+
```
131+
132+
## Key Dependencies
133+
134+
- **DataFusion**: Query execution engine (forked version with custom modifications)
135+
- **sqlparser-rs**: SQL parser (forked with CubeSQL-specific extensions)
136+
- **egg**: E-graph library for query optimization
137+
- **tokio**: Async runtime for network and I/O operations
138+
- **pgwire**: PostgreSQL wire protocol implementation
139+
140+
## Important Notes
141+
142+
- This codebase uses heavily modified forks of DataFusion and sqlparser-rs
143+
- Many clippy lints are disabled due to code generation and complex patterns
144+
- Integration tests require a running Cube.js instance
145+
- The rewrite engine is performance-critical and uses advanced optimization techniques
146+
- Protocol compatibility is paramount for BI tool support

rust/cubesql/cubesql/e2e/tests/postgres.rs

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -407,19 +407,17 @@ impl PostgresIntegrationTestSuite {
407407
Ok(())
408408
}
409409

410-
async fn test_prepare(&self) -> RunResult<()> {
410+
async fn test_prepare_autodetect(&self) -> RunResult<()> {
411411
// Unknown variables will be detected as TEXT
412412
// LIMIT has a typehint for i64
413413
let stmt = self
414414
.client
415415
.prepare("SELECT $1 as t1, $2 as t2 LIMIT $3")
416-
.await
417-
.unwrap();
416+
.await?;
418417

419418
self.client
420419
.query(&stmt, &[&"test1", &"test2", &0_i64])
421-
.await
422-
.unwrap();
420+
.await?;
423421

424422
Ok(())
425423
}
@@ -444,9 +442,9 @@ impl PostgresIntegrationTestSuite {
444442
}
445443

446444
async fn test_prepare_empty_query(&self) -> RunResult<()> {
447-
let stmt = self.client.prepare("").await.unwrap();
445+
let stmt = self.client.prepare("").await?;
448446

449-
self.client.query(&stmt, &[]).await.unwrap();
447+
self.client.query(&stmt, &[]).await?;
450448

451449
Ok(())
452450
}
@@ -1170,7 +1168,7 @@ impl AsyncTestSuite for PostgresIntegrationTestSuite {
11701168
async fn run(&mut self) -> RunResult<()> {
11711169
self.test_cancel_simple_query().await?;
11721170
self.test_cancel_execute_prepared().await?;
1173-
self.test_prepare().await?;
1171+
self.test_prepare_autodetect().await?;
11741172
self.test_extended_error().await?;
11751173
self.test_prepare_empty_query().await?;
11761174
self.test_stream_all().await?;

rust/cubesql/pg-srv/src/decoding.rs

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,26 @@ impl FromProtocolValue for bool {
102102
}
103103
}
104104

105+
impl FromProtocolValue for f64 {
106+
fn from_text(raw: &[u8]) -> Result<Self, ProtocolError> {
107+
let as_str = std::str::from_utf8(raw).map_err(|err| ProtocolError::ErrorResponse {
108+
source: ErrorResponse::error(ErrorCode::ProtocolViolation, err.to_string()),
109+
backtrace: Backtrace::capture(),
110+
})?;
111+
112+
as_str
113+
.parse::<f64>()
114+
.map_err(|err| ProtocolError::ErrorResponse {
115+
source: ErrorResponse::error(ErrorCode::ProtocolViolation, err.to_string()),
116+
backtrace: Backtrace::capture(),
117+
})
118+
}
119+
120+
fn from_binary(raw: &[u8]) -> Result<Self, ProtocolError> {
121+
Ok(BigEndian::read_f64(raw))
122+
}
123+
}
124+
105125
#[cfg(test)]
106126
mod tests {
107127
use crate::*;
@@ -132,6 +152,9 @@ mod tests {
132152
assert_test_decode(false, Format::Text)?;
133153
assert_test_decode(1_i64, Format::Text)?;
134154
assert_test_decode(100_i64, Format::Text)?;
155+
assert_test_decode(3.14_f64, Format::Text)?;
156+
assert_test_decode(-2.71828_f64, Format::Text)?;
157+
assert_test_decode(0.0_f64, Format::Text)?;
135158

136159
Ok(())
137160
}
@@ -143,6 +166,9 @@ mod tests {
143166
assert_test_decode(false, Format::Binary)?;
144167
assert_test_decode(1_i64, Format::Binary)?;
145168
assert_test_decode(100_i64, Format::Binary)?;
169+
assert_test_decode(3.14_f64, Format::Binary)?;
170+
assert_test_decode(-2.71828_f64, Format::Binary)?;
171+
assert_test_decode(0.0_f64, Format::Binary)?;
146172

147173
Ok(())
148174
}

rust/cubesql/pg-srv/src/protocol.rs

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -783,6 +783,9 @@ impl Bind {
783783
PgTypeId::INT8 => {
784784
BindValue::Int64(i64::from_protocol(raw_value, param_format)?)
785785
}
786+
PgTypeId::FLOAT8 => {
787+
BindValue::Float64(f64::from_protocol(raw_value, param_format)?)
788+
}
786789
_ => {
787790
return Err(ErrorResponse::error(
788791
ErrorCode::FeatureNotSupported,
@@ -1279,6 +1282,73 @@ mod tests {
12791282
Ok(())
12801283
}
12811284

1285+
#[tokio::test]
1286+
async fn test_frontend_message_parse_bind_float64() -> Result<(), ProtocolError> {
1287+
// Test text format float64
1288+
let buffer = parse_hex_dump(
1289+
r#"
1290+
42 00 00 00 19 00 73 30 00 00 01 00 00 00 01 00 B.....s0........
1291+
00 00 04 33 2e 31 34 00 00 00 00 ...3.14....
1292+
"#
1293+
.to_string(),
1294+
);
1295+
let mut cursor = Cursor::new(buffer);
1296+
1297+
let message = read_message(&mut cursor, MessageTagParserDefaultImpl::with_arc()).await?;
1298+
match message {
1299+
FrontendMessage::Bind(body) => {
1300+
assert_eq!(
1301+
body,
1302+
Bind {
1303+
portal: "".to_string(),
1304+
statement: "s0".to_string(),
1305+
parameter_formats: vec![Format::Text],
1306+
parameter_values: vec![Some(vec![51, 46, 49, 52])], // "3.14"
1307+
result_formats: vec![]
1308+
},
1309+
);
1310+
1311+
assert_eq!(
1312+
body.to_bind_values(&ParameterDescription::new(vec![PgTypeId::FLOAT8]))
1313+
.unwrap(),
1314+
vec![BindValue::Float64(3.14)]
1315+
);
1316+
}
1317+
_ => panic!("Wrong message, must be Bind"),
1318+
}
1319+
1320+
// Test binary format float64
1321+
let buffer = parse_hex_dump(
1322+
r#"
1323+
42 00 00 00 1e 00 73 30 00 00 01 00 01 00 01 00 B.....s0........
1324+
00 00 08 40 09 1e b8 51 eb 85 1f 00 00 00 00 ...@...Q.......
1325+
"#
1326+
.to_string(),
1327+
);
1328+
let mut cursor = Cursor::new(buffer);
1329+
1330+
let message = read_message(&mut cursor, MessageTagParserDefaultImpl::with_arc()).await?;
1331+
match message {
1332+
FrontendMessage::Bind(body) => {
1333+
assert_eq!(body.parameter_formats, vec![Format::Binary]);
1334+
1335+
let values = body
1336+
.to_bind_values(&ParameterDescription::new(vec![PgTypeId::FLOAT8]))
1337+
.unwrap();
1338+
1339+
match &values[0] {
1340+
BindValue::Float64(val) => {
1341+
assert!((val - 3.14).abs() < 0.0001);
1342+
}
1343+
_ => panic!("Expected Float64 bind value"),
1344+
}
1345+
}
1346+
_ => panic!("Wrong message, must be Bind"),
1347+
}
1348+
1349+
Ok(())
1350+
}
1351+
12821352
#[tokio::test]
12831353
async fn test_frontend_message_parse_describe() -> Result<(), ProtocolError> {
12841354
let buffer = parse_hex_dump(

0 commit comments

Comments
 (0)