refactor(plan): converge estimated plan capability and strategy
This commit is contained in:
jearhe authored and GitHub committed
1 parent
c5eb0a7a4a
commit
78e717bea6
2 files changed
+270
-135
No files matched your search
@@ -42,7 +42,7 @@ pub struct ExplainSqlBuildResult {
|
||||
/// Callers that hand a raw plan to a consumer which cannot inspect the SQL
|
||||
/// (the plugin Host API) need to know how to decode it. Keep this in sync with
|
||||
/// the generated statements in [`build_explain_sql`];
|
||||
/// `estimated_plan_format_matches_generated_sql` guards the pairing.
|
||||
/// `estimated_plan_strategy_matches_capability_sql_format_and_acquisition` guards the pairing.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum EstimatedPlanFormat {
|
||||
@@ -64,6 +64,51 @@ impl EstimatedPlanFormat {
|
||||
}
|
||||
}
|
||||
|
||||
/// The execution contract used to acquire an estimated plan.
|
||||
///
|
||||
/// This describes the acquisition path only; connection and driver execution
|
||||
/// remain owned by `dbx-core`.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum EstimatedPlanAcquisition {
|
||||
/// Execute the SQL generated by `build_explain_sql`.
|
||||
GeneratedSql,
|
||||
/// Ask the connected driver's existing explain facility for a plan.
|
||||
DriverNative,
|
||||
/// Capture SHOWPLAN on a dedicated SQL Server client session.
|
||||
SqlServerShowPlanSession,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum EstimatedPlanSqlGeneration {
|
||||
Postgres,
|
||||
Mysql,
|
||||
Plain,
|
||||
Oracle,
|
||||
SqlServer,
|
||||
Json,
|
||||
}
|
||||
|
||||
/// One complete Estimated Plan contract for a database dialect.
|
||||
///
|
||||
/// Capability, SQL generation, payload encoding and Plugin Host acquisition
|
||||
/// all derive from this mapping.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct EstimatedPlanStrategy {
|
||||
sql_generation: EstimatedPlanSqlGeneration,
|
||||
acquisition: EstimatedPlanAcquisition,
|
||||
format: EstimatedPlanFormat,
|
||||
}
|
||||
|
||||
impl EstimatedPlanStrategy {
|
||||
pub const fn acquisition(self) -> EstimatedPlanAcquisition {
|
||||
self.acquisition
|
||||
}
|
||||
|
||||
pub const fn format(self) -> EstimatedPlanFormat {
|
||||
self.format
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DroppedFilePreviewSqlOptions {
|
||||
@@ -73,9 +118,9 @@ pub struct DroppedFilePreviewSqlOptions {
|
||||
}
|
||||
|
||||
pub fn build_explain_sql(options: ExplainSqlOptions) -> ExplainSqlBuildResult {
|
||||
if !supports_explain_plan(options.database_type) {
|
||||
let Some(strategy) = estimated_plan_strategy(options.database_type) else {
|
||||
return explain_err("unsupported");
|
||||
}
|
||||
};
|
||||
|
||||
let source = strip_trailing_semicolons(options.sql.trim());
|
||||
if source.is_empty() {
|
||||
@@ -96,33 +141,30 @@ pub fn build_explain_sql(options: ExplainSqlOptions) -> ExplainSqlBuildResult {
|
||||
return explain_err("unsafe");
|
||||
}
|
||||
|
||||
let sql = match options.database_type {
|
||||
// ANALYZE executes the statement; is_safe_explain_sql has already limited
|
||||
// the source to SELECT/WITH/TABLE/VALUES. MongoDb shares the plain arm only.
|
||||
Some(DatabaseType::Postgres) if options.analyze == Some(true) => {
|
||||
let sql = match strategy.sql_generation {
|
||||
// ANALYZE executes the statement; the safety gate above has already
|
||||
// limited the source to the existing read-only subset.
|
||||
EstimatedPlanSqlGeneration::Postgres if options.analyze == Some(true) => {
|
||||
format!("EXPLAIN (ANALYZE, FORMAT JSON) {source}")
|
||||
}
|
||||
Some(DatabaseType::Postgres | DatabaseType::MongoDb) => {
|
||||
format!("EXPLAIN (FORMAT JSON) {source}")
|
||||
}
|
||||
Some(DatabaseType::Dameng | DatabaseType::Questdb) => {
|
||||
format!("EXPLAIN {source}")
|
||||
}
|
||||
Some(DatabaseType::Doris) => format!("EXPLAIN {source}"),
|
||||
Some(DatabaseType::Oracle) => format!("EXPLAIN PLAN FOR {source}"),
|
||||
EstimatedPlanSqlGeneration::Postgres => format!("EXPLAIN (FORMAT JSON) {source}"),
|
||||
EstimatedPlanSqlGeneration::Plain => format!("EXPLAIN {source}"),
|
||||
EstimatedPlanSqlGeneration::Oracle => format!("EXPLAIN PLAN FOR {source}"),
|
||||
// STATISTICS XML returns the same ShowPlanXML document plus per-operator
|
||||
// runtime counters, at the price of actually running the statement.
|
||||
Some(DatabaseType::SqlServer) if options.analyze == Some(true) => {
|
||||
EstimatedPlanSqlGeneration::SqlServer if options.analyze == Some(true) => {
|
||||
format!("SET STATISTICS XML ON;\nGO\n{source}\nGO\nSET STATISTICS XML OFF;")
|
||||
}
|
||||
Some(DatabaseType::SqlServer) => {
|
||||
EstimatedPlanSqlGeneration::SqlServer => {
|
||||
format!("SET SHOWPLAN_XML ON;\nGO\n{source}\nGO\nSET SHOWPLAN_XML OFF;")
|
||||
}
|
||||
Some(DatabaseType::Mysql) if options.format == Some(ExplainFormat::Standard) => {
|
||||
EstimatedPlanSqlGeneration::Mysql if options.format == Some(ExplainFormat::Standard) => {
|
||||
// MySQL 8.0.32+ may otherwise inherit TREE or JSON from the session-level explain_format.
|
||||
format!("EXPLAIN FORMAT=TRADITIONAL {source}")
|
||||
}
|
||||
_ => format!("EXPLAIN FORMAT=JSON {source}"),
|
||||
EstimatedPlanSqlGeneration::Mysql | EstimatedPlanSqlGeneration::Json => {
|
||||
format!("EXPLAIN FORMAT=JSON {source}")
|
||||
}
|
||||
};
|
||||
ExplainSqlBuildResult { ok: true, sql: Some(sql), reason: None }
|
||||
}
|
||||
@@ -146,34 +188,51 @@ pub fn build_dropped_file_preview_sql(options: DroppedFilePreviewSqlOptions) ->
|
||||
None
|
||||
}
|
||||
|
||||
/// How a consumer should decode the estimated plan `build_explain_sql` asks
|
||||
/// the server for. Dameng and Oracle plans reach DBX through the driver's
|
||||
/// native plan text rather than a generated `EXPLAIN` statement, and every
|
||||
/// agent plan is text.
|
||||
pub fn estimated_plan_format(database_type: Option<DatabaseType>) -> EstimatedPlanFormat {
|
||||
match database_type {
|
||||
Some(DatabaseType::SqlServer) => EstimatedPlanFormat::Xml,
|
||||
Some(DatabaseType::Dameng | DatabaseType::Questdb | DatabaseType::Doris | DatabaseType::Oracle) => {
|
||||
EstimatedPlanFormat::Text
|
||||
/// The estimated-plan contract for a database, if one is currently supported.
|
||||
pub fn estimated_plan_strategy(database_type: Option<DatabaseType>) -> Option<EstimatedPlanStrategy> {
|
||||
use EstimatedPlanAcquisition::{DriverNative, GeneratedSql, SqlServerShowPlanSession};
|
||||
use EstimatedPlanFormat::{Json, Text, Xml};
|
||||
use EstimatedPlanSqlGeneration::{Mysql, Oracle, Plain, Postgres, SqlServer};
|
||||
|
||||
let strategy = match database_type? {
|
||||
DatabaseType::Mysql => EstimatedPlanStrategy { sql_generation: Mysql, acquisition: GeneratedSql, format: Json },
|
||||
DatabaseType::Doris => EstimatedPlanStrategy { sql_generation: Plain, acquisition: GeneratedSql, format: Text },
|
||||
DatabaseType::Postgres => {
|
||||
EstimatedPlanStrategy { sql_generation: Postgres, acquisition: GeneratedSql, format: Json }
|
||||
}
|
||||
_ => EstimatedPlanFormat::Json,
|
||||
}
|
||||
DatabaseType::Questdb => {
|
||||
EstimatedPlanStrategy { sql_generation: Plain, acquisition: GeneratedSql, format: Text }
|
||||
}
|
||||
DatabaseType::Dameng => {
|
||||
EstimatedPlanStrategy { sql_generation: Plain, acquisition: DriverNative, format: Text }
|
||||
}
|
||||
DatabaseType::Oracle => {
|
||||
EstimatedPlanStrategy { sql_generation: Oracle, acquisition: DriverNative, format: Text }
|
||||
}
|
||||
DatabaseType::OceanbaseOracle => EstimatedPlanStrategy {
|
||||
sql_generation: EstimatedPlanSqlGeneration::Json,
|
||||
acquisition: GeneratedSql,
|
||||
format: Json,
|
||||
},
|
||||
DatabaseType::SqlServer => {
|
||||
EstimatedPlanStrategy { sql_generation: SqlServer, acquisition: SqlServerShowPlanSession, format: Xml }
|
||||
}
|
||||
// The former MongoDB SQL-generation arm was unreachable behind the old
|
||||
// capability gate; it remains unsupported and has no Estimated Plan strategy.
|
||||
_ => return None,
|
||||
};
|
||||
Some(strategy)
|
||||
}
|
||||
|
||||
/// Payload encoding defined by the estimated-plan strategy for this dialect.
|
||||
/// Keep the historical JSON fallback for unsupported/unspecified dialects;
|
||||
/// callers must check capability separately before using it.
|
||||
pub fn estimated_plan_format(database_type: Option<DatabaseType>) -> EstimatedPlanFormat {
|
||||
estimated_plan_strategy(database_type).map_or(EstimatedPlanFormat::Json, EstimatedPlanStrategy::format)
|
||||
}
|
||||
|
||||
pub fn supports_explain_plan(database_type: Option<DatabaseType>) -> bool {
|
||||
matches!(
|
||||
database_type,
|
||||
Some(
|
||||
DatabaseType::Mysql
|
||||
| DatabaseType::Doris
|
||||
| DatabaseType::Postgres
|
||||
| DatabaseType::Questdb
|
||||
| DatabaseType::Dameng
|
||||
| DatabaseType::Oracle
|
||||
| DatabaseType::OceanbaseOracle
|
||||
| DatabaseType::SqlServer
|
||||
)
|
||||
)
|
||||
estimated_plan_strategy(database_type).is_some()
|
||||
}
|
||||
|
||||
pub fn is_safe_explain_sql(sql: &str) -> bool {
|
||||
@@ -1398,9 +1457,11 @@ mod tests {
|
||||
for (db_type, expected) in [
|
||||
(DatabaseType::MongoDb, "EXPLAIN (FORMAT JSON) SELECT 1"),
|
||||
(DatabaseType::Mysql, "EXPLAIN FORMAT=JSON SELECT 1"),
|
||||
(DatabaseType::Dameng, "EXPLAIN SELECT 1"),
|
||||
(DatabaseType::Doris, "EXPLAIN SELECT 1"),
|
||||
(DatabaseType::Questdb, "EXPLAIN SELECT 1"),
|
||||
(DatabaseType::Dameng, "EXPLAIN SELECT 1"),
|
||||
(DatabaseType::Oracle, "EXPLAIN PLAN FOR SELECT 1"),
|
||||
(DatabaseType::OceanbaseOracle, "EXPLAIN FORMAT=JSON SELECT 1"),
|
||||
] {
|
||||
let analyzed = build_explain_sql(ExplainSqlOptions {
|
||||
database_type: Some(db_type),
|
||||
@@ -1416,8 +1477,8 @@ mod tests {
|
||||
});
|
||||
|
||||
assert_eq!(analyzed, plain, "{db_type:?} must ignore the analyze flag");
|
||||
// MongoDb is rejected earlier by supports_explain_plan, so it never
|
||||
// reaches the match arm it nominally shares with Postgres.
|
||||
// MongoDb is rejected by the strategy capability gate; its historical
|
||||
// SQL-generation arm remains unreachable.
|
||||
if db_type != DatabaseType::MongoDb {
|
||||
assert_eq!(analyzed.sql, Some(expected.to_string()), "{db_type:?}");
|
||||
} else {
|
||||
@@ -1570,44 +1631,71 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimated_plan_format_matches_generated_sql() {
|
||||
// Every dialect `supports_explain_plan` advertises must declare a payload
|
||||
// format that the statement `build_explain_sql` generates really produces.
|
||||
for database_type in
|
||||
DatabaseType::ALL.iter().copied().filter(|database_type| supports_explain_plan(Some(*database_type)))
|
||||
{
|
||||
let explain_sql = build_explain_sql(ExplainSqlOptions {
|
||||
fn estimated_plan_strategy_matches_capability_sql_format_and_acquisition() {
|
||||
use EstimatedPlanAcquisition::{DriverNative, GeneratedSql, SqlServerShowPlanSession};
|
||||
use EstimatedPlanFormat::{Json, Text, Xml};
|
||||
|
||||
let cases = [
|
||||
(DatabaseType::Mysql, "EXPLAIN FORMAT=JSON SELECT 1", Json, GeneratedSql),
|
||||
(DatabaseType::Doris, "EXPLAIN SELECT 1", Text, GeneratedSql),
|
||||
(DatabaseType::Postgres, "EXPLAIN (FORMAT JSON) SELECT 1", Json, GeneratedSql),
|
||||
(DatabaseType::Questdb, "EXPLAIN SELECT 1", Text, GeneratedSql),
|
||||
(DatabaseType::Dameng, "EXPLAIN SELECT 1", Text, DriverNative),
|
||||
(DatabaseType::Oracle, "EXPLAIN PLAN FOR SELECT 1", Text, DriverNative),
|
||||
(DatabaseType::OceanbaseOracle, "EXPLAIN FORMAT=JSON SELECT 1", Json, GeneratedSql),
|
||||
(
|
||||
DatabaseType::SqlServer,
|
||||
"SET SHOWPLAN_XML ON;\nGO\nSELECT 1\nGO\nSET SHOWPLAN_XML OFF;",
|
||||
Xml,
|
||||
SqlServerShowPlanSession,
|
||||
),
|
||||
];
|
||||
|
||||
for (database_type, expected_sql, expected_format, expected_acquisition) in cases {
|
||||
let strategy = estimated_plan_strategy(Some(database_type))
|
||||
.unwrap_or_else(|| panic!("{database_type:?} must have a complete Estimated Plan strategy"));
|
||||
assert!(supports_explain_plan(Some(database_type)), "{database_type:?}");
|
||||
assert_eq!(strategy.acquisition(), expected_acquisition, "{database_type:?}");
|
||||
assert_eq!(strategy.format(), expected_format, "{database_type:?}");
|
||||
assert_eq!(estimated_plan_format(Some(database_type)), expected_format, "{database_type:?}");
|
||||
|
||||
let result = build_explain_sql(ExplainSqlOptions {
|
||||
database_type: Some(database_type),
|
||||
format: None,
|
||||
analyze: None,
|
||||
sql: "SELECT 1".to_string(),
|
||||
})
|
||||
.sql
|
||||
.expect("a supported dialect always builds EXPLAIN SQL");
|
||||
|
||||
match estimated_plan_format(Some(database_type)) {
|
||||
EstimatedPlanFormat::Json => assert!(
|
||||
explain_sql.contains("FORMAT JSON") || explain_sql.contains("FORMAT=JSON"),
|
||||
"{database_type:?} declares a JSON plan but generated '{explain_sql}'"
|
||||
),
|
||||
EstimatedPlanFormat::Xml => assert!(
|
||||
explain_sql.contains("SHOWPLAN_XML"),
|
||||
"{database_type:?} declares an XML plan but generated '{explain_sql}'"
|
||||
),
|
||||
EstimatedPlanFormat::Text => assert!(
|
||||
explain_sql.starts_with("EXPLAIN ") && !explain_sql.contains("FORMAT"),
|
||||
"{database_type:?} declares a text plan but generated '{explain_sql}'"
|
||||
),
|
||||
}
|
||||
});
|
||||
assert!(result.ok, "{database_type:?}: {:?}", result.reason);
|
||||
assert_eq!(result.sql.as_deref(), Some(expected_sql), "{database_type:?}");
|
||||
}
|
||||
|
||||
assert_eq!(estimated_plan_format(None), EstimatedPlanFormat::Json);
|
||||
assert_eq!(estimated_plan_format(Some(DatabaseType::Dameng)), EstimatedPlanFormat::Text);
|
||||
assert_eq!(estimated_plan_format(Some(DatabaseType::Questdb)), EstimatedPlanFormat::Text);
|
||||
assert_eq!(estimated_plan_format(Some(DatabaseType::Doris)), EstimatedPlanFormat::Text);
|
||||
assert_eq!(estimated_plan_format(Some(DatabaseType::Oracle)), EstimatedPlanFormat::Text);
|
||||
assert_eq!(estimated_plan_format(Some(DatabaseType::SqlServer)), EstimatedPlanFormat::Xml);
|
||||
assert_eq!(estimated_plan_format(Some(DatabaseType::Mysql)), EstimatedPlanFormat::Json);
|
||||
for database_type in [
|
||||
None,
|
||||
Some(DatabaseType::Redis),
|
||||
Some(DatabaseType::MongoDb),
|
||||
Some(DatabaseType::Gaussdb),
|
||||
Some(DatabaseType::StarRocks),
|
||||
] {
|
||||
assert!(estimated_plan_strategy(database_type).is_none(), "{database_type:?}");
|
||||
assert!(!supports_explain_plan(database_type), "{database_type:?}");
|
||||
assert_eq!(
|
||||
build_explain_sql(ExplainSqlOptions {
|
||||
database_type,
|
||||
format: None,
|
||||
analyze: None,
|
||||
sql: "SELECT 1".to_string(),
|
||||
})
|
||||
.reason
|
||||
.as_deref(),
|
||||
Some("unsupported"),
|
||||
"{database_type:?}"
|
||||
);
|
||||
}
|
||||
|
||||
// Preserve the legacy fallback for callers that ask for a format without
|
||||
// first checking support; it does not enable plan capability.
|
||||
assert_eq!(estimated_plan_format(None), Json);
|
||||
assert_eq!(estimated_plan_format(Some(DatabaseType::Redis)), Json);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1697,6 +1785,26 @@ mod tests {
|
||||
ExplainSqlBuildResult { ok: false, sql: None, reason: Some("unsafe".to_string()) }
|
||||
);
|
||||
|
||||
for sql in [
|
||||
"DROP TABLE users;",
|
||||
"DELETE FROM users;",
|
||||
"INSERT INTO users (id) VALUES (1);",
|
||||
"UPDATE users SET name = 'x';",
|
||||
"SELECT 1; DROP TABLE users;",
|
||||
"CREATE INDEX idx ON users (id);",
|
||||
] {
|
||||
assert_eq!(
|
||||
build_explain_sql(ExplainSqlOptions {
|
||||
database_type: Some(DatabaseType::Mysql),
|
||||
format: None,
|
||||
analyze: None,
|
||||
sql: sql.to_string(),
|
||||
}),
|
||||
ExplainSqlBuildResult { ok: false, sql: None, reason: Some("unsafe".to_string()) },
|
||||
"{sql} must remain blocked by the explain safety gate"
|
||||
);
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
build_explain_sql(ExplainSqlOptions {
|
||||
database_type: Some(DatabaseType::Mysql),
|
||||
|
||||
Reference in new issue
Block a user