Merge remote-tracking branch 'mikael/sim-pragma' into avcm

This commit is contained in:
Pavan-Nambi
2025-10-18 16:41:39 +05:30
7 changed files with 74 additions and 5 deletions

View File

@@ -237,6 +237,7 @@ impl InteractionPlan {
Query::AlterTable(_) => stats.alter_table_count += 1,
Query::DropIndex(_) => stats.drop_index_count += 1,
Query::Placeholder => {}
Query::Pragma(_) => stats.pragma_count += 1,
}
}
for interactions in &self.plan {
@@ -771,6 +772,7 @@ pub(crate) struct InteractionStats {
pub rollback_count: u32,
pub alter_table_count: u32,
pub drop_index_count: u32,
pub pragma_count:u32,
}
impl Display for InteractionStats {

View File

@@ -302,7 +302,6 @@ impl Property {
let rows = insert.rows();
let row = &rows[*row_index];
match &query {
Query::Delete(Delete {
table: t,
@@ -1381,6 +1380,7 @@ pub(super) struct Remaining {
pub drop: u32,
pub alter_table: u32,
pub drop_index: u32,
pub pragma_count: u32,
}
pub(super) fn remaining(
@@ -1401,6 +1401,7 @@ pub(super) fn remaining(
let total_drop = (max_interactions * opts.drop_table_weight) / total_weight;
let total_alter_table = (max_interactions * opts.alter_table_weight) / total_weight;
let total_drop_index = (max_interactions * opts.drop_index) / total_weight;
let total_pragma = (max_interactions * opts.pragma_weight) / total_weight;
let remaining_select = total_select
.checked_sub(stats.select_count)
@@ -1421,6 +1422,9 @@ pub(super) fn remaining(
.checked_sub(stats.update_count)
.unwrap_or_default();
let remaining_drop = total_drop.checked_sub(stats.drop_count).unwrap_or_default();
let remaining_pragma = total_pragma
.checked_sub(stats.pragma_count)
.unwrap_or_default();
let remaining_alter_table = total_alter_table
.checked_sub(stats.alter_table_count)
@@ -1455,6 +1459,7 @@ pub(super) fn remaining(
update: remaining_update,
alter_table: remaining_alter_table,
drop_index: remaining_drop_index,
pragma_count: remaining_pragma,
}
}

View File

@@ -5,12 +5,14 @@ use crate::{
use rand::{
Rng,
distr::{Distribution, weighted::WeightedIndex},
seq::IndexedRandom,
};
use sql_generation::{
generation::{Arbitrary, ArbitraryFrom, GenerationContext, query::SelectFree},
model::{
query::{
Create, CreateIndex, Delete, DropIndex, Insert, Select, alter_table::AlterTable,
pragma::{Pragma, VacuumMode},
update::Update,
},
table::Table,
@@ -82,6 +84,18 @@ fn random_create_index<R: rand::Rng + ?Sized>(
Query::CreateIndex(create_index)
}
fn random_pragma<R: rand::Rng + ?Sized>(rng: &mut R, _conn_ctx: &impl GenerationContext) -> Query {
const ALL_MODES: [VacuumMode; 2] = [
VacuumMode::None,
// VacuumMode::Incremental, not implemented yer
VacuumMode::Full,
];
let mode = ALL_MODES.choose(rng).unwrap();
Query::Pragma(Pragma::AutoVacuumMode(mode.clone()))
}
fn random_alter_table<R: rand::Rng + ?Sized>(
rng: &mut R,
conn_ctx: &impl GenerationContext,
@@ -140,14 +154,13 @@ impl QueryDiscriminants {
QueryDiscriminants::Placeholder => {
unreachable!("Query Placeholders should not be generated")
}
QueryDiscriminants::Pragma => random_pragma,
}
}
fn weight(&self, remaining: &Remaining) -> u32 {
match self {
QueryDiscriminants::Create => remaining.create,
// remaining.select / 3 is for the random_expr generation
// have a max of 1 so that we always generate at least a non zero weight for `QueryDistribution`
QueryDiscriminants::Select => (remaining.select + remaining.select / 3).max(1),
QueryDiscriminants::Insert => remaining.insert,
QueryDiscriminants::Delete => remaining.delete,
@@ -164,6 +177,7 @@ impl QueryDiscriminants {
QueryDiscriminants::Placeholder => {
unreachable!("Query Placeholders should not be generated")
}
QueryDiscriminants::Pragma => remaining.pragma_count,
}
}
}

View File

@@ -9,6 +9,7 @@ use sql_generation::model::{
query::{
Create, CreateIndex, Delete, Drop, DropIndex, Insert, Select,
alter_table::{AlterTable, AlterTableType},
pragma::Pragma,
select::{CompoundOperator, FromClause, ResultColumn, SelectInner},
transaction::{Begin, Commit, Rollback},
update::Update,
@@ -34,6 +35,7 @@ pub enum Query {
Begin(Begin),
Commit(Commit),
Rollback(Rollback),
Pragma(Pragma),
/// Placeholder query that still needs to be generated
Placeholder,
}
@@ -81,8 +83,11 @@ impl Query {
| Query::DropIndex(DropIndex {
table_name: table, ..
}) => IndexSet::from_iter([table.clone()]),
Query::Begin(_) | Query::Commit(_) | Query::Rollback(_) => IndexSet::new(),
Query::Placeholder => IndexSet::new(),
Query::Begin(_)
| Query::Commit(_)
| Query::Rollback(_)
| Query::Placeholder
| Query::Pragma(_) => IndexSet::new(),
}
}
pub fn uses(&self) -> Vec<String> {
@@ -107,6 +112,7 @@ impl Query {
}) => vec![table.clone()],
Query::Begin(..) | Query::Commit(..) | Query::Rollback(..) => vec![],
Query::Placeholder => vec![],
Query::Pragma(_) => vec![],
}
}
@@ -147,6 +153,7 @@ impl Display for Query {
Self::Commit(commit) => write!(f, "{commit}"),
Self::Rollback(rollback) => write!(f, "{rollback}"),
Self::Placeholder => Ok(()),
Query::Pragma(pragma) => write!(f, "{pragma}"),
}
}
}
@@ -169,12 +176,14 @@ impl Shadow for Query {
Query::Commit(commit) => Ok(commit.shadow(env)),
Query::Rollback(rollback) => Ok(rollback.shadow(env)),
Query::Placeholder => Ok(vec![]),
Query::Pragma(Pragma::AutoVacuumMode(_)) => Ok(vec![]),
}
}
}
bitflags! {
pub struct QueryCapabilities: u32 {
const NONE = 0;
const CREATE = 1 << 0;
const SELECT = 1 << 1;
const INSERT = 1 << 2;
@@ -222,6 +231,7 @@ impl From<QueryDiscriminants> for QueryCapabilities {
QueryDiscriminants::Placeholder => {
unreachable!("QueryCapabilities do not apply to query Placeholder")
}
QueryDiscriminants::Pragma => QueryCapabilities::NONE,
}
}
}
@@ -237,6 +247,8 @@ impl QueryDiscriminants {
QueryDiscriminants::CreateIndex,
QueryDiscriminants::AlterTable,
QueryDiscriminants::DropIndex,
QueryDiscriminants::Pragma,
];
}

View File

@@ -26,6 +26,8 @@ pub struct QueryProfile {
pub alter_table_weight: u32,
#[garde(skip)]
pub drop_index: u32,
#[garde(skip)]
pub pragma_weight: u32,
}
impl Default for QueryProfile {
@@ -41,6 +43,7 @@ impl Default for QueryProfile {
drop_table_weight: 2,
alter_table_weight: 2,
drop_index: 2,
pragma_weight: 100,//TODO change this back to 2
}
}
}

View File

@@ -17,3 +17,4 @@ pub mod predicate;
pub mod select;
pub mod transaction;
pub mod update;
pub mod pragma;

View File

@@ -0,0 +1,32 @@
use std::fmt::Display;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum Pragma {
AutoVacuumMode(VacuumMode),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum VacuumMode {
None,
Incremental,
Full,
}
impl Display for Pragma {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Pragma::AutoVacuumMode(vacuum_mode) => {
let mode = match vacuum_mode {
VacuumMode::None => "none",
VacuumMode::Incremental => "incremental",
VacuumMode::Full => "full",
};
write!(f, "PRAGMA auto_vacuum={mode} ")?;
Ok(())
}
}
}
}