使用Rust和Axum实现整洁代码 - PropelAuth


在《Clean Code》一书中,我最喜欢的部分之一是看到一个代码片段开始时很粗糙且难以管理,然后看着它迭代改进。

在这篇文章中,我们将做同样的事情,但特别使用 Rust 和 Axum 的接近真实世界的示例。对于我们所做的每一次重构,我们还会指出为什么这种更改会改进代码。

我们将从一个单一的、混乱的 API 路由开始:

#[derive(Deserialize)]
struct CreateUrl {
    url: String,
}

#[derive(Serialize)]
struct CreateUrlResponse {
    id: String,
}

// Save a URL and return an ID associated with it
async fn save_url(
    headers: HeaderMap,
    State(pool): State<PgPool>,
    Json(create_url): Json<CreateUrl>,
) -> Response {
    // First grab an auth token from a custom header
    let token_header = headers.get("X-Auth-Token");
    let token = match token_header {
        None => return (StatusCode::UNAUTHORIZED, "Unauthorized").into_response(),
        Some(header) => header.to_str().unwrap(),
    };

    // Then verify the token is correct
    let verify_result = verify_auth_token(token).await;
    let user = match verify_result {
        Some(user) => user,
        None => return (StatusCode::UNAUTHORIZED, "Unauthorized").into_response(),
    };

    // Insert our URL into the database and get back an ID that the database generated
    let insert_result =
        sqlx::query!(
            "INSERT INTO urls (url, user_id) VALUE (lower($1), $2) RETURNING id",
            create_url.url, user.user_id()
        )
            .map(|row| row.id)
            .fetch_one(pool)
            .await;
    let id = match insert_result {
        Ok(id) => id,
        Err(_) => {
            return (
                StatusCode::INTERNAL_SERVER_ERROR,
                "An unexpected exception has occurred",
            )
                .into_response()
        }
    };

    (StatusCode::CREATED, Json(CreateUrlResponse { id })).into_response()
}

花点时间想想你不喜欢那个代码片段的所有事情,然后让我们深入研究并让它变得更好!

来自请求
您可能首先注意到的是,有 10 行代码专门用于获取标X-Auth-Token头并对其进行验证。解决这个问题的第一个合理尝试是将其转换为一种方法,以便可以重复使用:

pub async fn extract_and_verify_auth_token(headers: HeaderMap) -> Option<User>

我们可以称之为:

async fn save_url(
    headers: HeaderMap,
    State(pool): State<PgPool>,
    Json(create_url): Json<CreateUrl>,
) -> Response {
    let user = match extract_and_verify_auth_token(headers) {
        Some(user) => user,
        None => return (StatusCode::UNAUTHORIZED, "Unauthorized").into_response(),
    };
    // ...
}

但我们实际上可以做得更好。如果我们的函数签名是这样的:

async fn save_url(
    user: User, // <-- new
    State(pool): State<PgPool>,
    Json(create_url): Json<CreateUrl>
) -> Response

X-Auth-Token如果存在有效标头,我们会自动提取并验证它。如果X-Auth-Token不存在有效的标头,我们甚至都不会调用save_url.

这就是本质上的FromRequest作用。Axum 实际上两者都有FromRequest,FromRequestParts不同之处在于FromRequest会消耗请求的主体。因为我们只需要一个标题,所以我们可以使用FromRequestParts:

#[async_trait]
impl<S> FromRequestParts<S> for User
    where
        S: Send + Sync,
{
    type Rejection = (StatusCode, &'static str);

    async fn from_request_parts(parts: &mut Parts, _: &S) -> Result<Self, Self::Rejection> {
        let auth_header = parts.headers.get("X-Auth-Token")
            .and_then(|header| header.to_str().ok())
            .ok_or((StatusCode::UNAUTHORIZED, "Unauthorized"))?;

        verify_auth_token(auth_header).await
            .map_err(|_| (StatusCode::UNAUTHORIZED, "Unauthorized"))
        }
    }
}

现在我们的 API 看起来像这样:

// Save a URL and return an ID associated with it
async fn save_url(
    user: User,
    State(pool): State<PgPool>,
    Json(create_url): Json<CreateUrl>,
) -> Response {
    // Insert our URL into the database and get back an ID that the database generated
    let insert_result = sqlx::query!(
        "INSERT INTO urls (url, user_id) VALUE (lower($1), $2) RETURNING id", 
        create_url.url, user.user_id())        
        .map(|row| row.id)
        .fetch_one(pool)
        .await;
    let id = match insert_result {
        Ok(id) => id,
        Err(_) => return (StatusCode::INTERNAL_SERVER_ERROR, "An unexpected exception has occurred").into_response(),
    };

    (StatusCode::CREATED, Json(CreateUrlResponse { id })).into_response()
}

为什么这种改变是有益的?
这里最大的优势实际上与我们创建下一个 API 路由时发生的事情有关。我们不需要复制和粘贴那么大的代码块,只需添加User函数参数就可以了。

另一个好处是,当我考虑函数时save_url,我并不真正关心我们如何提取和验证令牌的细节。如果我想了解这一点,我可以通过阅读用户的来源来选择加入。我更关心我们如何保存 url 的细节,这里只剩下这些了。

FromRequest并且FromRequestParts都被称为提取器。由于FromRequest消耗了请求的主体,因此每个路由只能有一个FromRequest提取器——而且它必须是最后一个参数。

所以虽然下面这个函数签名:

async fn save_url(
    user: User,
    State(pool): State<PgPool>,
    Json(create_url): Json<CreateUrl>,
) -> Response

与这个函数签名:

async fn save_url(
    Json(create_url): Json<CreateUrl>,
    user: User,
    State(pool): State<PgPool>,
) -> Response

可能看起来一样,但是只有第一个会编译。感谢j_platte指出。如果您的路由遇到编译错误,axum-macros可以提供更详细的错误信息。

响应
Rust 的?运算符是我最喜欢的语言部分之一——它允许您简洁地传播错误或选项。但是,我们的save_url函数根本不使用它,而是选择匹配结果。

让我们解决这个问题:


async fn save_url(
    user: User,
    State(pool): State<PgPool>,
    Json(create_url): Json<CreateUrl>,
) -> Response {
    // Insert our URL into the database and get back an ID that the database generated
    let id = sqlx::query!(
        "INSERT INTO urls (url, user_id) VALUE (lower($1), $2) RETURNING id", 
        create_url.url, user.user_id())        
        .map(|row| row.id)
        .fetch_one(pool)
        .await?; // <-- 

    (StatusCode::CREATED, Json(CreateUrlResponse { id })).into_response()
}

这还不能编译,Rust 告诉我们原因:

the `?` operator can only be used in an async function that returns `Result` or `Option`

我们可以将返回类型更改为Result以表示它可能会失败。我们将定义我们自己的错误类型并实现,From<sqlx::Error>这样我们的 sql 错误将自动转换为我们的错误:

// Our Error
pub enum ApiError {
    DatabaseError(sqlx::Error),
}

// The ? operator will automatically convert sqlx::Error to ApiError
impl From<sqlx::Error> for ApiError {
    fn from(e: sqlx::Error) -> Self {
        SaveUrlError::DatabaseError(e)
    }
}

async fn save_url(
    user: User,
    State(pool): State<PgPool>,
    Json(create_url): Json<CreateUrl>,
) -> Result<Response, ApiError> {
    // Insert our URL into the database and get back an ID that the database generated
    let id = sqlx::query!(
        "INSERT INTO urls (url, user_id) VALUE (lower($1), $2) RETURNING id", 
        create_url.url, user.user_id)        
        .map(|row| row.id)
        .fetch_one(pool)
        .await?;

    Ok((StatusCode::CREATED, Json(CreateUrlResponse { id })).into_response())
}

我们很接近。唯一剩下的部分是 Axum 不知道如何处理 ApiError。我们可以通过实施来解决这个问题IntoResponse

impl IntoResponse for ApiError {
    fn into_response(self) -> Response {
        match self {
            ApiError::DatabaseError(_) => 
                (StatusCode::INTERNAL_SERVER_ERROR, "An unexpected exception has occured").into_response()
    }
}

为什么这种变化是有益的?
对于常见错误,这为我们节省了很多样板文件。否则,每次我们进行数据库调用时,我们都需要匹配结果并返回一个内部服务器错误。

值得一提的是,这种变化并不是普遍有益的。如果您有一个函数可能会以 3-4 种不同的方式失败,您可能希望明确显示您的路由如何将这些不同的错误转换为 HTTP 响应。

急需的重构
如果您一直在大喊“不要将 SQL 调用直接放在 API 路由中”,那么我有一些好消息要告诉您。

pub struct Db;

impl Db {
    pub async fn save_url(
        url: &str, 
        user: &User, 
        pool: &PgPool
    ) -> Result<String, sqlx::Error> {
        sqlx::query!(
            "INSERT INTO urls (url, user_id) VALUE (lower($1), $2) RETURNING id", 
            url, user.user_id())        
        .map(|row| row.id )
        .fetch_one(pool)
        .await
    }
}

(请注意,此函数最初返回一个 CreateUrlResponse - 但正如Kulinda 指出的那样,这是一种特定于路由的类型,应该由 API 负责构建。感谢他们的修复)

通过抽象所有这些,我们的路由变得更加简洁:

async fn save_url(
    user: User,
    State(pool): State<PgPool>,
    Json(create_url): Json<CreateUrl>,
) -> Result<Response, SaveUrlError> {

    let id = Db::save_url(&create_url.url, &user, &pool).await?;
    Ok((StatusCode::CREATED, Json(CreateUrlResponse { id })).into_response())
}

为什么这种变化是有益的?
我喜欢称之为“一目了然”的问题。如果你正在浏览最后的代码片段,而你只是在大声自言自语,你可能会说:

“好的,所以 save_url 路由接受一个 User 和 CreateUrl JSON,然后将其与用户一起保存到数据库并返回响应”

如果您查看路由中包含完整 SQL 查询的示例,您可能会说:

“好的,所以 save_url 路由接受一个 User 和 CreateUrl JSON,然后它执行一些 SQL……这个 sql 做了什么……哦,好的,它只保存 URL 和用户 ID 并返回 ID”

在此示例中,您可能能够足够快地解析 SQL 查询,但想象一下如果您有一个连接或任何远程复杂的东西。通过用简短描述替换 SQL 查询,save_url您可以立即了解它的作用,如果您想了解详细信息,则可以选择加入。

PropelAuth整个后端都是用 Rust 编写的。如果您正在考虑使用 Rust 编写后端并且需要内置多租户/B2B 支持的身份验证提供程序,您可以在此处了解更多关于我们的信息。


概括
最后,我们采用了相当冗长的代码片段并将其分解为关键组件。这不仅使代码一目了然地更容易阅读,而且在我们创建更多 API 路由时也会有所帮助。