first commit

This commit is contained in:
Masato Imai
2025-01-07 09:34:55 +00:00
commit 56aa5fca6a
10 changed files with 3765 additions and 0 deletions

1
.gitignore vendored Normal file
View File

@@ -0,0 +1 @@
/target

3229
Cargo.lock generated Normal file

File diff suppressed because it is too large Load Diff

14
Cargo.toml Normal file
View File

@@ -0,0 +1,14 @@
[package]
name = "cigarette-counter"
version = "0.1.0"
edition = "2021"
[dependencies]
poise = "0.6.1"
tokio = { version = "1.42.0", features = ["full"] }
sqlx = { version = "0.8", features = [ "runtime-tokio", "tls-native-tls", "postgres", "chrono" ] }
serde = { version = "1.0", features = ["derive"] }
chrono = { version = "0.4", features = ["serde"] }
native-tls = "0.2"
postgres-native-tls = "0.5"
anyhow = "1.0.95"

View File

@@ -0,0 +1,13 @@
DROP VIEW IF EXISTS daily_smoking_summary;
DROP TRIGGER IF EXISTS update_smoking_logs_updated_at ON smoking_logs;
DROP TRIGGER IF EXISTS update_users_updated_at ON users;
DROP FUNCTION IF EXISTS update_updated_at_column();
DROP INDEX IF EXISTS idx_smoking_logs_smoked_at;
DROP INDEX IF EXISTS idx_smoking_logs_discord_id;
DROP TABLE IF EXISTS smoking_logs;
DROP TABLE IF EXISTS smoking_types;
DROP TABLE IF EXISTS users;

View File

@@ -0,0 +1,66 @@
CREATE TABLE users (
discord_id VARCHAR(20) PRIMARY KEY,
username VARCHAR(100) NOT NULL,
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE smoking_types (
id SERIAL PRIMARY KEY,
type_name VARCHAR(50) NOT NULL,
description TEXT,
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);
INSERT INTO smoking_types (type_name, description)
VALUES
('traditional', '紙タバコ'),
('iqos', 'IQOS');
CREATE TABLE smoking_logs (
id SERIAL PRIMARY KEY,
discord_id VARCHAR(20) REFERENCES users(discord_id),
smoking_type_id INTEGER REFERENCES smoking_types(id),
quantity INTEGER NOT NULL CHECK (quantity > 0),
smoked_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX idx_smoking_logs_discord_id ON smoking_logs(discord_id);
CREATE INDEX idx_smoking_logs_smoked_at ON smoking_logs(smoked_at);
CREATE OR REPLACE FUNCTION update_updated_at_column()
RETURNS TRIGGER AS $$
BEGIN
NEW.updated_at = CURRENT_TIMESTAMP;
RETURN NEW;
END;
$$ language 'plpgsql';
CREATE TRIGGER update_users_updated_at
BEFORE UPDATE ON users
FOR EACH ROW
EXECUTE FUNCTION update_updated_at_column();
CREATE TRIGGER update_smoking_logs_updated_at
BEFORE UPDATE ON smoking_logs
FOR EACH ROW
EXECUTE FUNCTION update_updated_at_column();
CREATE VIEW daily_smoking_summary AS
SELECT
sl.discord_id,
u.username,
DATE(sl.smoked_at) as smoke_date,
st.type_name,
SUM(sl.quantity) as total_quantity
FROM smoking_logs sl
JOIN users u ON sl.discord_id = u.discord_id
JOIN smoking_types st ON sl.smoking_type_id = st.id
GROUP BY
sl.discord_id,
u.username,
DATE(sl.smoked_at),
st.type_name;

View File

@@ -0,0 +1 @@
DELETE FROM smoking_types WHERE type_name IN ('ploom', 'glo', 'other');

View File

@@ -0,0 +1,5 @@
INSERT INTO smoking_types (type_name, description)
VALUES
('ploom', 'プルーム'),
('glo', 'グロー'),
('other', 'その他');

99
src/commands.rs Normal file
View File

@@ -0,0 +1,99 @@
use crate::database::DailySmokingSummary;
use crate::{Context, Error};
use chrono::Local;
use poise::serenity_prelude::{self as serenity, CreateInteractionResponseMessage};
use poise::CreateReply;
async fn create_cigarette_buttons(
ctx: &Context<'_>,
uuid: &str,
) -> Result<Vec<serenity::CreateButton>, Error> {
let db = ctx.data().database.lock().await;
let cigarette_types = db.get_smoking_types().await?;
Ok(cigarette_types
.into_iter()
.map(|cigarette_type| {
serenity::CreateButton::new(format!("{}{}", uuid, cigarette_type.id))
.style(serenity::ButtonStyle::Primary)
.label(cigarette_type.description.unwrap_or_default())
})
.collect())
}
fn format_daily_summary(daily_summary: Vec<DailySmokingSummary>) -> String {
daily_summary
.into_iter()
.map(|summary| {
format!(
"\n{}: {}",
summary.description,
summary.total_quantity.unwrap_or_default()
)
})
.collect()
}
async fn handle_interaction(
ctx: &Context<'_>,
mci: &serenity::ComponentInteraction,
uuid: &str,
) -> Result<(), Error> {
let db = ctx.data().database.lock().await;
let user_id = mci.user.id.get().to_string();
let user = db.get_or_create_user(&user_id, &ctx.author().name).await?;
let cigarette_id = extract_cigarette_id(&mci.data.custom_id, uuid)?;
db.log_smoking(&user.discord_id, cigarette_id, 1).await?;
let daily_summary = db
.get_daily_summary(&user.discord_id, Local::now().date_naive())
.await?;
let reply_content = format!(
"記録しました。\n本日の累計本数{}",
format_daily_summary(daily_summary)
);
mci.create_response(
ctx,
serenity::CreateInteractionResponse::Message(
CreateInteractionResponseMessage::new().content(reply_content),
),
)
.await?;
Ok(())
}
fn extract_cigarette_id(custom_id: &str, uuid: &str) -> Result<i32, Error> {
i32::from_str_radix(custom_id.trim_start_matches(uuid), 10)
.map_err(|e| Error::from(format!("Failed to parse cigarette ID: {}", e)))
}
#[poise::command(prefix_command)]
pub async fn create_cigarette_ui(ctx: Context<'_>) -> Result<(), Error> {
let uuid = ctx.id().to_string();
let buttons = create_cigarette_buttons(&ctx, &uuid).await?;
let components = vec![serenity::CreateActionRow::Buttons(buttons)];
let reply = CreateReply::default()
.content("喫煙カウント")
.components(components);
ctx.send(reply).await?;
while let Some(mci) = serenity::ComponentInteractionCollector::new(ctx)
.channel_id(ctx.channel_id())
.filter({
let uuid = uuid.clone();
move |mci| mci.data.custom_id.starts_with(&uuid)
})
.await
{
handle_interaction(&ctx, &mci, &uuid).await?;
}
Ok(())
}

281
src/database.rs Normal file
View File

@@ -0,0 +1,281 @@
use chrono::{DateTime, NaiveDate, Utc};
use serde::{Deserialize, Serialize};
use sqlx::{postgres::PgPool, Error};
use std::sync::Arc;
#[derive(Debug, Serialize, Deserialize)]
pub struct User {
pub discord_id: String,
pub username: String,
pub created_at: Option<DateTime<Utc>>,
pub updated_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct SmokingType {
pub id: i32,
pub type_name: String,
pub description: Option<String>,
pub created_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct SmokingLog {
pub id: i32,
pub discord_id: String,
pub smoking_type_id: i32,
pub quantity: i32,
pub smoked_at: DateTime<Utc>,
pub created_at: Option<DateTime<Utc>>,
pub updated_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct DailySmokingSummary {
pub discord_id: String,
pub username: String,
pub smoke_date: NaiveDate,
pub type_name: String,
pub description: String,
pub total_quantity: Option<i64>,
}
pub struct Database {
pool: Arc<PgPool>,
}
impl Database {
pub fn new(pool: PgPool) -> Self {
Self {
pool: Arc::new(pool),
}
}
pub async fn create_user(&self, discord_id: &str, username: &str) -> Result<User, Error> {
let user = sqlx::query_as!(
User,
r#"
INSERT INTO users (discord_id, username)
VALUES ($1, $2)
RETURNING
discord_id as "discord_id!",
username as "username!",
created_at,
updated_at
"#,
discord_id,
username
)
.fetch_one(&*self.pool)
.await?;
Ok(user)
}
pub async fn get_or_create_user(
&self,
discord_id: &str,
username: &str,
) -> Result<User, Error> {
let mut tx = self.pool.begin().await?;
let user = sqlx::query_as!(
User,
r#"
SELECT
discord_id,
username,
created_at as "created_at!",
updated_at as "updated_at!"
FROM users
WHERE discord_id = $1
"#,
discord_id
)
.fetch_optional(&mut *tx)
.await?;
let user = match user {
Some(user) => {
if user.username != username {
sqlx::query_as!(
User,
r#"
UPDATE users
SET username = $2, updated_at = CURRENT_TIMESTAMP
WHERE discord_id = $1
RETURNING
discord_id,
username,
created_at as "created_at!",
updated_at as "updated_at!"
"#,
discord_id,
username
)
.fetch_one(&mut *tx)
.await?
} else {
user
}
}
None => {
sqlx::query_as!(
User,
r#"
INSERT INTO users (discord_id, username)
VALUES ($1, $2)
RETURNING
discord_id,
username,
created_at as "created_at!",
updated_at as "updated_at!"
"#,
discord_id,
username
)
.fetch_one(&mut *tx)
.await?
}
};
tx.commit().await?;
Ok(user)
}
pub async fn user_exists(&self, discord_id: &str) -> Result<bool, Error> {
let exists = sqlx::query_scalar!(
r#"
SELECT EXISTS(SELECT 1 FROM users WHERE discord_id = $1) as "exists!"
"#,
discord_id
)
.fetch_one(&*self.pool)
.await?;
Ok(exists)
}
pub async fn log_smoking(
&self,
discord_id: &str,
smoking_type_id: i32,
quantity: i32,
) -> Result<SmokingLog, Error> {
let log = sqlx::query_as!(
SmokingLog,
r#"
INSERT INTO smoking_logs (discord_id, smoking_type_id, quantity)
VALUES ($1, $2, $3)
RETURNING
id as "id!",
discord_id as "discord_id!",
smoking_type_id as "smoking_type_id!",
quantity as "quantity!",
smoked_at as "smoked_at!",
created_at,
updated_at
"#,
discord_id,
smoking_type_id,
quantity
)
.fetch_one(&*self.pool)
.await?;
Ok(log)
}
pub async fn get_daily_summary(
&self,
discord_id: &str,
date: NaiveDate,
) -> Result<Vec<DailySmokingSummary>, Error> {
let summary = sqlx::query_as!(
DailySmokingSummary,
r#"
SELECT
sl.discord_id as "discord_id!",
u.username as "username!",
DATE(sl.smoked_at) as "smoke_date!",
st.type_name as "type_name!",
st.description as "description!",
SUM(sl.quantity) as total_quantity
FROM smoking_logs sl
JOIN users u ON sl.discord_id = u.discord_id
JOIN smoking_types st ON sl.smoking_type_id = st.id
WHERE sl.discord_id = $1
AND DATE(sl.smoked_at) = $2
GROUP BY
sl.discord_id,
u.username,
DATE(sl.smoked_at),
st.type_name,
st.description
"#,
discord_id,
date
)
.fetch_all(&*self.pool)
.await?;
Ok(summary)
}
pub async fn get_smoking_type(&self, id: i32) -> Result<SmokingType, Error> {
let smoking_type = sqlx::query_as!(
SmokingType,
r#"
SELECT
id as "id!",
type_name as "type_name!",
description,
created_at
FROM smoking_types
WHERE id = $1
"#,
id
)
.fetch_one(&*self.pool)
.await?;
Ok(smoking_type)
}
pub async fn get_smoking_types(&self) -> Result<Vec<SmokingType>, Error> {
let types = sqlx::query_as!(
SmokingType,
r#"
SELECT
id as "id!",
type_name as "type_name!",
description,
created_at
FROM smoking_types
ORDER BY id
"#
)
.fetch_all(&*self.pool)
.await?;
Ok(types)
}
pub async fn smoking_type_exists(&self, id: i32) -> Result<bool, Error> {
let exists = sqlx::query_scalar!(
r#"
SELECT EXISTS(SELECT 1 FROM smoking_types WHERE id = $1) as "exists!"
"#,
id
)
.fetch_one(&*self.pool)
.await?;
Ok(exists)
}
}

56
src/main.rs Normal file
View File

@@ -0,0 +1,56 @@
mod commands;
mod database;
use std::sync::Arc;
use anyhow::Result;
use commands::create_cigarette_ui;
use database::Database;
use poise::{
serenity_prelude::{self as serenity, futures::lock::Mutex},
PrefixFrameworkOptions,
};
use sqlx::PgPool;
pub struct Data {
pub database: Arc<Mutex<Database>>,
}
pub type Error = Box<dyn std::error::Error + Send + Sync>;
pub type Context<'a> = poise::Context<'a, Data, Error>;
#[tokio::main]
async fn main() -> Result<()> {
let token = std::env::var("BOT_TOKEN").expect("missing BOT_TOKEN");
let intents = serenity::GatewayIntents::non_privileged() | serenity::GatewayIntents::all();
let pool =
PgPool::connect(&std::env::var("DATABASE_URL").expect("missing DATABASE_URL")).await?;
let db = Database::new(pool);
let framework = poise::Framework::builder()
.options(poise::FrameworkOptions {
commands: vec![create_cigarette_ui()],
prefix_options: PrefixFrameworkOptions {
prefix: Some(String::from("c:")),
..Default::default()
},
..Default::default()
})
.setup(|_ctx, _ready, _framework| {
Box::pin(async move {
Ok(Data {
database: Arc::new(Mutex::new(db)),
})
})
})
.build();
let client = serenity::ClientBuilder::new(token, intents)
.framework(framework)
.await;
client.unwrap().start().await?;
Ok(())
}