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://") || url.starts_with("postgresql://") => {
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 // Get the project root directory
41 let project_root = std::env::current_dir()?;
42 let data_dir = project_root.join("data");
43
44 // Create a data directory if it doesn't exist
45 if !data_dir.exists() {
46 std::fs::create_dir_all(&data_dir)?;
47 }
48
49 let db_path = data_dir.join("simplelink.db");
50 let sqlite_url = format!("sqlite://{}", db_path.display());
51
52 // Check if database exists and create it if it doesn't
53 if !Sqlite::database_exists(&sqlite_url).await.unwrap_or(false) {
54 info!("Creating new SQLite database at {}", db_path.display());
55 Sqlite::create_database(&sqlite_url).await?;
56 info!("Database created successfully");
57 } else {
58 info!("Database already exists");
59 }
60
61 let pool = sqlx::sqlite::SqlitePoolOptions::new()
62 .max_connections(5)
63 .connect(&sqlite_url)
64 .await?;
65
66 Ok(DatabasePool::Sqlite(pool))
67 }
68 }
69}
70
71pub async fn run_migrations(pool: &DatabasePool) -> Result<()> {
72 match pool {
73 DatabasePool::Postgres(pool) => {
74 // Use the root migrations directory for postgres
75 sqlx::migrate!().run(pool).await?;
76 }
77 DatabasePool::Sqlite(pool) => {
78 sqlx::migrate!("./migrations/sqlite").run(pool).await?;
79 }
80 }
81 Ok(())
82}
83
84pub async fn check_and_generate_admin_token(db: &DatabasePool) -> anyhow::Result<Option<String>> {
85 // Check if any users exist
86 let user_count = match db {
87 DatabasePool::Postgres(pool) => {
88 let mut tx = pool.begin().await?;
89 let count = sqlx::query_as::<Postgres, (i64,)>("SELECT COUNT(*)::bigint FROM users")
90 .fetch_one(&mut *tx)
91 .await?
92 .0;
93 tx.commit().await?;
94 count
95 }
96 DatabasePool::Sqlite(pool) => {
97 let mut tx = pool.begin().await?;
98 let count = sqlx::query_as::<Sqlite, (i64,)>("SELECT COUNT(*) FROM users")
99 .fetch_one(&mut *tx)
100 .await?
101 .0;
102 tx.commit().await?;
103 count
104 }
105 };
106
107 if user_count == 0 {
108 let token: String = (0..32)
109 .map(|_| {
110 let idx = rand::thread_rng().gen_range(0..62);
111 match idx {
112 0..=9 => (b'0' + idx as u8) as char,
113 10..=35 => (b'a' + (idx - 10) as u8) as char,
114 _ => (b'A' + (idx - 36) as u8) as char,
115 }
116 })
117 .collect();
118
119 // Get the project root directory
120 let project_root = std::env::current_dir()?;
121 let token_path = project_root.join("admin-setup-token.txt");
122
123 // Save token to file
124 let mut file = File::create(token_path)?;
125 writeln!(file, "{}", token)?;
126
127 info!("No users found - generated admin setup token");
128 info!("Token has been saved to admin-setup-token.txt");
129 info!("Use this token to create the admin user");
130 info!("Admin setup token: {}", token);
131
132 Ok(Some(token))
133 } else {
134 Ok(None)
135 }
136}