1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
|
use crate::parsing::get_hash_from_string;
use crate::schema::links;
use diesel::r2d2::{ConnectionManager, Pool};
use diesel::{
ExpressionMethods, OptionalExtension, PgConnection, QueryDsl, QueryResult, RunQueryDsl,
};
embed_migrations!("migrations");
#[derive(Queryable, Insertable)]
#[table_name = "links"]
pub struct ShortenedLink {
pub hash: String,
pub url: String,
}
// Creates a connection pool to the database
pub fn establish_connection() -> Pool<ConnectionManager<PgConnection>> {
let password =
std::env::var("POSTGRES_PASSWORD").expect("Cannot find POSTGRES_PASSWORD in environment.");
let database_url = format!("postgresql://postgres:{}@postgres/shorest", password);
let manager = ConnectionManager::<PgConnection>::new(database_url);
Pool::builder()
.build(manager)
.expect("Failed to create pool.")
}
/// Gets a link from the database.
pub fn get_link_from_database(
url_hash: &String,
connection: &PgConnection,
) -> actix_web::Result<Option<String>, diesel::result::Error> {
links::table
.filter(links::hash.eq(url_hash))
.select(links::url)
.first(connection)
.optional()
}
/// Inserts a link into the database.
fn add_link_to_database(
entry: ShortenedLink,
connection: &PgConnection,
) -> QueryResult<ShortenedLink> {
diesel::insert_into(links::table)
.values(&entry)
.get_result::<ShortenedLink>(connection)
}
/// Inserts a new link into the database, avoiding collisions.
pub fn add_link_to_database_safely(
mut hash: String,
user_url: &str,
connection: &PgConnection,
) -> Result<String, diesel::result::Error> {
// Check whether a link with the same hash exists within the database
match get_link_from_database(&hash, connection)? {
Some(other_url) => {
// We found a collision!
if user_url != other_url {
warn!("'{}' and '{}' have colliding hashes.", user_url, other_url);
// Try inserting again using the hash of the collided hash
hash =
add_link_to_database_safely(get_hash_from_string(&hash), user_url, connection)?;
}
}
None => {
// No link with the same hash exists, put it in the database
add_link_to_database(
ShortenedLink {
hash: hash.clone(),
url: user_url.to_string(),
},
connection,
)?;
}
};
Ok(hash)
}
pub fn run_migrations(pool: &PoolState) {
let migration_connection = &pool
.get()
.expect("Could not connect to database to run migrations");
embedded_migrations::run_with_output(migration_connection, &mut std::io::stdout())
.expect("Failed to run migrations.");
}
pub type PoolState = Pool<ConnectionManager<PgConnection>>;
|
