From 484146aba99e43c95d9a140de56e2f8219ff7046 Mon Sep 17 00:00:00 2001 From: Milky0217 Date: Wed, 25 Mar 2026 13:17:36 +0800 Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E4=BB=98=E7=B3=BB=E7=BB=9F=E6=A0=B8?= =?UTF-8?q?=E5=BF=83=EF=BC=9A=E9=85=8D=E9=A2=9D=E9=99=90=E5=88=B6=E4=B8=8E?= =?UTF-8?q?=E7=94=A8=E6=88=B7=E4=BB=98=E8=B4=B9=E7=8A=B6=E6=80=81=E7=AE=A1?= =?UTF-8?q?=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 添加 tokio 依赖用于异步测试 - 启用 User 结构体，添加 is_paid/is_admin/paid_expires_at 字段 - 添加 UpdatePaymentRequest 请求体 - insert_weather_data 集成配额检查逻辑 - 新增 get_user_by_id、count_user_weather_data、update_user_payment_status - 添加 payment_fields 数据库迁移脚本 --- Cargo.lock | 13 ++++ Cargo.toml | 4 ++ migrations/001_add_payment_fields.sql | 3 + src/db.rs | 99 ++++++++++++++++++++++----- src/models.rs | 17 ++++- 5 files changed, 117 insertions(+), 19 deletions(-) create mode 100644 migrations/001_add_payment_fields.sql diff --git a/Cargo.lock b/Cargo.lock index 1f52136..db8ade5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2079,6 +2079,7 @@ dependencies = [ "serde", "serde_json", "sqlx", + "tokio", ] [[package]] @@ -2718,9 +2719,21 @@ dependencies = [ "signal-hook-registry", "slab", "socket2 0.6.0", + "tokio-macros", "windows-sys 0.59.0", ] +[[package]] +name = "tokio-macros" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e06d43f1345a3bcd39f6a56dbb7dcab2ba47e68e8ac134855e7e2bdbaf8cab8" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tokio-native-tls" version = "0.3.1" diff --git a/Cargo.toml b/Cargo.toml index 84115f2..e21f4bd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,3 +18,7 @@ reqwest = { version = "0.12.23", features=["json"]} serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.143" sqlx = { version = "0.8.6", features = ["runtime-tokio-rustls", "postgres", "chrono", "uuid"] } +tokio = { version = "1", features = ["full"] } + +[dev-dependencies] +tokio = { version = "1", features = ["full"] } diff --git a/migrations/001_add_payment_fields.sql b/migrations/001_add_payment_fields.sql new file mode 100644 index 0000000..928b0b7 --- /dev/null +++ b/migrations/001_add_payment_fields.sql @@ -0,0 +1,3 @@ +ALTER TABLE users ADD COLUMN IF NOT EXISTS is_paid BOOLEAN NOT NULL DEFAULT false; +ALTER TABLE users ADD COLUMN IF NOT EXISTS is_admin BOOLEAN NOT NULL DEFAULT false; +ALTER TABLE users ADD COLUMN IF NOT EXISTS paid_expires_at TIMESTAMPTZ DEFAULT NULL; diff --git a/src/db.rs b/src/db.rs index 9cde540..ccf663a 100644 --- a/src/db.rs +++ b/src/db.rs @@ -1,28 +1,30 @@ use sqlx::{PgPool, postgres::PgPoolOptions}; use std::env; use std::error::Error; +use chrono::Utc; // 从 models 模块引入 WeatherData 结构体 -use crate::models::{WeatherData, WeatherDataBrief, WeatherListResponse}; +use crate::models::{User, WeatherData, WeatherDataBrief, WeatherListResponse}; // 用于插入weather_data的数据 -pub async fn insert_weather_data(pool: &PgPool, weather_data: &WeatherData) -> Result { - // 1. 根据 openid 查询 user_id (这部分逻辑不变) - let user_id = match sqlx::query_as::<_, (i32,)>("SELECT id FROM users WHERE openid = $1") - .bind(&weather_data.openid) - .fetch_optional(pool) - .await - { - Ok(Some((id,))) => id, - Ok(None) => { - return Err(format!("未找到openid为 {} 的用户", weather_data.openid)); +pub async fn insert_weather_data(pool: &PgPool, weather_data: &WeatherData, user_id: i32) -> Result { + // 配额检查：非付费用户数据条数限制 + let user = get_user_by_id(pool, user_id).await?; + let is_paid_active = user.is_paid && (user.paid_expires_at.is_none() || user.paid_expires_at.unwrap() > Utc::now()); + + if !is_paid_active { + let current_count = count_user_weather_data(pool, user_id).await?; + let limit: i64 = env::var("FREE_USER_DATA_LIMIT") + .unwrap_or_else(|_| "20".to_string()) + .parse() + .unwrap_or(20); + + if current_count >= limit { + return Err("数据条数已达上限，请升级为付费用户".to_string()); } - Err(e) => { - return Err(format!("查询用户失败: {}", e)); - } - }; + } - // 2. 准备插入数据的 SQL 语句 (SQL本身不变) + // 准备插入数据的 SQL 语句 let insert_query = r#" INSERT INTO weather_data ( user_id, title, date, hour, min, longitude, latitude, daysincejanfirst, @@ -242,3 +244,68 @@ pub async fn delete_weather_data( Err(e) => Err(format!("删除天气数据失败: {}", e)), } } + +// 根据用户ID获取用户信息 +pub async fn get_user_by_id(pool: &PgPool, user_id: i32) -> Result { + let query = r#" + SELECT + id, name, openid, phone, type, desc, is_paid, is_admin, paid_expires_at + FROM users + WHERE id = $1 + "#; + + let row = match sqlx::query_as::<_, User>(query) + .bind(user_id) + .fetch_optional(pool) + .await + { + Ok(Some(row)) => row, + Ok(None) => return Err(format!("未找到ID为 {} 的用户", user_id)), + Err(e) => return Err(format!("查询用户信息失败: {}", e)), + }; + + Ok(row) +} + +// 统计用户的天气数据条数 +pub async fn count_user_weather_data(pool: &PgPool, user_id: i32) -> Result { + let query = r#" + SELECT COUNT(*) FROM weather_data WHERE user_id = $1 + "#; + + let count = match sqlx::query_as::<_, (i64,)>(query) + .bind(user_id) + .fetch_one(pool) + .await + { + Ok((count,)) => count, + Err(e) => return Err(format!("查询天气数据条数失败: {}", e)), + }; + + Ok(count) +} + +// 更新用户付费状态 +pub async fn update_user_payment_status( + pool: &PgPool, + user_id: i32, + is_paid: bool, + paid_expires_at: Option>, +) -> Result<(), String> { + let query = r#" + UPDATE users + SET is_paid = $1, paid_expires_at = $2 + WHERE id = $3 + "#; + + match sqlx::query(query) + .bind(is_paid) + .bind(paid_expires_at) + .bind(user_id) + .execute(pool) + .await + { + Ok(_) => Ok(()), + Err(e) => Err(format!("更新用户付费状态失败: {}", e)), + } +} diff --git a/src/models.rs b/src/models.rs index e4db6d5..be16b37 100644 --- a/src/models.rs +++ b/src/models.rs @@ -272,8 +272,7 @@ pub struct WeatherData { pub wind_speed_suitability: String, } -/* -#[derive(Debug, Serialize, FromRow)] +#[derive(Debug, Serialize, Deserialize, FromRow)] pub struct User { pub id: i32, #[serde(rename = "name")] @@ -289,5 +288,17 @@ pub struct User { #[serde(rename = "desc")] #[sqlx(rename = "desc")] pub description: Option, + #[sqlx(rename = "is_paid")] + pub is_paid: bool, + #[sqlx(rename = "is_admin")] + pub is_admin: bool, + #[sqlx(rename = "paid_expires_at")] + pub paid_expires_at: Option>, +} + +// 管理员更新用户付费状态的请求体 +#[derive(Debug, Deserialize)] +pub struct UpdatePaymentRequest { + pub is_paid: bool, + pub paid_expires_at: Option, // ISO 8601 格式 } -*/