+61
-74
Cargo.lock
+61
-74
Cargo.lock
···
838
838
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
839
839
840
840
[[package]]
841
+
name = "foldhash"
842
+
version = "0.1.4"
843
+
source = "registry+https://github.com/rust-lang/crates.io-index"
844
+
checksum = "a0d2fde1f7b3d48b8395d5f2de76c18a528bd6a9cdde438df747bfcba3e05d6f"
845
+
846
+
[[package]]
841
847
name = "foreign-types"
842
848
version = "0.3.2"
843
849
source = "registry+https://github.com/rust-lang/crates.io-index"
···
862
868
]
863
869
864
870
[[package]]
871
+
name = "futures"
872
+
version = "0.3.31"
873
+
source = "registry+https://github.com/rust-lang/crates.io-index"
874
+
checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876"
875
+
dependencies = [
876
+
"futures-channel",
877
+
"futures-core",
878
+
"futures-executor",
879
+
"futures-io",
880
+
"futures-sink",
881
+
"futures-task",
882
+
"futures-util",
883
+
]
884
+
885
+
[[package]]
865
886
name = "futures-channel"
866
887
version = "0.3.31"
867
888
source = "registry+https://github.com/rust-lang/crates.io-index"
···
906
927
checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6"
907
928
908
929
[[package]]
930
+
name = "futures-macro"
931
+
version = "0.3.31"
932
+
source = "registry+https://github.com/rust-lang/crates.io-index"
933
+
checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650"
934
+
dependencies = [
935
+
"proc-macro2",
936
+
"quote",
937
+
"syn",
938
+
]
939
+
940
+
[[package]]
909
941
name = "futures-sink"
910
942
version = "0.3.31"
911
943
source = "registry+https://github.com/rust-lang/crates.io-index"
···
923
955
source = "registry+https://github.com/rust-lang/crates.io-index"
924
956
checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81"
925
957
dependencies = [
958
+
"futures-channel",
926
959
"futures-core",
927
960
"futures-io",
961
+
"futures-macro",
928
962
"futures-sink",
929
963
"futures-task",
930
964
"memchr",
···
983
1017
984
1018
[[package]]
985
1019
name = "hashbrown"
986
-
version = "0.14.5"
1020
+
version = "0.15.2"
987
1021
source = "registry+https://github.com/rust-lang/crates.io-index"
988
-
checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
1022
+
checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289"
989
1023
dependencies = [
990
-
"ahash",
991
1024
"allocator-api2",
1025
+
"equivalent",
1026
+
"foldhash",
992
1027
]
993
1028
994
1029
[[package]]
995
-
name = "hashbrown"
996
-
version = "0.15.2"
997
-
source = "registry+https://github.com/rust-lang/crates.io-index"
998
-
checksum = "bf151400ff0baff5465007dd2f3e717f3fe502074ca563069ce3a6629d07b289"
999
-
1000
-
[[package]]
1001
1030
name = "hashlink"
1002
-
version = "0.9.1"
1031
+
version = "0.10.0"
1003
1032
source = "registry+https://github.com/rust-lang/crates.io-index"
1004
-
checksum = "6ba4ff7128dee98c7dc9794b6a411377e1404dba1c97deb8d1a55297bd25d8af"
1033
+
checksum = "7382cf6263419f2d8df38c55d7da83da5c18aef87fc7a7fc1fb1e344edfe14c1"
1005
1034
dependencies = [
1006
-
"hashbrown 0.14.5",
1035
+
"hashbrown",
1007
1036
]
1008
1037
1009
1038
[[package]]
···
1249
1278
checksum = "8c9c992b02b5b4c94ea26e32fe5bccb7aa7d9f390ab5c1221ff895bc7ea8b652"
1250
1279
dependencies = [
1251
1280
"equivalent",
1252
-
"hashbrown 0.15.2",
1281
+
"hashbrown",
1253
1282
]
1254
1283
1255
1284
[[package]]
···
1414
1443
]
1415
1444
1416
1445
[[package]]
1417
-
name = "minimal-lexical"
1418
-
version = "0.2.1"
1419
-
source = "registry+https://github.com/rust-lang/crates.io-index"
1420
-
checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a"
1421
-
1422
-
[[package]]
1423
1446
name = "miniz_oxide"
1424
1447
version = "0.8.3"
1425
1448
source = "registry+https://github.com/rust-lang/crates.io-index"
···
1458
1481
]
1459
1482
1460
1483
[[package]]
1461
-
name = "nom"
1462
-
version = "7.1.3"
1463
-
source = "registry+https://github.com/rust-lang/crates.io-index"
1464
-
checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a"
1465
-
dependencies = [
1466
-
"memchr",
1467
-
"minimal-lexical",
1468
-
]
1469
-
1470
-
[[package]]
1471
1484
name = "nu-ansi-term"
1472
1485
version = "0.46.0"
1473
1486
source = "registry+https://github.com/rust-lang/crates.io-index"
···
2109
2122
"chrono",
2110
2123
"clap",
2111
2124
"dotenv",
2125
+
"futures",
2112
2126
"jsonwebtoken",
2113
2127
"lazy_static",
2114
2128
"mime_guess",
···
2173
2187
]
2174
2188
2175
2189
[[package]]
2176
-
name = "sqlformat"
2177
-
version = "0.2.6"
2178
-
source = "registry+https://github.com/rust-lang/crates.io-index"
2179
-
checksum = "7bba3a93db0cc4f7bdece8bb09e77e2e785c20bfebf79eb8340ed80708048790"
2180
-
dependencies = [
2181
-
"nom",
2182
-
"unicode_categories",
2183
-
]
2184
-
2185
-
[[package]]
2186
2190
name = "sqlx"
2187
-
version = "0.8.1"
2191
+
version = "0.8.3"
2188
2192
source = "registry+https://github.com/rust-lang/crates.io-index"
2189
-
checksum = "fcfa89bea9500db4a0d038513d7a060566bfc51d46d1c014847049a45cce85e8"
2193
+
checksum = "4410e73b3c0d8442c5f99b425d7a435b5ee0ae4167b3196771dd3f7a01be745f"
2190
2194
dependencies = [
2191
2195
"sqlx-core",
2192
2196
"sqlx-macros",
···
2197
2201
2198
2202
[[package]]
2199
2203
name = "sqlx-core"
2200
-
version = "0.8.1"
2204
+
version = "0.8.3"
2201
2205
source = "registry+https://github.com/rust-lang/crates.io-index"
2202
-
checksum = "d06e2f2bd861719b1f3f0c7dbe1d80c30bf59e76cf019f07d9014ed7eefb8e08"
2206
+
checksum = "6a007b6936676aa9ab40207cde35daab0a04b823be8ae004368c0793b96a61e0"
2203
2207
dependencies = [
2204
-
"atoi",
2205
-
"byteorder",
2206
2208
"bytes",
2207
2209
"chrono",
2208
2210
"crc",
2209
2211
"crossbeam-queue",
2210
2212
"either",
2211
2213
"event-listener",
2212
-
"futures-channel",
2213
2214
"futures-core",
2214
2215
"futures-intrusive",
2215
2216
"futures-io",
2216
2217
"futures-util",
2217
-
"hashbrown 0.14.5",
2218
+
"hashbrown",
2218
2219
"hashlink",
2219
-
"hex",
2220
2220
"indexmap",
2221
2221
"log",
2222
2222
"memchr",
2223
2223
"native-tls",
2224
2224
"once_cell",
2225
-
"paste",
2226
2225
"percent-encoding",
2227
2226
"serde",
2228
2227
"serde_json",
2229
2228
"sha2",
2230
2229
"smallvec",
2231
-
"sqlformat",
2232
-
"thiserror 1.0.69",
2230
+
"thiserror 2.0.11",
2233
2231
"tokio",
2234
2232
"tokio-stream",
2235
2233
"tracing",
2236
2234
"url",
2237
-
"uuid",
2238
2235
]
2239
2236
2240
2237
[[package]]
2241
2238
name = "sqlx-macros"
2242
-
version = "0.8.1"
2239
+
version = "0.8.3"
2243
2240
source = "registry+https://github.com/rust-lang/crates.io-index"
2244
-
checksum = "2f998a9defdbd48ed005a89362bd40dd2117502f15294f61c8d47034107dbbdc"
2241
+
checksum = "3112e2ad78643fef903618d78cf0aec1cb3134b019730edb039b69eaf531f310"
2245
2242
dependencies = [
2246
2243
"proc-macro2",
2247
2244
"quote",
···
2252
2249
2253
2250
[[package]]
2254
2251
name = "sqlx-macros-core"
2255
-
version = "0.8.1"
2252
+
version = "0.8.3"
2256
2253
source = "registry+https://github.com/rust-lang/crates.io-index"
2257
-
checksum = "3d100558134176a2629d46cec0c8891ba0be8910f7896abfdb75ef4ab6f4e7ce"
2254
+
checksum = "4e9f90acc5ab146a99bf5061a7eb4976b573f560bc898ef3bf8435448dd5e7ad"
2258
2255
dependencies = [
2259
2256
"dotenvy",
2260
2257
"either",
···
2278
2275
2279
2276
[[package]]
2280
2277
name = "sqlx-mysql"
2281
-
version = "0.8.1"
2278
+
version = "0.8.3"
2282
2279
source = "registry+https://github.com/rust-lang/crates.io-index"
2283
-
checksum = "936cac0ab331b14cb3921c62156d913e4c15b74fb6ec0f3146bd4ef6e4fb3c12"
2280
+
checksum = "4560278f0e00ce64938540546f59f590d60beee33fffbd3b9cd47851e5fff233"
2284
2281
dependencies = [
2285
2282
"atoi",
2286
2283
"base64 0.22.1",
···
2314
2311
"smallvec",
2315
2312
"sqlx-core",
2316
2313
"stringprep",
2317
-
"thiserror 1.0.69",
2314
+
"thiserror 2.0.11",
2318
2315
"tracing",
2319
-
"uuid",
2320
2316
"whoami",
2321
2317
]
2322
2318
2323
2319
[[package]]
2324
2320
name = "sqlx-postgres"
2325
-
version = "0.8.1"
2321
+
version = "0.8.3"
2326
2322
source = "registry+https://github.com/rust-lang/crates.io-index"
2327
-
checksum = "9734dbce698c67ecf67c442f768a5e90a49b2a4d61a9f1d59f73874bd4cf0710"
2323
+
checksum = "c5b98a57f363ed6764d5b3a12bfedf62f07aa16e1856a7ddc2a0bb190a959613"
2328
2324
dependencies = [
2329
2325
"atoi",
2330
2326
"base64 0.22.1",
···
2336
2332
"etcetera",
2337
2333
"futures-channel",
2338
2334
"futures-core",
2339
-
"futures-io",
2340
2335
"futures-util",
2341
2336
"hex",
2342
2337
"hkdf",
···
2354
2349
"smallvec",
2355
2350
"sqlx-core",
2356
2351
"stringprep",
2357
-
"thiserror 1.0.69",
2352
+
"thiserror 2.0.11",
2358
2353
"tracing",
2359
-
"uuid",
2360
2354
"whoami",
2361
2355
]
2362
2356
2363
2357
[[package]]
2364
2358
name = "sqlx-sqlite"
2365
-
version = "0.8.1"
2359
+
version = "0.8.3"
2366
2360
source = "registry+https://github.com/rust-lang/crates.io-index"
2367
-
checksum = "a75b419c3c1b1697833dd927bdc4c6545a620bc1bbafabd44e1efbe9afcd337e"
2361
+
checksum = "f85ca71d3a5b24e64e1d08dd8fe36c6c95c339a896cc33068148906784620540"
2368
2362
dependencies = [
2369
2363
"atoi",
2370
2364
"chrono",
···
2382
2376
"sqlx-core",
2383
2377
"tracing",
2384
2378
"url",
2385
-
"uuid",
2386
2379
]
2387
2380
2388
2381
[[package]]
···
2705
2698
version = "0.1.3"
2706
2699
source = "registry+https://github.com/rust-lang/crates.io-index"
2707
2700
checksum = "e70f2a8b45122e719eb623c01822704c4e0907e7e426a05927e1a1cfff5b75d0"
2708
-
2709
-
[[package]]
2710
-
name = "unicode_categories"
2711
-
version = "0.1.1"
2712
-
source = "registry+https://github.com/rust-lang/crates.io-index"
2713
-
checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
2714
2701
2715
2702
[[package]]
2716
2703
name = "untrusted"
+2
-1
Cargo.toml
+2
-1
Cargo.toml
···
14
14
actix-files = "0.6"
15
15
actix-cors = "0.6"
16
16
tokio = { version = "1.36", features = ["full"] }
17
-
sqlx = { version = "0.8", features = ["runtime-tokio-native-tls", "postgres", "uuid", "chrono"] }
17
+
sqlx = { version = "0.8", features = ["runtime-tokio-native-tls", "postgres", "sqlite", "chrono"] }
18
18
serde = { version = "1.0", features = ["derive"] }
19
19
serde_json = "1.0"
20
20
anyhow = "1.0"
···
31
31
argon2 = "0.5.3"
32
32
rand = { version = "0.8", features = ["std"] }
33
33
mime_guess = "2.0.5"
34
+
futures = "0.3.31"
+42
migrations/sqlite/20250125000000_init.sql
+42
migrations/sqlite/20250125000000_init.sql
···
1
+
-- Enable foreign key support
2
+
PRAGMA foreign_keys = ON;
3
+
4
+
-- Add Migration Version
5
+
CREATE TABLE IF NOT EXISTS _sqlx_migrations (
6
+
version INTEGER PRIMARY KEY,
7
+
description TEXT NOT NULL,
8
+
installed_on TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
9
+
);
10
+
11
+
-- Create users table
12
+
CREATE TABLE users (
13
+
id INTEGER PRIMARY KEY AUTOINCREMENT,
14
+
email VARCHAR(255) NOT NULL UNIQUE,
15
+
password_hash TEXT NOT NULL
16
+
);
17
+
18
+
-- Create links table
19
+
CREATE TABLE links (
20
+
id INTEGER PRIMARY KEY AUTOINCREMENT,
21
+
original_url TEXT NOT NULL,
22
+
short_code VARCHAR(8) NOT NULL UNIQUE,
23
+
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
24
+
clicks INTEGER NOT NULL DEFAULT 0,
25
+
user_id INTEGER,
26
+
FOREIGN KEY (user_id) REFERENCES users(id)
27
+
);
28
+
29
+
-- Create clicks table
30
+
CREATE TABLE clicks (
31
+
id INTEGER PRIMARY KEY AUTOINCREMENT,
32
+
link_id INTEGER,
33
+
source TEXT,
34
+
query_source TEXT,
35
+
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
36
+
FOREIGN KEY (link_id) REFERENCES links(id)
37
+
);
38
+
39
+
-- Create indexes
40
+
CREATE INDEX idx_short_code ON links(short_code);
41
+
CREATE INDEX idx_user_id ON links(user_id);
42
+
CREATE INDEX idx_link_id ON clicks(link_id);
+433
-152
src/handlers.rs
+433
-152
src/handlers.rs
···
2
2
use crate::{
3
3
error::AppError,
4
4
models::{
5
-
AuthResponse, Claims, ClickStats, CreateLink, Link, LoginRequest, RegisterRequest,
6
-
SourceStats, User, UserResponse,
5
+
AuthResponse, Claims, ClickStats, CreateLink, DatabasePool, Link, LoginRequest,
6
+
RegisterRequest, SourceStats, User, UserResponse,
7
7
},
8
8
AppState,
9
9
};
···
16
16
use jsonwebtoken::{encode, EncodingKey, Header};
17
17
use lazy_static::lazy_static;
18
18
use regex::Regex;
19
+
use sqlx::{Postgres, Sqlite};
19
20
20
21
lazy_static! {
21
22
static ref VALID_CODE_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_-]{1,32}$").unwrap();
···
27
28
payload: web::Json<CreateLink>,
28
29
) -> Result<impl Responder, AppError> {
29
30
tracing::debug!("Creating short URL with user_id: {}", user.user_id);
30
-
31
31
validate_url(&payload.url)?;
32
32
33
33
let short_code = if let Some(ref custom_code) = payload.custom_code {
34
34
validate_custom_code(custom_code)?;
35
35
36
-
tracing::debug!("Checking if custom code {} exists", custom_code);
37
-
// Check if code is already taken
38
-
if let Some(_) = sqlx::query_as::<_, Link>("SELECT * FROM links WHERE short_code = $1")
39
-
.bind(custom_code)
40
-
.fetch_optional(&state.db)
41
-
.await?
42
-
{
36
+
// Check if code exists using match on pool type
37
+
let exists = match &state.db {
38
+
DatabasePool::Postgres(pool) => {
39
+
sqlx::query_as::<_, Link>("SELECT * FROM links WHERE short_code = $1")
40
+
.bind(custom_code)
41
+
.fetch_optional(pool)
42
+
.await?
43
+
}
44
+
DatabasePool::Sqlite(pool) => {
45
+
sqlx::query_as::<_, Link>("SELECT * FROM links WHERE short_code = ?1")
46
+
.bind(custom_code)
47
+
.fetch_optional(pool)
48
+
.await?
49
+
}
50
+
};
51
+
52
+
if exists.is_some() {
43
53
return Err(AppError::InvalidInput(
44
54
"Custom code already taken".to_string(),
45
55
));
46
56
}
47
-
48
57
custom_code.clone()
49
58
} else {
50
59
generate_short_code()
51
60
};
52
61
53
-
// Start transaction
54
-
let mut tx = state.db.begin().await?;
62
+
// Start transaction based on pool type
63
+
let result = match &state.db {
64
+
DatabasePool::Postgres(pool) => {
65
+
let mut tx = pool.begin().await?;
55
66
56
-
tracing::debug!("Inserting new link with short_code: {}", short_code);
57
-
let link = sqlx::query_as::<_, Link>(
58
-
"INSERT INTO links (original_url, short_code, user_id) VALUES ($1, $2, $3) RETURNING *",
59
-
)
60
-
.bind(&payload.url)
61
-
.bind(&short_code)
62
-
.bind(user.user_id)
63
-
.fetch_one(&mut *tx)
64
-
.await?;
67
+
let link = sqlx::query_as::<_, Link>(
68
+
"INSERT INTO links (original_url, short_code, user_id) VALUES ($1, $2, $3) RETURNING *"
69
+
)
70
+
.bind(&payload.url)
71
+
.bind(&short_code)
72
+
.bind(user.user_id)
73
+
.fetch_one(&mut *tx)
74
+
.await?;
75
+
76
+
if let Some(ref source) = payload.source {
77
+
sqlx::query("INSERT INTO clicks (link_id, source) VALUES ($1, $2)")
78
+
.bind(link.id)
79
+
.bind(source)
80
+
.execute(&mut *tx)
81
+
.await?;
82
+
}
83
+
84
+
tx.commit().await?;
85
+
link
86
+
}
87
+
DatabasePool::Sqlite(pool) => {
88
+
let mut tx = pool.begin().await?;
65
89
66
-
if let Some(ref source) = payload.source {
67
-
tracing::debug!("Adding click source: {}", source);
68
-
sqlx::query("INSERT INTO clicks (link_id, source) VALUES ($1, $2)")
69
-
.bind(link.id)
70
-
.bind(source)
71
-
.execute(&mut *tx)
90
+
let link = sqlx::query_as::<_, Link>(
91
+
"INSERT INTO links (original_url, short_code, user_id) VALUES (?1, ?2, ?3) RETURNING *"
92
+
)
93
+
.bind(&payload.url)
94
+
.bind(&short_code)
95
+
.bind(user.user_id)
96
+
.fetch_one(&mut *tx)
72
97
.await?;
73
-
}
74
98
75
-
tx.commit().await?;
76
-
Ok(HttpResponse::Created().json(link))
99
+
if let Some(ref source) = payload.source {
100
+
sqlx::query("INSERT INTO clicks (link_id, source) VALUES (?1, ?2)")
101
+
.bind(link.id)
102
+
.bind(source)
103
+
.execute(&mut *tx)
104
+
.await?;
105
+
}
106
+
107
+
tx.commit().await?;
108
+
link
109
+
}
110
+
};
111
+
112
+
Ok(HttpResponse::Created().json(result))
77
113
}
78
114
79
115
fn validate_custom_code(code: &str) -> Result<(), AppError> {
···
120
156
.and_then(|q| web::Query::<std::collections::HashMap<String, String>>::from_query(q).ok())
121
157
.and_then(|params| params.get("source").cloned());
122
158
123
-
let mut tx = state.db.begin().await?;
124
-
125
-
let link = sqlx::query_as::<_, Link>(
126
-
"UPDATE links SET clicks = clicks + 1 WHERE short_code = $1 RETURNING *",
127
-
)
128
-
.bind(&short_code)
129
-
.fetch_optional(&mut *tx)
130
-
.await?;
159
+
let link = match &state.db {
160
+
DatabasePool::Postgres(pool) => {
161
+
let mut tx = pool.begin().await?;
162
+
let link = sqlx::query_as::<_, Link>(
163
+
"UPDATE links SET clicks = clicks + 1 WHERE short_code = $1 RETURNING *",
164
+
)
165
+
.bind(&short_code)
166
+
.fetch_optional(&mut *tx)
167
+
.await?;
168
+
tx.commit().await?;
169
+
link
170
+
}
171
+
DatabasePool::Sqlite(pool) => {
172
+
let mut tx = pool.begin().await?;
173
+
let link = sqlx::query_as::<_, Link>(
174
+
"UPDATE links SET clicks = clicks + 1 WHERE short_code = ?1 RETURNING *",
175
+
)
176
+
.bind(&short_code)
177
+
.fetch_optional(&mut *tx)
178
+
.await?;
179
+
tx.commit().await?;
180
+
link
181
+
}
182
+
};
131
183
132
184
match link {
133
185
Some(link) => {
134
-
// Record click with both user agent and query source
135
-
let user_agent = req
136
-
.headers()
137
-
.get("user-agent")
138
-
.and_then(|h| h.to_str().ok())
139
-
.unwrap_or("unknown")
140
-
.to_string();
186
+
// Handle click recording based on database type
187
+
match &state.db {
188
+
DatabasePool::Postgres(pool) => {
189
+
let mut tx = pool.begin().await?;
190
+
let user_agent = req
191
+
.headers()
192
+
.get("user-agent")
193
+
.and_then(|h| h.to_str().ok())
194
+
.unwrap_or("unknown")
195
+
.to_string();
196
+
197
+
sqlx::query(
198
+
"INSERT INTO clicks (link_id, source, query_source) VALUES ($1, $2, $3)",
199
+
)
200
+
.bind(link.id)
201
+
.bind(user_agent)
202
+
.bind(query_source)
203
+
.execute(&mut *tx)
204
+
.await?;
205
+
206
+
tx.commit().await?;
207
+
}
208
+
DatabasePool::Sqlite(pool) => {
209
+
let mut tx = pool.begin().await?;
210
+
let user_agent = req
211
+
.headers()
212
+
.get("user-agent")
213
+
.and_then(|h| h.to_str().ok())
214
+
.unwrap_or("unknown")
215
+
.to_string();
141
216
142
-
sqlx::query("INSERT INTO clicks (link_id, source, query_source) VALUES ($1, $2, $3)")
143
-
.bind(link.id)
144
-
.bind(user_agent)
145
-
.bind(query_source)
146
-
.execute(&mut *tx)
147
-
.await?;
217
+
sqlx::query(
218
+
"INSERT INTO clicks (link_id, source, query_source) VALUES (?1, ?2, ?3)",
219
+
)
220
+
.bind(link.id)
221
+
.bind(user_agent)
222
+
.bind(query_source)
223
+
.execute(&mut *tx)
224
+
.await?;
148
225
149
-
tx.commit().await?;
226
+
tx.commit().await?;
227
+
}
228
+
};
150
229
151
230
Ok(HttpResponse::TemporaryRedirect()
152
231
.append_header(("Location", link.original_url))
···
160
239
state: web::Data<AppState>,
161
240
user: AuthenticatedUser,
162
241
) -> Result<impl Responder, AppError> {
163
-
let links = sqlx::query_as::<_, Link>(
164
-
"SELECT * FROM links WHERE user_id = $1 ORDER BY created_at DESC",
165
-
)
166
-
.bind(user.user_id)
167
-
.fetch_all(&state.db)
168
-
.await?;
242
+
let links = match &state.db {
243
+
DatabasePool::Postgres(pool) => {
244
+
sqlx::query_as::<_, Link>(
245
+
"SELECT * FROM links WHERE user_id = $1 ORDER BY created_at DESC",
246
+
)
247
+
.bind(user.user_id)
248
+
.fetch_all(pool)
249
+
.await?
250
+
}
251
+
DatabasePool::Sqlite(pool) => {
252
+
sqlx::query_as::<_, Link>(
253
+
"SELECT * FROM links WHERE user_id = ?1 ORDER BY created_at DESC",
254
+
)
255
+
.bind(user.user_id)
256
+
.fetch_all(pool)
257
+
.await?
258
+
}
259
+
};
169
260
170
261
Ok(HttpResponse::Ok().json(links))
171
262
}
172
263
173
264
pub async fn health_check(state: web::Data<AppState>) -> impl Responder {
174
-
match sqlx::query("SELECT 1").execute(&state.db).await {
175
-
Ok(_) => HttpResponse::Ok().json("Healthy"),
176
-
Err(_) => HttpResponse::ServiceUnavailable().json("Database unavailable"),
265
+
let is_healthy = match &state.db {
266
+
DatabasePool::Postgres(pool) => sqlx::query("SELECT 1").execute(pool).await.is_ok(),
267
+
DatabasePool::Sqlite(pool) => sqlx::query("SELECT 1").execute(pool).await.is_ok(),
268
+
};
269
+
270
+
if is_healthy {
271
+
HttpResponse::Ok().json("Healthy")
272
+
} else {
273
+
HttpResponse::ServiceUnavailable().json("Database unavailable")
177
274
}
178
275
}
179
276
···
190
287
payload: web::Json<RegisterRequest>,
191
288
) -> Result<impl Responder, AppError> {
192
289
// Check if any users exist
193
-
let user_count = sqlx::query!("SELECT COUNT(*) as count FROM users")
194
-
.fetch_one(&state.db)
195
-
.await?
196
-
.count
197
-
.unwrap_or(0);
290
+
let user_count = match &state.db {
291
+
DatabasePool::Postgres(pool) => {
292
+
let mut tx = pool.begin().await?;
293
+
let count = sqlx::query_as::<Postgres, (i64,)>("SELECT COUNT(*)::bigint FROM users")
294
+
.fetch_one(&mut *tx)
295
+
.await?
296
+
.0;
297
+
tx.commit().await?;
298
+
count
299
+
}
300
+
DatabasePool::Sqlite(pool) => {
301
+
let mut tx = pool.begin().await?;
302
+
let count = sqlx::query_as::<Sqlite, (i64,)>("SELECT COUNT(*) FROM users")
303
+
.fetch_one(&mut *tx)
304
+
.await?
305
+
.0;
306
+
tx.commit().await?;
307
+
count
308
+
}
309
+
};
198
310
199
311
// If users exist, registration is closed - no exceptions
200
312
if user_count > 0 {
···
210
322
}
211
323
212
324
// Check if email already exists
213
-
let exists = sqlx::query!("SELECT id FROM users WHERE email = $1", payload.email)
214
-
.fetch_optional(&state.db)
215
-
.await?;
325
+
let exists = match &state.db {
326
+
DatabasePool::Postgres(pool) => {
327
+
let mut tx = pool.begin().await?;
328
+
let exists =
329
+
sqlx::query_as::<Postgres, (i32,)>("SELECT id FROM users WHERE email = $1")
330
+
.bind(&payload.email)
331
+
.fetch_optional(&mut *tx)
332
+
.await?;
333
+
tx.commit().await?;
334
+
exists
335
+
}
336
+
DatabasePool::Sqlite(pool) => {
337
+
let mut tx = pool.begin().await?;
338
+
let exists = sqlx::query_as::<Sqlite, (i32,)>("SELECT id FROM users WHERE email = ?")
339
+
.bind(&payload.email)
340
+
.fetch_optional(&mut *tx)
341
+
.await?;
342
+
tx.commit().await?;
343
+
exists
344
+
}
345
+
};
216
346
217
347
if exists.is_some() {
218
348
return Err(AppError::Auth("Email already registered".to_string()));
···
225
355
.map_err(|e| AppError::Auth(e.to_string()))?
226
356
.to_string();
227
357
228
-
let user = sqlx::query_as!(
229
-
User,
230
-
"INSERT INTO users (email, password_hash) VALUES ($1, $2) RETURNING *",
231
-
payload.email,
232
-
password_hash
233
-
)
234
-
.fetch_one(&state.db)
235
-
.await?;
358
+
// Insert new user
359
+
let user = match &state.db {
360
+
DatabasePool::Postgres(pool) => {
361
+
let mut tx = pool.begin().await?;
362
+
let user = sqlx::query_as::<Postgres, User>(
363
+
"INSERT INTO users (email, password_hash) VALUES ($1, $2) RETURNING *",
364
+
)
365
+
.bind(&payload.email)
366
+
.bind(&password_hash)
367
+
.fetch_one(&mut *tx)
368
+
.await?;
369
+
tx.commit().await?;
370
+
user
371
+
}
372
+
DatabasePool::Sqlite(pool) => {
373
+
let mut tx = pool.begin().await?;
374
+
let user = sqlx::query_as::<Sqlite, User>(
375
+
"INSERT INTO users (email, password_hash) VALUES (?, ?) RETURNING *",
376
+
)
377
+
.bind(&payload.email)
378
+
.bind(&password_hash)
379
+
.fetch_one(&mut *tx)
380
+
.await?;
381
+
tx.commit().await?;
382
+
user
383
+
}
384
+
};
236
385
237
386
let claims = Claims::new(user.id);
238
387
let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
···
256
405
state: web::Data<AppState>,
257
406
payload: web::Json<LoginRequest>,
258
407
) -> Result<impl Responder, AppError> {
259
-
let user = sqlx::query_as!(User, "SELECT * FROM users WHERE email = $1", payload.email)
260
-
.fetch_optional(&state.db)
261
-
.await?
262
-
.ok_or_else(|| AppError::Auth("Invalid credentials".to_string()))?;
408
+
let user = match &state.db {
409
+
DatabasePool::Postgres(pool) => {
410
+
let mut tx = pool.begin().await?;
411
+
let user = sqlx::query_as::<Postgres, User>("SELECT * FROM users WHERE email = $1")
412
+
.bind(&payload.email)
413
+
.fetch_optional(&mut *tx)
414
+
.await?;
415
+
tx.commit().await?;
416
+
user
417
+
}
418
+
DatabasePool::Sqlite(pool) => {
419
+
let mut tx = pool.begin().await?;
420
+
let user = sqlx::query_as::<Sqlite, User>("SELECT * FROM users WHERE email = ?")
421
+
.bind(&payload.email)
422
+
.fetch_optional(&mut *tx)
423
+
.await?;
424
+
tx.commit().await?;
425
+
user
426
+
}
427
+
}
428
+
.ok_or_else(|| AppError::Auth("Invalid credentials".to_string()))?;
263
429
264
430
let argon2 = Argon2::default();
265
431
let parsed_hash =
···
297
463
) -> Result<impl Responder, AppError> {
298
464
let link_id = path.into_inner();
299
465
300
-
// Start transaction
301
-
let mut tx = state.db.begin().await?;
466
+
match &state.db {
467
+
DatabasePool::Postgres(pool) => {
468
+
let mut tx = pool.begin().await?;
302
469
303
-
// Verify the link belongs to the user
304
-
let link = sqlx::query!(
305
-
"SELECT id FROM links WHERE id = $1 AND user_id = $2",
306
-
link_id,
307
-
user.user_id
308
-
)
309
-
.fetch_optional(&mut *tx)
310
-
.await?;
470
+
// Verify the link belongs to the user
471
+
let link = sqlx::query_as::<Postgres, (i32,)>(
472
+
"SELECT id FROM links WHERE id = $1 AND user_id = $2",
473
+
)
474
+
.bind(link_id)
475
+
.bind(user.user_id)
476
+
.fetch_optional(&mut *tx)
477
+
.await?;
311
478
312
-
if link.is_none() {
313
-
return Err(AppError::NotFound);
314
-
}
479
+
if link.is_none() {
480
+
return Err(AppError::NotFound);
481
+
}
315
482
316
-
// Delete associated clicks first due to foreign key constraint
317
-
sqlx::query!("DELETE FROM clicks WHERE link_id = $1", link_id)
318
-
.execute(&mut *tx)
319
-
.await?;
483
+
// Delete associated clicks first due to foreign key constraint
484
+
sqlx::query("DELETE FROM clicks WHERE link_id = $1")
485
+
.bind(link_id)
486
+
.execute(&mut *tx)
487
+
.await?;
320
488
321
-
// Delete the link
322
-
sqlx::query!("DELETE FROM links WHERE id = $1", link_id)
323
-
.execute(&mut *tx)
324
-
.await?;
489
+
// Delete the link
490
+
sqlx::query("DELETE FROM links WHERE id = $1")
491
+
.bind(link_id)
492
+
.execute(&mut *tx)
493
+
.await?;
494
+
495
+
tx.commit().await?;
496
+
}
497
+
DatabasePool::Sqlite(pool) => {
498
+
let mut tx = pool.begin().await?;
499
+
500
+
// Verify the link belongs to the user
501
+
let link = sqlx::query_as::<Sqlite, (i32,)>(
502
+
"SELECT id FROM links WHERE id = ? AND user_id = ?",
503
+
)
504
+
.bind(link_id)
505
+
.bind(user.user_id)
506
+
.fetch_optional(&mut *tx)
507
+
.await?;
508
+
509
+
if link.is_none() {
510
+
return Err(AppError::NotFound);
511
+
}
512
+
513
+
// Delete associated clicks first due to foreign key constraint
514
+
sqlx::query("DELETE FROM clicks WHERE link_id = ?")
515
+
.bind(link_id)
516
+
.execute(&mut *tx)
517
+
.await?;
518
+
519
+
// Delete the link
520
+
sqlx::query("DELETE FROM links WHERE id = ?")
521
+
.bind(link_id)
522
+
.execute(&mut *tx)
523
+
.await?;
325
524
326
-
tx.commit().await?;
525
+
tx.commit().await?;
526
+
}
527
+
}
327
528
328
529
Ok(HttpResponse::NoContent().finish())
329
530
}
···
336
537
let link_id = path.into_inner();
337
538
338
539
// Verify the link belongs to the user
339
-
let link = sqlx::query!(
340
-
"SELECT id FROM links WHERE id = $1 AND user_id = $2",
341
-
link_id,
342
-
user.user_id
343
-
)
344
-
.fetch_optional(&state.db)
345
-
.await?;
540
+
let link = match &state.db {
541
+
DatabasePool::Postgres(pool) => {
542
+
let mut tx = pool.begin().await?;
543
+
let link = sqlx::query_as::<Postgres, (i32,)>(
544
+
"SELECT id FROM links WHERE id = $1 AND user_id = $2",
545
+
)
546
+
.bind(link_id)
547
+
.bind(user.user_id)
548
+
.fetch_optional(&mut *tx)
549
+
.await?;
550
+
tx.commit().await?;
551
+
link
552
+
}
553
+
DatabasePool::Sqlite(pool) => {
554
+
let mut tx = pool.begin().await?;
555
+
let link = sqlx::query_as::<Sqlite, (i32,)>(
556
+
"SELECT id FROM links WHERE id = ? AND user_id = ?",
557
+
)
558
+
.bind(link_id)
559
+
.bind(user.user_id)
560
+
.fetch_optional(&mut *tx)
561
+
.await?;
562
+
tx.commit().await?;
563
+
link
564
+
}
565
+
};
346
566
347
567
if link.is_none() {
348
568
return Err(AppError::NotFound);
349
569
}
350
570
351
-
let clicks = sqlx::query_as!(
352
-
ClickStats,
353
-
r#"
354
-
SELECT
355
-
DATE(created_at)::date as "date!",
356
-
COUNT(*)::bigint as "clicks!"
357
-
FROM clicks
358
-
WHERE link_id = $1
359
-
GROUP BY DATE(created_at)
360
-
ORDER BY DATE(created_at) ASC -- Changed from DESC to ASC
361
-
LIMIT 30
362
-
"#,
363
-
link_id
364
-
)
365
-
.fetch_all(&state.db)
366
-
.await?;
571
+
let clicks = match &state.db {
572
+
DatabasePool::Postgres(pool) => {
573
+
sqlx::query_as::<Postgres, ClickStats>(
574
+
r#"
575
+
SELECT
576
+
DATE(created_at)::date as "date!",
577
+
COUNT(*)::bigint as "clicks!"
578
+
FROM clicks
579
+
WHERE link_id = $1
580
+
GROUP BY DATE(created_at)
581
+
ORDER BY DATE(created_at) ASC
582
+
LIMIT 30
583
+
"#,
584
+
)
585
+
.bind(link_id)
586
+
.fetch_all(pool)
587
+
.await?
588
+
}
589
+
DatabasePool::Sqlite(pool) => {
590
+
sqlx::query_as::<Sqlite, ClickStats>(
591
+
r#"
592
+
SELECT
593
+
DATE(created_at) as "date!",
594
+
COUNT(*) as "clicks!"
595
+
FROM clicks
596
+
WHERE link_id = ?
597
+
GROUP BY DATE(created_at)
598
+
ORDER BY DATE(created_at) ASC
599
+
LIMIT 30
600
+
"#,
601
+
)
602
+
.bind(link_id)
603
+
.fetch_all(pool)
604
+
.await?
605
+
}
606
+
};
367
607
368
608
Ok(HttpResponse::Ok().json(clicks))
369
609
}
···
376
616
let link_id = path.into_inner();
377
617
378
618
// Verify the link belongs to the user
379
-
let link = sqlx::query!(
380
-
"SELECT id FROM links WHERE id = $1 AND user_id = $2",
381
-
link_id,
382
-
user.user_id
383
-
)
384
-
.fetch_optional(&state.db)
385
-
.await?;
619
+
let link = match &state.db {
620
+
DatabasePool::Postgres(pool) => {
621
+
let mut tx = pool.begin().await?;
622
+
let link = sqlx::query_as::<Postgres, (i32,)>(
623
+
"SELECT id FROM links WHERE id = $1 AND user_id = $2",
624
+
)
625
+
.bind(link_id)
626
+
.bind(user.user_id)
627
+
.fetch_optional(&mut *tx)
628
+
.await?;
629
+
tx.commit().await?;
630
+
link
631
+
}
632
+
DatabasePool::Sqlite(pool) => {
633
+
let mut tx = pool.begin().await?;
634
+
let link = sqlx::query_as::<Sqlite, (i32,)>(
635
+
"SELECT id FROM links WHERE id = ? AND user_id = ?",
636
+
)
637
+
.bind(link_id)
638
+
.bind(user.user_id)
639
+
.fetch_optional(&mut *tx)
640
+
.await?;
641
+
tx.commit().await?;
642
+
link
643
+
}
644
+
};
386
645
387
646
if link.is_none() {
388
647
return Err(AppError::NotFound);
389
648
}
390
649
391
-
let sources = sqlx::query_as!(
392
-
SourceStats,
393
-
r#"
394
-
SELECT
395
-
query_source as "source!",
396
-
COUNT(*)::bigint as "count!"
397
-
FROM clicks
398
-
WHERE link_id = $1
399
-
AND query_source IS NOT NULL
400
-
AND query_source != ''
401
-
GROUP BY query_source
402
-
ORDER BY COUNT(*) DESC
403
-
LIMIT 10
404
-
"#,
405
-
link_id
406
-
)
407
-
.fetch_all(&state.db)
408
-
.await?;
650
+
let sources = match &state.db {
651
+
DatabasePool::Postgres(pool) => {
652
+
sqlx::query_as::<Postgres, SourceStats>(
653
+
r#"
654
+
SELECT
655
+
query_source as "source!",
656
+
COUNT(*)::bigint as "count!"
657
+
FROM clicks
658
+
WHERE link_id = $1
659
+
AND query_source IS NOT NULL
660
+
AND query_source != ''
661
+
GROUP BY query_source
662
+
ORDER BY COUNT(*) DESC
663
+
LIMIT 10
664
+
"#,
665
+
)
666
+
.bind(link_id)
667
+
.fetch_all(pool)
668
+
.await?
669
+
}
670
+
DatabasePool::Sqlite(pool) => {
671
+
sqlx::query_as::<Sqlite, SourceStats>(
672
+
r#"
673
+
SELECT
674
+
query_source as "source!",
675
+
COUNT(*) as "count!"
676
+
FROM clicks
677
+
WHERE link_id = ?
678
+
AND query_source IS NOT NULL
679
+
AND query_source != ''
680
+
GROUP BY query_source
681
+
ORDER BY COUNT(*) DESC
682
+
LIMIT 10
683
+
"#,
684
+
)
685
+
.bind(link_id)
686
+
.fetch_all(pool)
687
+
.await?
688
+
}
689
+
};
409
690
410
691
Ok(HttpResponse::Ok().json(sources))
411
692
}
+86
-8
src/lib.rs
+86
-8
src/lib.rs
···
1
+
use anyhow::Result;
1
2
use rand::Rng;
2
-
use sqlx::PgPool;
3
+
use sqlx::migrate::MigrateDatabase;
4
+
use sqlx::postgres::PgPoolOptions;
5
+
use sqlx::{Postgres, Sqlite};
3
6
use std::fs::File;
4
7
use std::io::Write;
5
8
use tracing::info;
6
9
10
+
use models::DatabasePool;
11
+
7
12
pub mod auth;
8
13
pub mod error;
9
14
pub mod handlers;
···
11
16
12
17
#[derive(Clone)]
13
18
pub struct AppState {
14
-
pub db: PgPool,
19
+
pub db: DatabasePool,
15
20
pub admin_token: Option<String>,
16
21
}
17
22
18
-
pub async fn check_and_generate_admin_token(pool: &sqlx::PgPool) -> anyhow::Result<Option<String>> {
23
+
pub async fn create_db_pool() -> Result<DatabasePool> {
24
+
let database_url = std::env::var("DATABASE_URL").ok();
25
+
26
+
match database_url {
27
+
Some(url) if url.starts_with("postgres://") => {
28
+
info!("Using PostgreSQL database");
29
+
let pool = PgPoolOptions::new()
30
+
.max_connections(5)
31
+
.acquire_timeout(std::time::Duration::from_secs(3))
32
+
.connect(&url)
33
+
.await?;
34
+
35
+
Ok(DatabasePool::Postgres(pool))
36
+
}
37
+
_ => {
38
+
info!("No PostgreSQL connection string found, using SQLite");
39
+
40
+
// Create a data directory if it doesn't exist
41
+
let data_dir = std::path::Path::new("data");
42
+
if !data_dir.exists() {
43
+
std::fs::create_dir_all(data_dir)?;
44
+
}
45
+
46
+
let db_path = data_dir.join("simplelink.db");
47
+
let sqlite_url = format!("sqlite://{}", db_path.display());
48
+
49
+
// Check if database exists and create it if it doesn't
50
+
if !Sqlite::database_exists(&sqlite_url).await.unwrap_or(false) {
51
+
info!("Creating new SQLite database at {}", db_path.display());
52
+
Sqlite::create_database(&sqlite_url).await?;
53
+
info!("Database created successfully");
54
+
} else {
55
+
info!("Database already exists");
56
+
}
57
+
58
+
let pool = sqlx::sqlite::SqlitePoolOptions::new()
59
+
.max_connections(5)
60
+
.connect(&sqlite_url)
61
+
.await?;
62
+
63
+
Ok(DatabasePool::Sqlite(pool))
64
+
}
65
+
}
66
+
}
67
+
68
+
pub async fn run_migrations(pool: &DatabasePool) -> Result<()> {
69
+
match pool {
70
+
DatabasePool::Postgres(pool) => {
71
+
// Use the root migrations directory for postgres
72
+
sqlx::migrate!().run(pool).await?;
73
+
}
74
+
DatabasePool::Sqlite(pool) => {
75
+
sqlx::migrate!("./migrations/sqlite").run(pool).await?;
76
+
}
77
+
}
78
+
Ok(())
79
+
}
80
+
81
+
pub async fn check_and_generate_admin_token(db: &DatabasePool) -> anyhow::Result<Option<String>> {
19
82
// Check if any users exist
20
-
let user_count = sqlx::query!("SELECT COUNT(*) as count FROM users")
21
-
.fetch_one(pool)
22
-
.await?
23
-
.count
24
-
.unwrap_or(0);
83
+
let user_count = match db {
84
+
DatabasePool::Postgres(pool) => {
85
+
let mut tx = pool.begin().await?;
86
+
let count = sqlx::query_as::<Postgres, (i64,)>("SELECT COUNT(*)::bigint FROM users")
87
+
.fetch_one(&mut *tx)
88
+
.await?
89
+
.0;
90
+
tx.commit().await?;
91
+
count
92
+
}
93
+
DatabasePool::Sqlite(pool) => {
94
+
let mut tx = pool.begin().await?;
95
+
let count = sqlx::query_as::<Sqlite, (i64,)>("SELECT COUNT(*) FROM users")
96
+
.fetch_one(&mut *tx)
97
+
.await?
98
+
.0;
99
+
tx.commit().await?;
100
+
count
101
+
}
102
+
};
25
103
26
104
if user_count == 0 {
27
105
// Generate a random token using simple characters
+3
-12
src/main.rs
+3
-12
src/main.rs
···
3
3
use anyhow::Result;
4
4
use rust_embed::RustEmbed;
5
5
use simplelink::check_and_generate_admin_token;
6
+
use simplelink::{create_db_pool, run_migrations};
6
7
use simplelink::{handlers, AppState};
7
-
use sqlx::postgres::PgPoolOptions;
8
8
use tracing::info;
9
9
10
10
#[derive(RustEmbed)]
···
31
31
// Initialize logging
32
32
tracing_subscriber::fmt::init();
33
33
34
-
// Database connection string from environment
35
-
let database_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set");
36
-
37
34
// Create database connection pool
38
-
let pool = PgPoolOptions::new()
39
-
.max_connections(5)
40
-
.acquire_timeout(std::time::Duration::from_secs(3))
41
-
.connect(&database_url)
42
-
.await?;
43
-
44
-
// Run database migrations
45
-
sqlx::migrate!("./migrations").run(&pool).await?;
35
+
let pool = create_db_pool().await?;
36
+
run_migrations(&pool).await?;
46
37
47
38
let admin_token = check_and_generate_admin_token(&pool).await?;
48
39
+74
-2
src/models.rs
+74
-2
src/models.rs
···
1
-
use std::time::{SystemTime, UNIX_EPOCH};
2
-
1
+
use anyhow::Result;
3
2
use chrono::NaiveDate;
3
+
use futures::future::BoxFuture;
4
4
use serde::{Deserialize, Serialize};
5
+
use sqlx::postgres::PgRow;
6
+
use sqlx::sqlite::SqliteRow;
5
7
use sqlx::FromRow;
8
+
use sqlx::Pool;
9
+
use sqlx::Postgres;
10
+
use sqlx::Sqlite;
11
+
use sqlx::Transaction;
12
+
use std::time::{SystemTime, UNIX_EPOCH};
13
+
14
+
#[derive(Clone)]
15
+
pub enum DatabasePool {
16
+
Postgres(Pool<Postgres>),
17
+
Sqlite(Pool<Sqlite>),
18
+
}
19
+
20
+
impl DatabasePool {
21
+
pub async fn begin(&self) -> Result<Box<dyn std::any::Any + Send>> {
22
+
match self {
23
+
DatabasePool::Postgres(pool) => Ok(Box::new(pool.begin().await?)),
24
+
DatabasePool::Sqlite(pool) => Ok(Box::new(pool.begin().await?)),
25
+
}
26
+
}
27
+
28
+
pub async fn fetch_optional<T>(&self, pg_query: &str, sqlite_query: &str) -> Result<Option<T>>
29
+
where
30
+
T: for<'r> FromRow<'r, PgRow> + for<'r> FromRow<'r, SqliteRow> + Send + Sync + Unpin,
31
+
{
32
+
match self {
33
+
DatabasePool::Postgres(pool) => {
34
+
Ok(sqlx::query_as(pg_query).fetch_optional(pool).await?)
35
+
}
36
+
DatabasePool::Sqlite(pool) => {
37
+
Ok(sqlx::query_as(sqlite_query).fetch_optional(pool).await?)
38
+
}
39
+
}
40
+
}
41
+
42
+
pub async fn execute(&self, pg_query: &str, sqlite_query: &str) -> Result<()> {
43
+
match self {
44
+
DatabasePool::Postgres(pool) => {
45
+
sqlx::query(pg_query).execute(pool).await?;
46
+
Ok(())
47
+
}
48
+
DatabasePool::Sqlite(pool) => {
49
+
sqlx::query(sqlite_query).execute(pool).await?;
50
+
Ok(())
51
+
}
52
+
}
53
+
}
54
+
55
+
pub async fn transaction<'a, F, R>(&'a self, f: F) -> Result<R>
56
+
where
57
+
F: for<'c> Fn(&'c mut Transaction<'_, Postgres>) -> BoxFuture<'c, Result<R>>
58
+
+ for<'c> Fn(&'c mut Transaction<'_, Sqlite>) -> BoxFuture<'c, Result<R>>
59
+
+ Copy,
60
+
R: Send + 'static,
61
+
{
62
+
match self {
63
+
DatabasePool::Postgres(pool) => {
64
+
let mut tx = pool.begin().await?;
65
+
let result = f(&mut tx).await?;
66
+
tx.commit().await?;
67
+
Ok(result)
68
+
}
69
+
DatabasePool::Sqlite(pool) => {
70
+
let mut tx = pool.begin().await?;
71
+
let result = f(&mut tx).await?;
72
+
tx.commit().await?;
73
+
Ok(result)
74
+
}
75
+
}
76
+
}
77
+
}
6
78
7
79
#[derive(Debug, Serialize, Deserialize)]
8
80
pub struct Claims {