picsorter/database.py

import sqlite3
from datetime import datetime


class Database:
    def __init__(self):
        self.db_name = 'images.db'
        self.__create_tables()

    def __create_tables(self):
        conn = sqlite3.connect(self.db_name)
        c = conn.cursor()
        c.executescript("""
            CREATE TABLE IF NOT EXISTS danbooru (
                id INTEGER PRIMARY KEY NOT NULL UNIQUE,
                tags TEXT NOT NULL,
                created_at TIMESTAMP
            );
            """)
        conn.commit()
        conn.close()

    def is_exists(self, _id) -> bool:
        conn = sqlite3.connect(self.db_name)
        c = conn.cursor()
        c.execute("SELECT EXISTS(SELECT 1 FROM danbooru WHERE id=?)", (_id, ))
        result = c.fetchone()[0]
        conn.close()
        return bool(result)
    
    def add(self, _id, tags):
        conn = sqlite3.connect(self.db_name)
        c = conn.cursor()
        sql = 'INSERT INTO danbooru(id, tags, created_at) VALUES (?,?,?)'
        c.execute(sql, (_id, tags, datetime.now()))
        conn.commit()
        conn.close()