Skip to content

Commit

Permalink
improvement: add url-shortener example to Axum examples
Browse files Browse the repository at this point in the history
  • Loading branch information
sentinel1909 committed Feb 4, 2025
1 parent d14e062 commit db48e81
Show file tree
Hide file tree
Showing 8 changed files with 238 additions and 0 deletions.
24 changes: 24 additions & 0 deletions axum/url-shortener/Cargo.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
[package]
name = "url-shortener"
version = "0.1.0"
edition = "2021"

[dependencies]
axum = "0.8.1"
axum-extra = { version = "0.10.0", features = [ "typed-header" ] }
axum-macros = "0.5.0"
nanoid = "0.4.0"
regex = "1.11.1"
shuttle-axum = "0.51.0"
shuttle-runtime = { version = "0.51.0", default-features = false }
shuttle-shared-db = { version = "0.51.0", features = ["postgres", "sqlx"] }
sqlx = "0.8.3"
tokio = "1.43.0"
tower = "0.5.2"
tower-http = { version = "0.6.2", features = ["request-id", "trace", "util"] }
tracing = { version = "0.1.41", features = ["log"] }
tracing-bunyan-formatter = "0.3.10"
tracing-log = "0.2.0"
tracing-subscriber = { version = "0.3.19", features = ["registry", "env-filter"] }
url = "2.5.4"
uuid = { version = "1.12.1", features = ["v4"] }
13 changes: 13 additions & 0 deletions axum/url-shortener/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
# URL Shortener

A URL shortener that you can use from your terminal - built with Shuttle, Axum, and Postgres.

## How to use it

POST a URL like so:

```bash
curl -d 'https://shuttle.dev/' http://localhost:8000/
```

You will get the shortened URL back (something like http://localhost:8000/0fvAo2). Visiting it will redirect you to the original URL.
2 changes: 2 additions & 0 deletions axum/url-shortener/migrations/20250201_url-shortener.down.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
-- Add down migration script here
DROP TABLE IF EXISTS urls;
5 changes: 5 additions & 0 deletions axum/url-shortener/migrations/20250201_url-shortener.up.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
-- Add up migration script here
CREATE TABLE urls (
id VARCHAR(6) PRIMARY KEY,
url VARCHAR NOT NULL
);
32 changes: 32 additions & 0 deletions axum/url-shortener/src/main.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
// src/bin/main.rs

use shuttle_runtime::CustomError;
use telemetry::{get_subscriber, init_subscriber};
use sqlx::PgPool;
use startup::Application;

mod routes;
mod startup;
mod telemetry;

#[shuttle_runtime::main]
async fn main(#[shuttle_shared_db::Postgres] pool: PgPool) -> shuttle_axum::ShuttleAxum {
// initialize tracing
let subscriber = get_subscriber("url-shortener-v1".into(), "info".into(), std::io::stdout);
init_subscriber(subscriber);

// run the database migrations
tracing::info!("Running database migrations...");
sqlx::migrate!("./migrations")
.run(&pool)
.await
.map_err(|err| {
let msg = format!("Unable to run the database migrations: {}", err);
CustomError::new(err).context(msg)
})?;

tracing::info!("Building the application...");
let Application(router) = Application::build(pool);

Ok(router.into())
}
65 changes: 65 additions & 0 deletions axum/url-shortener/src/routes.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
// src/lib/routes/routes.rs

use axum::{
extract::{Path, State},
http::StatusCode,
response::{IntoResponse, Redirect},
};
use axum_extra::{headers::Host, TypedHeader};
use axum_macros::debug_handler;
use sqlx::{Error, PgPool};
use tracing::instrument;
use url::Url;

// health_check handler
pub async fn health_check() -> impl IntoResponse {
StatusCode::OK
}

// redirect endpoint handler
#[debug_handler]
#[instrument(name = "redirect" skip(state))]
pub async fn get_redirect(
State(state): State<PgPool>,
Path(id): Path<String>,
) -> Result<impl IntoResponse, StatusCode> {
let url: (String,) = sqlx::query_as("SELECT url FROM urls WHERE id = $1")
.bind(id)
.fetch_one(&state)
.await
.map_err(|e| match e {
Error::RowNotFound => {
tracing::error!("shortened URL not found in the database...");
StatusCode::NOT_FOUND
}
_ => StatusCode::INTERNAL_SERVER_ERROR,
})?;
tracing::info!("shortened URL retrieved, redirecting...");
Ok(Redirect::permanent(&url.0))
}

// shorten endpoint handler
#[debug_handler]
#[instrument(name = "shorten" skip(state))]
pub async fn post_shorten(
State(state): State<PgPool>,
TypedHeader(header): TypedHeader<Host>,
url: String,
) -> Result<impl IntoResponse, StatusCode> {
let id = &nanoid::nanoid!(6);
let p_url = Url::parse(&url).map_err(|_| {
tracing::error!("Unable to parse URL");
StatusCode::UNPROCESSABLE_ENTITY
})?;
let host = header.hostname();
sqlx::query("INSERT INTO urls (id, url) VALUES ($1, $2)")
.bind(id)
.bind(p_url.as_str())
.execute(&state)
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let response_body = format!("https://{}/{}\n", host, id);

tracing::info!("URL shortened and saved successfully...");
Ok((StatusCode::OK, response_body).into_response())
}
51 changes: 51 additions & 0 deletions axum/url-shortener/src/startup.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
// src/lib/startup.rs

use crate::routes::{get_redirect, health_check, post_shorten};
use crate::telemetry::MakeRequestUuid;
use axum::{
http::HeaderName,
routing::{get, post},
Router,
};
use sqlx::PgPool;
use tower::ServiceBuilder;
use tower_http::{
request_id::{PropagateRequestIdLayer, SetRequestIdLayer},
trace::{DefaultMakeSpan, DefaultOnResponse, TraceLayer},
};
use tracing::Level;

pub struct Application(pub Router);

impl Application {
// builds the router for the application
pub fn build(pool: PgPool) -> Self {
// define the tracing layer
let trace_layer = TraceLayer::new_for_http()
.make_span_with(
DefaultMakeSpan::new()
.include_headers(true)
.level(Level::INFO),
)
.on_response(DefaultOnResponse::new().include_headers(true));
let x_request_id = HeaderName::from_static("x-request-id");

// build the router, with state and tracing
let router = Router::new()
.route("/health_check", get(health_check))
.route("/{id}", get(get_redirect))
.route("/", post(post_shorten))
.with_state(pool)
.layer(
ServiceBuilder::new()
.layer(SetRequestIdLayer::new(
x_request_id.clone(),
MakeRequestUuid,
))
.layer(trace_layer)
.layer(PropagateRequestIdLayer::new(x_request_id)),
);

Self(router)
}
}
46 changes: 46 additions & 0 deletions axum/url-shortener/src/telemetry.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
// src/lib/telemetry.rs

use axum::http::Request;
use tower_http::request_id::{MakeRequestId, RequestId};
use tracing::subscriber::set_global_default;
use tracing::Subscriber;
use tracing_bunyan_formatter::{BunyanFormattingLayer, JsonStorageLayer};
use tracing_log::LogTracer;
use tracing_subscriber::fmt::MakeWriter;
use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Registry};
use uuid::Uuid;

#[derive(Clone)]
pub struct MakeRequestUuid;

impl MakeRequestId for MakeRequestUuid {
fn make_request_id<B>(&mut self, _: &Request<B>) -> Option<RequestId> {
let request_id = Uuid::new_v4().to_string();

Some(RequestId::new(request_id.parse().unwrap()))
}
}

pub fn get_subscriber<Sink>(
name: String,
env_filter: String,
sink: Sink,
) -> impl Subscriber + Sync + Send
where
Sink: for<'a> MakeWriter<'a> + Send + Sync + 'static,
{
let env_filter =
EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(env_filter));
let formatting_layer = BunyanFormattingLayer::new(name, sink);

Registry::default()
.with(env_filter)
.with(JsonStorageLayer)
.with(formatting_layer)
}

pub fn init_subscriber(subscriber: impl Subscriber + Sync + Send) {
// Redirect logs to subscriber
LogTracer::init().expect("Failed to set logger");
set_global_default(subscriber).expect("Failed to set subscriber");
}

0 comments on commit db48e81

Please sign in to comment.