A very performant and light (2mb in memory) link shortener and tracker. Written in Rust and React and uses Postgres/SQLite.
nekomimi.pet
1use anyhow::Result;
2use rand::Rng;
3use sqlx::migrate::MigrateDatabase;
4use sqlx::postgres::PgPoolOptions;
5use sqlx::{Postgres, Sqlite};
6use std::fs::File;
7use std::io::Write;
8use tracing::info;
9
10use models::DatabasePool;
11
12pub mod auth;
13pub mod error;
14pub mod handlers;
15pub mod models;
16
17#[derive(Clone)]
18pub struct AppState {
19 pub db: DatabasePool,
20 pub admin_token: Option<String>,
21}
22
23pub 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
68pub 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
81pub async fn check_and_generate_admin_token(db: &DatabasePool) -> anyhow::Result<Option<String>> {
82 // Check if any users exist
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 };
103
104 if user_count == 0 {
105 // Generate a random token using simple characters
106 let token: String = (0..32)
107 .map(|_| {
108 let idx = rand::thread_rng().gen_range(0..62);
109 match idx {
110 0..=9 => (b'0' + idx as u8) as char,
111 10..=35 => (b'a' + (idx - 10) as u8) as char,
112 _ => (b'A' + (idx - 36) as u8) as char,
113 }
114 })
115 .collect();
116
117 // Save token to file
118 let mut file = File::create("admin-setup-token.txt")?;
119 writeln!(file, "{}", token)?;
120
121 info!("No users found - generated admin setup token");
122 info!("Token has been saved to admin-setup-token.txt");
123 info!("Use this token to create the admin user");
124 info!("Admin setup token: {}", token);
125
126 Ok(Some(token))
127 } else {
128 Ok(None)
129 }
130}