first commit

.gitignore

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

Cargo.toml

@@ -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"

migrations/20250107040228_cigarette.down.sql

@@ -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;
+

migrations/20250107040228_cigarette.up.sql

@@ -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;
+

migrations/20250107084522_cigarette.down.sql

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

migrations/20250107084522_cigarette.up.sql

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

src/commands.rs

@@ -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(())
+}

src/database.rs

@@ -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)
+    }
+}

src/main.rs

@@ -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(())
+}