···
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,
···
use jsonwebtoken::{encode, EncodingKey, Header};
use lazy_static::lazy_static;
19
+
use sqlx::{Postgres, Sqlite};
static ref VALID_CODE_REGEX: Regex = Regex::new(r"^[a-zA-Z0-9_-]{1,32}$").unwrap();
···
payload: web::Json<CreateLink>,
) -> Result<impl Responder, AppError> {
tracing::debug!("Creating short URL with user_id: {}", user.user_id);
validate_url(&payload.url)?;
let short_code = if let Some(ref custom_code) = payload.custom_code {
validate_custom_code(custom_code)?;
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")
40
-
.fetch_optional(&state.db)
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")
41
+
.fetch_optional(pool)
44
+
DatabasePool::Sqlite(pool) => {
45
+
sqlx::query_as::<_, Link>("SELECT * FROM links WHERE short_code = ?1")
47
+
.fetch_optional(pool)
52
+
if exists.is_some() {
return Err(AppError::InvalidInput(
"Custom code already taken".to_string(),
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?;
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 *",
63
-
.fetch_one(&mut *tx)
67
+
let link = sqlx::query_as::<_, Link>(
68
+
"INSERT INTO links (original_url, short_code, user_id) VALUES ($1, $2, $3) RETURNING *"
73
+
.fetch_one(&mut *tx)
76
+
if let Some(ref source) = payload.source {
77
+
sqlx::query("INSERT INTO clicks (link_id, source) VALUES ($1, $2)")
87
+
DatabasePool::Sqlite(pool) => {
88
+
let mut tx = pool.begin().await?;
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)")
90
+
let link = sqlx::query_as::<_, Link>(
91
+
"INSERT INTO links (original_url, short_code, user_id) VALUES (?1, ?2, ?3) RETURNING *"
96
+
.fetch_one(&mut *tx)
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)")
107
+
tx.commit().await?;
112
+
Ok(HttpResponse::Created().json(result))
fn validate_custom_code(code: &str) -> Result<(), AppError> {
···
.and_then(|q| web::Query::<std::collections::HashMap<String, String>>::from_query(q).ok())
.and_then(|params| params.get("source").cloned());
123
-
let mut tx = state.db.begin().await?;
125
-
let link = sqlx::query_as::<_, Link>(
126
-
"UPDATE links SET clicks = clicks + 1 WHERE short_code = $1 RETURNING *",
129
-
.fetch_optional(&mut *tx)
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 *",
166
+
.fetch_optional(&mut *tx)
168
+
tx.commit().await?;
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 *",
177
+
.fetch_optional(&mut *tx)
179
+
tx.commit().await?;
134
-
// Record click with both user agent and query source
135
-
let user_agent = req
138
-
.and_then(|h| h.to_str().ok())
139
-
.unwrap_or("unknown")
186
+
// Handle click recording based on database type
188
+
DatabasePool::Postgres(pool) => {
189
+
let mut tx = pool.begin().await?;
190
+
let user_agent = req
193
+
.and_then(|h| h.to_str().ok())
194
+
.unwrap_or("unknown")
198
+
"INSERT INTO clicks (link_id, source, query_source) VALUES ($1, $2, $3)",
202
+
.bind(query_source)
206
+
tx.commit().await?;
208
+
DatabasePool::Sqlite(pool) => {
209
+
let mut tx = pool.begin().await?;
210
+
let user_agent = req
213
+
.and_then(|h| h.to_str().ok())
214
+
.unwrap_or("unknown")
142
-
sqlx::query("INSERT INTO clicks (link_id, source, query_source) VALUES ($1, $2, $3)")
145
-
.bind(query_source)
218
+
"INSERT INTO clicks (link_id, source, query_source) VALUES (?1, ?2, ?3)",
222
+
.bind(query_source)
149
-
tx.commit().await?;
226
+
tx.commit().await?;
Ok(HttpResponse::TemporaryRedirect()
.append_header(("Location", link.original_url))
···
state: web::Data<AppState>,
) -> Result<impl Responder, AppError> {
163
-
let links = sqlx::query_as::<_, Link>(
164
-
"SELECT * FROM links WHERE user_id = $1 ORDER BY created_at DESC",
166
-
.bind(user.user_id)
167
-
.fetch_all(&state.db)
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",
247
+
.bind(user.user_id)
251
+
DatabasePool::Sqlite(pool) => {
252
+
sqlx::query_as::<_, Link>(
253
+
"SELECT * FROM links WHERE user_id = ?1 ORDER BY created_at DESC",
255
+
.bind(user.user_id)
Ok(HttpResponse::Ok().json(links))
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(),
271
+
HttpResponse::Ok().json("Healthy")
273
+
HttpResponse::ServiceUnavailable().json("Database unavailable")
···
payload: web::Json<RegisterRequest>,
) -> Result<impl Responder, AppError> {
// Check if any users exist
193
-
let user_count = sqlx::query!("SELECT COUNT(*) as count FROM users")
194
-
.fetch_one(&state.db)
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)
297
+
tx.commit().await?;
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)
306
+
tx.commit().await?;
// If users exist, registration is closed - no exceptions
···
// Check if email already exists
213
-
let exists = sqlx::query!("SELECT id FROM users WHERE email = $1", payload.email)
214
-
.fetch_optional(&state.db)
325
+
let exists = match &state.db {
326
+
DatabasePool::Postgres(pool) => {
327
+
let mut tx = pool.begin().await?;
329
+
sqlx::query_as::<Postgres, (i32,)>("SELECT id FROM users WHERE email = $1")
330
+
.bind(&payload.email)
331
+
.fetch_optional(&mut *tx)
333
+
tx.commit().await?;
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)
342
+
tx.commit().await?;
return Err(AppError::Auth("Email already registered".to_string()));
···
.map_err(|e| AppError::Auth(e.to_string()))?
228
-
let user = sqlx::query_as!(
230
-
"INSERT INTO users (email, password_hash) VALUES ($1, $2) RETURNING *",
234
-
.fetch_one(&state.db)
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 *",
365
+
.bind(&payload.email)
366
+
.bind(&password_hash)
367
+
.fetch_one(&mut *tx)
369
+
tx.commit().await?;
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 *",
377
+
.bind(&payload.email)
378
+
.bind(&password_hash)
379
+
.fetch_one(&mut *tx)
381
+
tx.commit().await?;
let claims = Claims::new(user.id);
let secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| "default_secret".to_string());
···
state: web::Data<AppState>,
payload: web::Json<LoginRequest>,
) -> Result<impl Responder, AppError> {
259
-
let user = sqlx::query_as!(User, "SELECT * FROM users WHERE email = $1", payload.email)
260
-
.fetch_optional(&state.db)
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)
415
+
tx.commit().await?;
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)
424
+
tx.commit().await?;
428
+
.ok_or_else(|| AppError::Auth("Invalid credentials".to_string()))?;
let argon2 = Argon2::default();
···
) -> Result<impl Responder, AppError> {
let link_id = path.into_inner();
300
-
// Start transaction
301
-
let mut tx = state.db.begin().await?;
467
+
DatabasePool::Postgres(pool) => {
468
+
let mut tx = pool.begin().await?;
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",
309
-
.fetch_optional(&mut *tx)
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",
475
+
.bind(user.user_id)
476
+
.fetch_optional(&mut *tx)
312
-
if link.is_none() {
313
-
return Err(AppError::NotFound);
479
+
if link.is_none() {
480
+
return Err(AppError::NotFound);
316
-
// Delete associated clicks first due to foreign key constraint
317
-
sqlx::query!("DELETE FROM clicks WHERE link_id = $1", link_id)
483
+
// Delete associated clicks first due to foreign key constraint
484
+
sqlx::query("DELETE FROM clicks WHERE link_id = $1")
322
-
sqlx::query!("DELETE FROM links WHERE id = $1", link_id)
490
+
sqlx::query("DELETE FROM links WHERE id = $1")
495
+
tx.commit().await?;
497
+
DatabasePool::Sqlite(pool) => {
498
+
let mut tx = pool.begin().await?;
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 = ?",
505
+
.bind(user.user_id)
506
+
.fetch_optional(&mut *tx)
509
+
if link.is_none() {
510
+
return Err(AppError::NotFound);
513
+
// Delete associated clicks first due to foreign key constraint
514
+
sqlx::query("DELETE FROM clicks WHERE link_id = ?")
520
+
sqlx::query("DELETE FROM links WHERE id = ?")
326
-
tx.commit().await?;
525
+
tx.commit().await?;
Ok(HttpResponse::NoContent().finish())
···
let link_id = path.into_inner();
// Verify the link belongs to the user
339
-
let link = sqlx::query!(
340
-
"SELECT id FROM links WHERE id = $1 AND user_id = $2",
344
-
.fetch_optional(&state.db)
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",
547
+
.bind(user.user_id)
548
+
.fetch_optional(&mut *tx)
550
+
tx.commit().await?;
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 = ?",
559
+
.bind(user.user_id)
560
+
.fetch_optional(&mut *tx)
562
+
tx.commit().await?;
return Err(AppError::NotFound);
351
-
let clicks = sqlx::query_as!(
355
-
DATE(created_at)::date as "date!",
356
-
COUNT(*)::bigint as "clicks!"
359
-
GROUP BY DATE(created_at)
360
-
ORDER BY DATE(created_at) ASC -- Changed from DESC to ASC
365
-
.fetch_all(&state.db)
571
+
let clicks = match &state.db {
572
+
DatabasePool::Postgres(pool) => {
573
+
sqlx::query_as::<Postgres, ClickStats>(
576
+
DATE(created_at)::date as "date!",
577
+
COUNT(*)::bigint as "clicks!"
580
+
GROUP BY DATE(created_at)
581
+
ORDER BY DATE(created_at) ASC
589
+
DatabasePool::Sqlite(pool) => {
590
+
sqlx::query_as::<Sqlite, ClickStats>(
593
+
DATE(created_at) as "date!",
594
+
COUNT(*) as "clicks!"
597
+
GROUP BY DATE(created_at)
598
+
ORDER BY DATE(created_at) ASC
Ok(HttpResponse::Ok().json(clicks))
···
let link_id = path.into_inner();
// Verify the link belongs to the user
379
-
let link = sqlx::query!(
380
-
"SELECT id FROM links WHERE id = $1 AND user_id = $2",
384
-
.fetch_optional(&state.db)
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",
626
+
.bind(user.user_id)
627
+
.fetch_optional(&mut *tx)
629
+
tx.commit().await?;
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 = ?",
638
+
.bind(user.user_id)
639
+
.fetch_optional(&mut *tx)
641
+
tx.commit().await?;
return Err(AppError::NotFound);
391
-
let sources = sqlx::query_as!(
395
-
query_source as "source!",
396
-
COUNT(*)::bigint as "count!"
399
-
AND query_source IS NOT NULL
400
-
AND query_source != ''
401
-
GROUP BY query_source
402
-
ORDER BY COUNT(*) DESC
407
-
.fetch_all(&state.db)
650
+
let sources = match &state.db {
651
+
DatabasePool::Postgres(pool) => {
652
+
sqlx::query_as::<Postgres, SourceStats>(
655
+
query_source as "source!",
656
+
COUNT(*)::bigint as "count!"
659
+
AND query_source IS NOT NULL
660
+
AND query_source != ''
661
+
GROUP BY query_source
662
+
ORDER BY COUNT(*) DESC
670
+
DatabasePool::Sqlite(pool) => {
671
+
sqlx::query_as::<Sqlite, SourceStats>(
674
+
query_source as "source!",
675
+
COUNT(*) as "count!"
678
+
AND query_source IS NOT NULL
679
+
AND query_source != ''
680
+
GROUP BY query_source
681
+
ORDER BY COUNT(*) DESC
Ok(HttpResponse::Ok().json(sources))
nekomimi.pet