Skip to main content

pgdrift_db/
fixtures.rs

1use serde_json::json;
2use sqlx::PgPool;
3
4/// Create a test users table with consistent JSONB schema
5pub async fn create_users_consistent(pool: &PgPool) -> Result<(), sqlx::Error> {
6    sqlx::query(
7        r#"
8        CREATE TABLE IF NOT EXISTS users (
9            id SERIAL PRIMARY KEY,
10            metadata JSONB NOT NULL
11        )
12        "#,
13    )
14    .execute(pool)
15    .await?;
16
17    // Insert 5000 rows with consistent schema
18    for i in 0..5000 {
19        let metadata = json!({
20            "email": format!("user{}@example.com", i),
21            "age": 25 + (i % 50),
22            "country": "USA",
23            "preferences": {
24                "theme": "dark",
25                "notifications": true,
26                "language": "en"
27            },
28            "tags": ["active", "verified"],
29            "created_at": "2025-01-01",
30            "status": "active"
31        });
32
33        sqlx::query("INSERT INTO users (metadata) VALUES ($1)")
34            .bind(metadata)
35            .execute(pool)
36            .await?;
37    }
38
39    Ok(())
40}
41
42/// Create users table with type inconsistency (age: string vs number)
43pub async fn create_users_type_inconsistency(pool: &PgPool) -> Result<(), sqlx::Error> {
44    sqlx::query(
45        r#"
46        CREATE TABLE IF NOT EXISTS users_mixed_types (
47            id SERIAL PRIMARY KEY,
48            metadata JSONB NOT NULL
49        )
50        "#,
51    )
52    .execute(pool)
53    .await?;
54
55    // Insert 5000 rows: 92% with age as string, 8% as number
56    for i in 0..5000 {
57        let metadata = if i % 100 < 8 {
58            // 8% - age as number
59            json!({
60                "email": format!("user{}@example.com", i),
61                "age": 25 + (i % 50),
62                "country": "USA",
63                "preferences": {
64                    "theme": "dark",
65                    "notifications": true
66                }
67            })
68        } else {
69            // 92% - age as string
70            json!({
71                "email": format!("user{}@example.com", i),
72                "age": format!("{}", 25 + (i % 50)),
73                "country": "USA",
74                "preferences": {
75                    "theme": "dark",
76                    "notifications": true
77                }
78            })
79        };
80
81        sqlx::query("INSERT INTO users_mixed_types (metadata) VALUES ($1)")
82            .bind(metadata)
83            .execute(pool)
84            .await?;
85    }
86
87    Ok(())
88}
89
90/// Create users table with ghost keys and sparse fields
91pub async fn create_users_ghost_keys(pool: &PgPool) -> Result<(), sqlx::Error> {
92    sqlx::query(
93        r#"
94        CREATE TABLE IF NOT EXISTS users_sparse (
95            id SERIAL PRIMARY KEY,
96            metadata JSONB NOT NULL
97        )
98        "#,
99    )
100    .execute(pool)
101    .await?;
102
103    // Insert 5000 rows with sparse fields
104    for i in 0..5000 {
105        let mut metadata = json!({
106            "email": format!("user{}@example.com", i),
107            "name": format!("User {}", i),
108            "country": "USA"
109        });
110
111        // Ghost key: premium_feature - appears in <1% of rows (20 out of 5000 = 0.4%)
112        if i < 20 {
113            metadata["premium_feature"] = json!({
114                "active": true,
115                "tier": "pro",
116                "expiry": "2025-12-31"
117            });
118        }
119
120        // Missing key: billing_address - present in 95% of rows
121        if i >= 250 {
122            // First 250 rows don't have it
123            metadata["billing_address"] = json!({
124                "street": "123 Main St",
125                "city": "San Francisco",
126                "state": "CA"
127            });
128        }
129
130        // Occasional nested array
131        if i % 100 == 0 {
132            metadata["addresses"] = json!([
133                {"type": "home", "city": "SF"},
134                {"type": "work", "city": "Oakland"}
135            ]);
136        }
137
138        sqlx::query("INSERT INTO users_sparse (metadata) VALUES ($1)")
139            .bind(metadata)
140            .execute(pool)
141            .await?;
142    }
143
144    Ok(())
145}
146
147/// Create users table with deeply nested schema
148pub async fn create_users_nested(pool: &PgPool) -> Result<(), sqlx::Error> {
149    sqlx::query(
150        r#"
151        CREATE TABLE IF NOT EXISTS users_nested (
152            id SERIAL PRIMARY KEY,
153            metadata JSONB NOT NULL
154        )
155        "#,
156    )
157    .execute(pool)
158    .await?;
159
160    // Insert 5000 rows with deep nesting
161    for i in 0..5000 {
162        let metadata = json!({
163            "user": {
164                "profile": {
165                    "personal": {
166                        "name": {
167                            "first": "John",
168                            "last": "Doe"
169                        },
170                        "contact": {
171                            "email": format!("user{}@example.com", i),
172                            "phone": "+1-555-0100"
173                        }
174                    },
175                    "settings": {
176                        "privacy": {
177                            "profile_public": true,
178                            "email_visible": false
179                        }
180                    }
181                },
182                "subscriptions": [
183                    {"type": "email", "enabled": true},
184                    {"type": "sms", "enabled": false}
185                ]
186            }
187        });
188
189        sqlx::query("INSERT INTO users_nested (metadata) VALUES ($1)")
190            .bind(metadata)
191            .execute(pool)
192            .await?;
193    }
194
195    Ok(())
196}
197
198/// Create a simple products table for additional test scenarios
199pub async fn create_products_schema_evolution(pool: &PgPool) -> Result<(), sqlx::Error> {
200    sqlx::query(
201        r#"
202        CREATE TABLE IF NOT EXISTS products (
203            id SERIAL PRIMARY KEY,
204            data JSONB NOT NULL
205        )
206        "#,
207    )
208    .execute(pool)
209    .await?;
210
211    // Insert 5000 rows with schema evolution (old vs new format)
212    for i in 0..5000 {
213        let data = if i < 2500 {
214            // Old schema format
215            json!({
216                "product_id": i,
217                "name": format!("Product {}", i),
218                "price": 99.99,
219                "inventory": 100
220            })
221        } else {
222            // New schema format (with additional fields)
223            json!({
224                "product_id": i,
225                "name": format!("Product {}", i),
226                "price": 99.99,
227                "inventory": 100,
228                "sku": format!("SKU-{}", i),
229                "category": "electronics",
230                "tags": ["new", "featured"],
231                "metadata": {
232                    "version": "2.0",
233                    "updated_at": "2025-01-01"
234                }
235            })
236        };
237
238        sqlx::query("INSERT INTO products (data) VALUES ($1)")
239            .bind(data)
240            .execute(pool)
241            .await?;
242    }
243
244    Ok(())
245}
246
247/// Clean up all test tables
248pub async fn cleanup(pool: &PgPool) -> Result<(), sqlx::Error> {
249    sqlx::query("DROP TABLE IF EXISTS users")
250        .execute(pool)
251        .await?;
252    sqlx::query("DROP TABLE IF EXISTS users_mixed_types")
253        .execute(pool)
254        .await?;
255    sqlx::query("DROP TABLE IF EXISTS users_sparse")
256        .execute(pool)
257        .await?;
258    sqlx::query("DROP TABLE IF EXISTS users_nested")
259        .execute(pool)
260        .await?;
261    sqlx::query("DROP TABLE IF EXISTS products")
262        .execute(pool)
263        .await?;
264    Ok(())
265}