Rewrote Crabstero from Java to Python.
- Replaced Javacord with discord.py. - Replaced Redis backend with SQLite via aiosqlite. - Replaced Gradle build with pyproject.toml and uv. - Added setuptools-scm for automatic versioning from git tags. - Added argparse CLI with systemd credential support for the bot token. - Replaced per-channel ingestion tasks with a bounded queue and worker pool. - Removed Dockerfile and Gitea Actions workflow.
This commit is contained in:
@@ -0,0 +1,327 @@
|
||||
# Copyright 2026 Logan Fick
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
Provides async SQLite database access for all of Crabstero's persistent storage needs.
|
||||
|
||||
Uses aiosqlite for native async access. All methods are async def.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Self
|
||||
|
||||
import aiosqlite
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
AUTO_COMMIT_WRITE_THRESHOLD = 100 # Commit after this many DB write operations.
|
||||
AUTO_COMMIT_TIMEOUT_SECONDS = (
|
||||
60.0 # Commit after this many seconds since first uncommitted write.
|
||||
)
|
||||
|
||||
# SQL statements for creating the database schema.
|
||||
_SCHEMA = """
|
||||
-- Markov chain starting words.
|
||||
-- Each row represents one occurrence. Duplicates represent frequency weight.
|
||||
CREATE TABLE IF NOT EXISTS markov_start_words (
|
||||
channel_id INTEGER NOT NULL,
|
||||
word TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_start_channel ON markov_start_words(channel_id, word);
|
||||
|
||||
-- Markov chain word transitions.
|
||||
-- Each row represents one occurrence. Duplicates represent frequency weight.
|
||||
CREATE TABLE IF NOT EXISTS markov_transitions (
|
||||
channel_id INTEGER NOT NULL,
|
||||
word TEXT NOT NULL,
|
||||
next_word TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_transitions_channel_word ON markov_transitions(channel_id, word);
|
||||
|
||||
-- Image URLs per channel.
|
||||
CREATE TABLE IF NOT EXISTS channel_images (
|
||||
channel_id INTEGER NOT NULL,
|
||||
url TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_images_channel ON channel_images(channel_id);
|
||||
|
||||
-- Flags for channels, servers, and users.
|
||||
CREATE TABLE IF NOT EXISTS flags (
|
||||
entity_type TEXT NOT NULL,
|
||||
entity_id TEXT NOT NULL,
|
||||
flag_name TEXT NOT NULL,
|
||||
PRIMARY KEY (entity_type, entity_id, flag_name)
|
||||
);
|
||||
|
||||
-- Tracks which channels have been bulk-ingested.
|
||||
CREATE TABLE IF NOT EXISTS ingested_channels (
|
||||
channel_id INTEGER NOT NULL PRIMARY KEY
|
||||
);
|
||||
"""
|
||||
|
||||
|
||||
class Database:
|
||||
"""
|
||||
Manages all SQLite database operations for Crabstero.
|
||||
|
||||
Uses aiosqlite for native async access. A single connection is held open for the lifetime of
|
||||
the bot process with WAL mode enabled for concurrent read performance.
|
||||
"""
|
||||
|
||||
def __init__(self, connection: aiosqlite.Connection) -> None:
|
||||
"""
|
||||
Initializes the Database wrapper with an already-opened aiosqlite connection.
|
||||
|
||||
:param connection: An open aiosqlite connection.
|
||||
"""
|
||||
self._connection = connection
|
||||
self._pending_writes = 0
|
||||
self._flush_task: asyncio.Task[None] | None = None
|
||||
|
||||
@classmethod
|
||||
async def connect(cls, path: str) -> Self:
|
||||
"""
|
||||
Opens a new SQLite database at the given path, configures it for performance, and creates
|
||||
the schema if it does not already exist.
|
||||
|
||||
:param path: The file path to the SQLite database.
|
||||
:return: A new Database instance ready for use.
|
||||
"""
|
||||
connection = await aiosqlite.connect(path)
|
||||
|
||||
# Enable WAL mode for better concurrent read performance.
|
||||
await connection.execute("PRAGMA journal_mode=WAL")
|
||||
|
||||
# Set synchronous to NORMAL for a balance between safety and speed.
|
||||
await connection.execute("PRAGMA synchronous=NORMAL")
|
||||
|
||||
# Create the schema tables and indexes if they do not already exist.
|
||||
# executescript auto-commits, so no explicit commit is needed.
|
||||
await connection.executescript(_SCHEMA)
|
||||
|
||||
return cls(connection)
|
||||
|
||||
async def commit(self) -> None:
|
||||
"""Commits pending writes and resets the flush timer."""
|
||||
await self._connection.commit()
|
||||
self._pending_writes = 0
|
||||
if self._flush_task is not None:
|
||||
self._flush_task.cancel()
|
||||
self._flush_task = None
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Cancels the flush timer, commits any pending writes, then closes the database connection."""
|
||||
if self._flush_task is not None:
|
||||
self._flush_task.cancel()
|
||||
try:
|
||||
await self._flush_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._flush_task = None
|
||||
if self._pending_writes > 0:
|
||||
await self._connection.commit()
|
||||
self._pending_writes = 0
|
||||
await self._connection.close()
|
||||
|
||||
async def add_start_words_batch(self, rows: list[tuple[int, str]]) -> None:
|
||||
"""
|
||||
Inserts a batch of starting words into the markov_start_words table.
|
||||
|
||||
:param rows: A list of (channel_id, word) tuples to insert.
|
||||
"""
|
||||
await self._connection.executemany(
|
||||
"INSERT INTO markov_start_words (channel_id, word) VALUES (?, ?)", rows
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def get_random_start_word(self, channel_id: int) -> str | None:
|
||||
"""
|
||||
Returns a random starting word for a given channel, weighted by occurrence frequency.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:return: A random starting word, or None if no starting words exist for this channel.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT word FROM markov_start_words WHERE channel_id = ? ORDER BY RANDOM() LIMIT 1",
|
||||
(channel_id,),
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
async def add_transitions_batch(self, rows: list[tuple[int, str, str]]) -> None:
|
||||
"""
|
||||
Inserts a batch of word transitions into the markov_transitions table.
|
||||
|
||||
:param rows: A list of (channel_id, word, next_word) tuples to insert.
|
||||
"""
|
||||
await self._connection.executemany(
|
||||
"INSERT INTO markov_transitions (channel_id, word, next_word) VALUES (?, ?, ?)",
|
||||
rows,
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def get_random_next_word(self, channel_id: int, word: str) -> str | None:
|
||||
"""
|
||||
Returns a random next word for a given word in a given channel, weighted by occurrence
|
||||
frequency.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:param word: The current word to find a transition for.
|
||||
:return: A random next word, or None if no transitions exist.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT next_word FROM markov_transitions WHERE channel_id = ? AND word = ? ORDER BY RANDOM() LIMIT 1",
|
||||
(channel_id, word),
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
async def get_random_completing_next_word(
|
||||
self, channel_id: int, word: str
|
||||
) -> str | None:
|
||||
"""
|
||||
Returns a random next word that ends a sentence for a given word in a given channel,
|
||||
weighted by occurrence frequency.
|
||||
|
||||
A completing word is one whose last character is '.', '!', '?', or '§'.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:param word: The current word to find a completing transition for.
|
||||
:return: A random completing next word, or None if no completing transitions exist.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT next_word FROM markov_transitions WHERE channel_id = ? AND word = ? AND SUBSTR(next_word, -1, 1) IN ('.', '!', '?', '§') ORDER BY RANDOM() LIMIT 1",
|
||||
(channel_id, word),
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
async def add_image(self, channel_id: int, url: str) -> None:
|
||||
"""
|
||||
Stores an image URL for a given channel.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:param url: The image URL to store.
|
||||
"""
|
||||
await self._connection.execute(
|
||||
"INSERT INTO channel_images (channel_id, url) VALUES (?, ?)",
|
||||
(channel_id, url),
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def get_random_image(self, channel_id: int) -> str | None:
|
||||
"""
|
||||
Returns a random image URL for a given channel.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:return: A random image URL, or None if no images exist for this channel.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT url FROM channel_images WHERE channel_id = ? ORDER BY RANDOM() LIMIT 1",
|
||||
(channel_id,),
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
async def set_flag(self, entity_type: str, entity_id: str, flag_name: str) -> None:
|
||||
"""
|
||||
Sets a flag on a given entity. If the flag is already set, this is a no-op.
|
||||
|
||||
:param entity_type: The type of entity ("channel", "server", or "user").
|
||||
:param entity_id: The Discord ID of the entity.
|
||||
:param flag_name: The name of the flag to set.
|
||||
"""
|
||||
await self._connection.execute(
|
||||
"INSERT OR IGNORE INTO flags (entity_type, entity_id, flag_name) VALUES (?, ?, ?)",
|
||||
(entity_type, entity_id, flag_name),
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def clear_flag(
|
||||
self, entity_type: str, entity_id: str, flag_name: str
|
||||
) -> None:
|
||||
"""
|
||||
Clears a flag on a given entity. If the flag is not set, this is a no-op.
|
||||
|
||||
:param entity_type: The type of entity ("channel", "server", or "user").
|
||||
:param entity_id: The Discord ID of the entity.
|
||||
:param flag_name: The name of the flag to clear.
|
||||
"""
|
||||
await self._connection.execute(
|
||||
"DELETE FROM flags WHERE entity_type = ? AND entity_id = ? AND flag_name = ?",
|
||||
(entity_type, entity_id, flag_name),
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def is_flag_set(
|
||||
self, entity_type: str, entity_id: str, flag_name: str
|
||||
) -> bool:
|
||||
"""
|
||||
Checks whether a flag is set on a given entity.
|
||||
|
||||
:param entity_type: The type of entity ("channel", "server", or "user").
|
||||
:param entity_id: The Discord ID of the entity.
|
||||
:param flag_name: The name of the flag to check.
|
||||
:return: True if the flag is set, False otherwise.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT 1 FROM flags WHERE entity_type = ? AND entity_id = ? AND flag_name = ?",
|
||||
(entity_type, entity_id, flag_name),
|
||||
) as cursor:
|
||||
return await cursor.fetchone() is not None
|
||||
|
||||
async def is_channel_ingested(self, channel_id: int) -> bool:
|
||||
"""
|
||||
Checks whether a channel has already been bulk-ingested.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
:return: True if the channel has been ingested, False otherwise.
|
||||
"""
|
||||
async with self._connection.execute(
|
||||
"SELECT 1 FROM ingested_channels WHERE channel_id = ?",
|
||||
(channel_id,),
|
||||
) as cursor:
|
||||
return await cursor.fetchone() is not None
|
||||
|
||||
async def mark_channel_ingested(self, channel_id: int) -> None:
|
||||
"""
|
||||
Marks a channel as having been bulk-ingested.
|
||||
|
||||
:param channel_id: The Discord channel ID.
|
||||
"""
|
||||
await self._connection.execute(
|
||||
"INSERT OR IGNORE INTO ingested_channels (channel_id) VALUES (?)",
|
||||
(channel_id,),
|
||||
)
|
||||
await self._maybe_commit()
|
||||
|
||||
async def _maybe_commit(self) -> None:
|
||||
"""Tracks a pending write. Commits if the threshold is reached, otherwise starts a flush timer."""
|
||||
self._pending_writes += 1
|
||||
if self._pending_writes >= AUTO_COMMIT_WRITE_THRESHOLD:
|
||||
await self.commit()
|
||||
elif self._flush_task is None:
|
||||
self._flush_task = asyncio.create_task(self._flush_after_timeout())
|
||||
|
||||
async def _flush_after_timeout(self) -> None:
|
||||
"""Background task that commits after the timeout elapses."""
|
||||
try:
|
||||
await asyncio.sleep(AUTO_COMMIT_TIMEOUT_SECONDS)
|
||||
if self._pending_writes > 0:
|
||||
await self.commit()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("Error in database flush timer.")
|
||||
Reference in New Issue
Block a user