···
+
AuthResponse, Claims, ClickStats, CreateLink, DatabasePool, Link, LoginRequest,
+
RegisterRequest, SourceStats, User, UserResponse,
···
use jsonwebtoken::{encode, EncodingKey, Header};
use lazy_static::lazy_static;
+
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)?;
+
// Check if code exists using match on pool type
+
let exists = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
sqlx::query_as::<_, Link>("SELECT * FROM links WHERE short_code = $1")
+
DatabasePool::Sqlite(pool) => {
+
sqlx::query_as::<_, Link>("SELECT * FROM links WHERE short_code = ?1")
return Err(AppError::InvalidInput(
"Custom code already taken".to_string(),
+
// Start transaction based on pool type
+
let result = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
let mut tx = pool.begin().await?;
+
let link = sqlx::query_as::<_, Link>(
+
"INSERT INTO links (original_url, short_code, user_id) VALUES ($1, $2, $3) RETURNING *"
+
if let Some(ref source) = payload.source {
+
sqlx::query("INSERT INTO clicks (link_id, source) VALUES ($1, $2)")
+
DatabasePool::Sqlite(pool) => {
+
let mut tx = pool.begin().await?;
+
let link = sqlx::query_as::<_, Link>(
+
"INSERT INTO links (original_url, short_code, user_id) VALUES (?1, ?2, ?3) RETURNING *"
+
if let Some(ref source) = payload.source {
+
sqlx::query("INSERT INTO clicks (link_id, source) VALUES (?1, ?2)")
+
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());
+
let link = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
let mut tx = pool.begin().await?;
+
let link = sqlx::query_as::<_, Link>(
+
"UPDATE links SET clicks = clicks + 1 WHERE short_code = $1 RETURNING *",
+
.fetch_optional(&mut *tx)
+
DatabasePool::Sqlite(pool) => {
+
let mut tx = pool.begin().await?;
+
let link = sqlx::query_as::<_, Link>(
+
"UPDATE links SET clicks = clicks + 1 WHERE short_code = ?1 RETURNING *",
+
.fetch_optional(&mut *tx)
+
// Handle click recording based on database type
+
DatabasePool::Postgres(pool) => {
+
let mut tx = pool.begin().await?;
+
.and_then(|h| h.to_str().ok())
+
"INSERT INTO clicks (link_id, source, query_source) VALUES ($1, $2, $3)",
+
DatabasePool::Sqlite(pool) => {
+
let mut tx = pool.begin().await?;
+
.and_then(|h| h.to_str().ok())
+
"INSERT INTO clicks (link_id, source, query_source) VALUES (?1, ?2, ?3)",
Ok(HttpResponse::TemporaryRedirect()
.append_header(("Location", link.original_url))
···
state: web::Data<AppState>,
) -> Result<impl Responder, AppError> {
+
let links = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
sqlx::query_as::<_, Link>(
+
"SELECT * FROM links WHERE user_id = $1 ORDER BY created_at DESC",
+
DatabasePool::Sqlite(pool) => {
+
sqlx::query_as::<_, Link>(
+
"SELECT * FROM links WHERE user_id = ?1 ORDER BY created_at DESC",
Ok(HttpResponse::Ok().json(links))
pub async fn health_check(state: web::Data<AppState>) -> impl Responder {
+
let is_healthy = match &state.db {
+
DatabasePool::Postgres(pool) => sqlx::query("SELECT 1").execute(pool).await.is_ok(),
+
DatabasePool::Sqlite(pool) => sqlx::query("SELECT 1").execute(pool).await.is_ok(),
+
HttpResponse::Ok().json("Healthy")
+
HttpResponse::ServiceUnavailable().json("Database unavailable")
···
payload: web::Json<RegisterRequest>,
) -> Result<impl Responder, AppError> {
// Check if any users exist
+
let user_count = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
let mut tx = pool.begin().await?;
+
let count = sqlx::query_as::<Postgres, (i64,)>("SELECT COUNT(*)::bigint FROM users")
+
DatabasePool::Sqlite(pool) => {
+
let mut tx = pool.begin().await?;
+
let count = sqlx::query_as::<Sqlite, (i64,)>("SELECT COUNT(*) FROM users")
// If users exist, registration is closed - no exceptions
···
// Check if email already exists
+
let exists = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
let mut tx = pool.begin().await?;
+
sqlx::query_as::<Postgres, (i32,)>("SELECT id FROM users WHERE email = $1")
+
.fetch_optional(&mut *tx)
+
DatabasePool::Sqlite(pool) => {
+
let mut tx = pool.begin().await?;
+
let exists = sqlx::query_as::<Sqlite, (i32,)>("SELECT id FROM users WHERE email = ?")
+
.fetch_optional(&mut *tx)
return Err(AppError::Auth("Email already registered".to_string()));
···
.map_err(|e| AppError::Auth(e.to_string()))?
+
let user = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
let mut tx = pool.begin().await?;
+
let user = sqlx::query_as::<Postgres, User>(
+
"INSERT INTO users (email, password_hash) VALUES ($1, $2) RETURNING *",
+
DatabasePool::Sqlite(pool) => {
+
let mut tx = pool.begin().await?;
+
let user = sqlx::query_as::<Sqlite, User>(
+
"INSERT INTO users (email, password_hash) VALUES (?, ?) RETURNING *",
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> {
+
let user = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
let mut tx = pool.begin().await?;
+
let user = sqlx::query_as::<Postgres, User>("SELECT * FROM users WHERE email = $1")
+
.fetch_optional(&mut *tx)
+
DatabasePool::Sqlite(pool) => {
+
let mut tx = pool.begin().await?;
+
let user = sqlx::query_as::<Sqlite, User>("SELECT * FROM users WHERE email = ?")
+
.fetch_optional(&mut *tx)
+
.ok_or_else(|| AppError::Auth("Invalid credentials".to_string()))?;
let argon2 = Argon2::default();
···
) -> Result<impl Responder, AppError> {
let link_id = path.into_inner();
+
DatabasePool::Postgres(pool) => {
+
let mut tx = pool.begin().await?;
+
// Verify the link belongs to the user
+
let link = sqlx::query_as::<Postgres, (i32,)>(
+
"SELECT id FROM links WHERE id = $1 AND user_id = $2",
+
.fetch_optional(&mut *tx)
+
return Err(AppError::NotFound);
+
// Delete associated clicks first due to foreign key constraint
+
sqlx::query("DELETE FROM clicks WHERE link_id = $1")
+
sqlx::query("DELETE FROM links WHERE id = $1")
+
DatabasePool::Sqlite(pool) => {
+
let mut tx = pool.begin().await?;
+
// Verify the link belongs to the user
+
let link = sqlx::query_as::<Sqlite, (i32,)>(
+
"SELECT id FROM links WHERE id = ? AND user_id = ?",
+
.fetch_optional(&mut *tx)
+
return Err(AppError::NotFound);
+
// Delete associated clicks first due to foreign key constraint
+
sqlx::query("DELETE FROM clicks WHERE link_id = ?")
+
sqlx::query("DELETE FROM links WHERE id = ?")
Ok(HttpResponse::NoContent().finish())
···
let link_id = path.into_inner();
// Verify the link belongs to the user
+
let link = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
let mut tx = pool.begin().await?;
+
let link = sqlx::query_as::<Postgres, (i32,)>(
+
"SELECT id FROM links WHERE id = $1 AND user_id = $2",
+
.fetch_optional(&mut *tx)
+
DatabasePool::Sqlite(pool) => {
+
let mut tx = pool.begin().await?;
+
let link = sqlx::query_as::<Sqlite, (i32,)>(
+
"SELECT id FROM links WHERE id = ? AND user_id = ?",
+
.fetch_optional(&mut *tx)
return Err(AppError::NotFound);
+
let clicks = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
sqlx::query_as::<Postgres, ClickStats>(
+
DATE(created_at)::date as "date!",
+
COUNT(*)::bigint as "clicks!"
+
GROUP BY DATE(created_at)
+
ORDER BY DATE(created_at) ASC
+
DatabasePool::Sqlite(pool) => {
+
sqlx::query_as::<Sqlite, ClickStats>(
+
DATE(created_at) as "date!",
+
GROUP BY DATE(created_at)
+
ORDER BY DATE(created_at) ASC
Ok(HttpResponse::Ok().json(clicks))
···
let link_id = path.into_inner();
// Verify the link belongs to the user
+
let link = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
let mut tx = pool.begin().await?;
+
let link = sqlx::query_as::<Postgres, (i32,)>(
+
"SELECT id FROM links WHERE id = $1 AND user_id = $2",
+
.fetch_optional(&mut *tx)
+
DatabasePool::Sqlite(pool) => {
+
let mut tx = pool.begin().await?;
+
let link = sqlx::query_as::<Sqlite, (i32,)>(
+
"SELECT id FROM links WHERE id = ? AND user_id = ?",
+
.fetch_optional(&mut *tx)
return Err(AppError::NotFound);
+
let sources = match &state.db {
+
DatabasePool::Postgres(pool) => {
+
sqlx::query_as::<Postgres, SourceStats>(
+
query_source as "source!",
+
COUNT(*)::bigint as "count!"
+
AND query_source IS NOT NULL
+
DatabasePool::Sqlite(pool) => {
+
sqlx::query_as::<Sqlite, SourceStats>(
+
query_source as "source!",
+
AND query_source IS NOT NULL
Ok(HttpResponse::Ok().json(sources))
nekomimi.pet