feat: 支持记录访问日志

ArronYR · Mar 8, 2024 · 0d4f042 · 0d4f042
1 parent d66b3c5
commit 0d4f042
Show file tree

Hide file tree

Showing 10 changed files with 186 additions and 13 deletions.
diff --git a/src/bin/short_url.rs b/src/bin/short_url.rs
@@ -1,4 +1,5 @@
 use actix_cors::Cors;
+use actix_web::rt::spawn;
 use actix_web::{
     get, http, middleware, post, web, App, Error, HttpRequest, HttpResponse, HttpServer, Responder,
 };
@@ -9,10 +10,15 @@ use num_traits::ToPrimitive;
 use sea_orm::{ConnectionTrait, Database, DatabaseConnection, DbConn};
 use serde_json::json;
 use short_url::middleware::validator::ApiValidateMiddleware;
-use short_url::models::link::{
-    ChangeExpiredReq, ChangeStatusReq, GenerateReq, LinkStatusEnum, Model as LinkModel,
-    SearchParams,
+use short_url::models::link::SearchRecordItem;
+use short_url::models::{
+    access_log::Model as AccessLogModel,
+    link::{
+        ChangeExpiredReq, ChangeStatusReq, GenerateReq, LinkStatusEnum, Model as LinkModel,
+        SearchParams,
+    },
 };
+use short_url::service::access_log::AccessLogService;
 use short_url::service::link::LinkService;
 use short_url::utils::helpers::{
     generate_short_id, is_reasonable_timestamp, is_valid_url, md5_hex, validate_header_token,
@@ -51,6 +57,10 @@ fn init_config() -> Config {
         Err(_) => 60,
     };
     let api_secret = env::var("API_SECRET").unwrap_or_else(|_| API_SECRET.to_string());
+    let access_log = match env::var("ACCESS_LOG") {
+        Ok(value) => value.parse::<bool>().unwrap_or(true),
+        Err(_) => true,
+    };
 
     Config {
         port,
@@ -65,6 +75,7 @@ fn init_config() -> Config {
         cache_max_cap,
         cache_live_time,
         api_secret,
+        access_log,
     }
 }
 
@@ -100,7 +111,32 @@ async fn redirect(
 ) -> impl Responder {
     let db_conn = &state.db_conn;
     let cache = &state.cache;
+    let config = &state.config;
+
     let short_id = short_id.as_str();
+    // 写入日志
+    if config.access_log {
+        let short_id = short_id.to_string().clone();
+        let db_conn = db_conn.clone();
+        let headers = request.headers().clone();
+
+        spawn(async move {
+            let req_headers = headers
+                .iter()
+                .map(|(k, v)| format!("{}: {:?}", k, v))
+                .collect::<Vec<String>>()
+                .join("\n");
+            let model = AccessLogModel {
+                id: 0,
+                short_id: short_id.clone(),
+                req_headers,
+                create_time: NaiveDateTime::default(),
+            };
+            if let Err(e) = AccessLogService::add(&db_conn, model).await {
+                error!("add id: {} access log error: {:?}", short_id, e);
+            };
+        });
+    }
 
     // 如果缓存存在
     if let Some(cached) = cache.get(short_id).await {
@@ -163,6 +199,7 @@ async fn redirect(
 #[get("/api/search")]
 async fn search(req: HttpRequest, state: web::Data<AppState>) -> Result<HttpResponse, Error> {
     let db_conn = &state.db_conn;
+    let config = &state.config;
 
     // get params
     let params = web::Query::<SearchParams>::from_query(req.query_string()).unwrap();
@@ -178,8 +215,29 @@ async fn search(req: HttpRequest, state: web::Data<AppState>) -> Result<HttpResp
     let (links, pages) = LinkService::search(&db_conn, params)
         .await
         .expect("Cannot find links in page");
+
+    // 查找PV数据（如果开启了ACCESS_LOG）
+    let mut pv_map: HashMap<String, i64> = HashMap::new();
+    if config.access_log {
+        let ids: Vec<String> = links.iter().map(|r| r.short_id.clone()).collect();
+        pv_map = AccessLogService::batch_query_pv(&db_conn, ids).await;
+    }
+
+    let mut records: Vec<SearchRecordItem> = Vec::new();
+    for link in links {
+        records.push(SearchRecordItem {
+            id: link.id,
+            short_id: link.short_id.clone(),
+            original_url: link.original_url,
+            expired_ts: link.expired_ts,
+            status: link.status,
+            create_time: link.create_time,
+            pv: *pv_map.get(link.short_id.as_str()).unwrap_or(&0i64),
+        })
+    }
+
     let json = json!({
-        "records": links,
+        "records": records,
         "pages": pages,
         "size": size
     });

diff --git a/src/db.sql b/src/db.sql
@@ -10,4 +10,15 @@ CREATE TABLE IF NOT EXISTS `link`
     UNIQUE KEY `uniq_short_url` (`short_id`)
 ) ENGINE = InnoDB
   DEFAULT CHARSET = utf8mb4
-  COLLATE = utf8mb4_bin COMMENT ='链接记录';
+  COLLATE = utf8mb4_bin COMMENT ='链接记录';
+
+CREATE TABLE IF NOT EXISTS `access_log`
+(
+    `id`          bigint(20) unsigned NOT NULL AUTO_INCREMENT COMMENT '主键',
+    `short_id`    varchar(50)         NOT NULL COMMENT '短链ID',
+    `req_headers` longtext COMMENT '请求头',
+    `create_time` datetime            NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
+    PRIMARY KEY (`id`),
+    KEY `idx_short_id` (`short_id`) USING BTREE
+) ENGINE = InnoDB
+  DEFAULT CHARSET = utf8mb4 COMMENT ='访问日志';
diff --git a/src/lib.rs b/src/lib.rs
@@ -21,6 +21,7 @@ pub struct Config {
     pub cache_max_cap: u64,
     pub cache_live_time: u64,
     pub api_secret: String,
+    pub access_log: bool,
 }
 
 #[derive(Debug, Clone)]

diff --git a/src/models/access_log.rs b/src/models/access_log.rs
@@ -0,0 +1,19 @@
+use chrono::NaiveDateTime;
+use sea_orm::entity::prelude::*;
+use sea_orm::{ActiveModelBehavior, DeriveEntityModel, DeriveRelation, EnumIter};
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel, Deserialize, Serialize)]
+#[sea_orm(table_name = "access_log")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: u64,
+    pub short_id: String,
+    pub req_headers: String,
+    pub create_time: NaiveDateTime,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}
diff --git a/src/models/link.rs b/src/models/link.rs
@@ -54,3 +54,14 @@ pub enum LinkStatusEnum {
     Normal = 0,
     Disabled = 1,
 }
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct SearchRecordItem {
+    pub id: u64,
+    pub short_id: String,
+    pub original_url: String,
+    pub expired_ts: i64,
+    pub status: i16,
+    pub create_time: NaiveDateTime,
+    pub pv: i64,
+}
diff --git a/src/models/mod.rs b/src/models/mod.rs
@@ -1 +1,2 @@
+pub mod access_log;
 pub mod link;
diff --git a/src/service/access_log.rs b/src/service/access_log.rs
@@ -0,0 +1,58 @@
+use crate::models::access_log;
+use log::error;
+use sea_orm::{ActiveModelTrait, DatabaseBackend, DbConn, DbErr, FromQueryResult, Set, Statement};
+use std::collections::HashMap;
+
+pub struct AccessLogService;
+
+impl AccessLogService {
+    // 添加
+    pub async fn add(
+        db: &DbConn,
+        data: access_log::Model,
+    ) -> Result<access_log::ActiveModel, DbErr> {
+        access_log::ActiveModel {
+            short_id: Set(data.short_id.to_owned()),
+            req_headers: Set(data.req_headers.to_owned()),
+            ..Default::default()
+        }
+        .save(db)
+        .await
+    }
+
+    pub async fn batch_query_pv(db: &DbConn, short_ids: Vec<String>) -> HashMap<String, i64> {
+        #[derive(Debug, FromQueryResult)]
+        struct GroupByResult {
+            short_id: String,
+            total: i64,
+        }
+
+        let mut pv_map: HashMap<String, i64> = HashMap::new();
+        if short_ids.is_empty() {
+            return pv_map;
+        }
+        let id_string = short_ids
+            .iter()
+            .map(|s| format!("'{}'", s))
+            .collect::<Vec<_>>()
+            .join(",");
+
+        let result = GroupByResult::find_by_statement(Statement::from_sql_and_values(
+            DatabaseBackend::MySql,
+            format!("SELECT short_id, COUNT(1) AS total FROM `access_log` WHERE `short_id` IN ({}) GROUP BY `short_id`", id_string),
+            [],
+        ))
+        .all(db)
+        .await;
+
+        if result.is_err() {
+            error!("batch_query_pv error: {:?}", result.err());
+            return pv_map;
+        }
+
+        for row in result.unwrap() {
+            pv_map.insert(row.short_id, row.total);
+        }
+        pv_map
+    }
+}
diff --git a/src/service/mod.rs b/src/service/mod.rs
@@ -1 +1,2 @@
+pub mod access_log;
 pub mod link;
diff --git a/web/src/api/typings.d.ts b/web/src/api/typings.d.ts
@@ -10,6 +10,7 @@ declare namespace API {
         original_url?: string;
         status?: number;
         expired_ts?: number;
+        pv?: number;
     }
 }
 

diff --git a/web/src/components/LinkTable.tsx b/web/src/components/LinkTable.tsx
@@ -60,7 +60,7 @@ export default function LinkTable() {
         {
             field: 'short_id',
             headerName: '短链接',
-            width: 240,
+            width: 180,
             sortable: false,
             renderCell: (props) => {
                 const url = `${baseUrl}/${props.row.short_id}`
@@ -71,18 +71,16 @@ export default function LinkTable() {
                         align={'center'}
                         display={'flex'}
                         justifyContent={"space-between"}
-                        width={'90%'}
                     >
-                        {url}
-                        <Link color={'inherit'} href={url} target={'_blank'}>访问</Link>
+                        <Link color={'inherit'} underline={'none'} href={url} target={'_blank'}>{url}</Link>
                     </Typography>
                 )
             }
         },
         {
             field: 'original_url',
             headerName: '原链接',
-            minWidth: 400,
+            minWidth: 360,
             cellClassName: 'cell-cls-name',
             sortable: false,
             renderCell: (props) => {
@@ -94,16 +92,22 @@ export default function LinkTable() {
                         align={'center'}
                         display={"flex"}
                         justifyContent={"space-between"}
-                        width={'90%'}
+                        width={'100%'}
                     >
                         <Tooltip title={props.row.original_url} arrow={true} placement={"top"}>
                             <Typography
                                 variant={'body2'}
                                 noWrap={true}
                                 textOverflow={"ellipsis"}
-                            >{props.row.original_url}</Typography>
+                            >
+                                <Link
+                                    color={"inherit"}
+                                    href={props.row.original_url}
+                                    target={"_blank"}
+                                    underline={'none'}
+                                >{props.row.original_url}</Link>
+                            </Typography>
                         </Tooltip>
-                        <Link sx={{ml: 4}} color={"inherit"} href={props.row.original_url} target={"_blank"}>访问</Link>
                     </Typography>
                 )
             }
@@ -132,6 +136,14 @@ export default function LinkTable() {
             sortable: false,
             valueGetter: ({value}) => value ? moment(value).format(DT_FORMAT.DATETIME) : '永久',
         },
+        {
+            field: 'pv',
+            headerName: 'PV',
+            minWidth: 40,
+            cellClassName: 'cell-cls-name',
+            sortable: false,
+            valueGetter: ({value}) => value ?? 0,
+        },
         {
             field: 'actions',
             type: 'actions',