Browse Source

Add database stuff, add asset store in guild

rewrite
noirscape 1 year ago
parent
commit
44f0e261ab
14 changed files with 113 additions and 17 deletions
  1. +6
    -1
      .gitignore
  2. +1
    -0
      TODO
  3. +5
    -3
      client.py
  4. +2
    -1
      config.yml.example
  5. +48
    -0
      db.py
  6. +2
    -0
      db_models/__init__.py
  7. +11
    -0
      db_models/asset_store.py
  8. +2
    -0
      db_models/base.py
  9. +1
    -1
      main.py
  10. +2
    -1
      models/config.py
  11. +21
    -5
      models/management_guild.py
  12. +10
    -4
      modules/admin/guild_log.py
  13. +2
    -1
      requirements.txt
  14. +0
    -0
      storage/.gitkeep

+ 6
- 1
.gitignore View File

@ -1,2 +1,7 @@
# App related content
secure.yml
config.yml
config.yml
storage
# Exclude IDE stuff
.vscode

+ 1
- 0
TODO View File

@ -0,0 +1 @@
- Set up alembic

+ 5
- 3
client.py View File

@ -8,9 +8,10 @@ import discord
# internal files
from models.management_guild import ManagementGuild
from db import DBOperationsManager
class KirigiriBot(commands.Bot):
def __init__(self, management_guild: int, post_ready_callback=None, **kwargs):
def __init__(self, db_path, management_guild: int, post_ready_callback=None, **kwargs):
"""
Initialize KirigiriBot. Accepts the same arguments as commands.Bot, in addition to a couple of extra ones.
@ -20,10 +21,11 @@ class KirigiriBot(commands.Bot):
"""
super().__init__(**kwargs)
self.management_guild_id = management_guild
self.db_manager = DBOperationsManager(db_path)
self._cache_ready = asyncio.Event()
self._post_ready_callback = post_ready_callback
async def get_management_guild(self) -> discord.Guild:
async def get_management_guild(self) -> ManagementGuild:
if not self.is_cache_ready():
await self.wait_for_cache_to_ready()
return self.management_guild
@ -39,7 +41,7 @@ class KirigiriBot(commands.Bot):
Initializes the bot. Sets up the required information for the bot to fully function.
"""
# define management_guild
self.management_guild = ManagementGuild(self.get_guild(self.management_guild_id))
self.management_guild = ManagementGuild(self, self.get_guild(self.management_guild_id))
# mark the cache as ready
self._cache_ready.set()


+ 2
- 1
config.yml.example View File

@ -1,2 +1,3 @@
prefix:
management_guild:
management_guild:
db_path: "sqlite:///storage/db.sqlite"

+ 48
- 0
db.py View File

@ -0,0 +1,48 @@
# stdlib
from typing import Union
# external libraries
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
# internal files
import db_models
class DatabaseManager():
def __init__(self, db_path):
self.engine = create_engine(db_path)
db_models.Base.metadata.bind = self.engine
self.db_session = sessionmaker(bind=self.engine)
class DBAssetStoreManager():
def __init__(self, db_manager):
self.db_manager = db_manager
def get_asset_url(self, id) -> Union[None, str]:
"""
Get an asset's url by the ID.
Returns None if the ID is not found.
"""
session = self.db_manager.db_session()
result = session.query(db_models.AssetStore).filter(db_models.AssetStore.id == id).scalar()
if result: # if this isn't none (scalar returns none if no results)
result = result.url # ...we map the result's url to the result variable
return result
def store_asset(self, id, url) -> str:
"""
Store or update an asset by ID.
Returns the URL of the asset.
"""
new_asset = db_models.AssetStore(id=id, url=url)
session = self.db_manager.db_session()
session.merge(new_asset)
session.commit()
return url
class DBOperationsManager():
def __init__(self, db_path):
self.db_manager = DatabaseManager(db_path)
self.asset_store = DBAssetStoreManager(self.db_manager)

+ 2
- 0
db_models/__init__.py View File

@ -0,0 +1,2 @@
from .base import Base
from .asset_store import AssetStore

+ 11
- 0
db_models/asset_store.py View File

@ -0,0 +1,11 @@
# external libraries
from sqlalchemy import Column, Integer, Text
# internal files
from .base import Base
class AssetStore(Base):
__tablename__ = "asset_store"
id = Column(Integer, primary_key=True)
url = Column(Text)

+ 2
- 0
db_models/base.py View File

@ -0,0 +1,2 @@
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()

+ 1
- 1
main.py View File

@ -25,7 +25,7 @@ def main():
args = parse_cli_args()
config, token = load_config()
if args["run"]:
bot = KirigiriBot(config.management_guild_id, command_prefix=config.prefix)
bot = KirigiriBot(config.db_path, config.management_guild_id, command_prefix=config.prefix)
bot.run(token)
if __name__ == "__main__":

+ 2
- 1
models/config.py View File

@ -1,4 +1,5 @@
class Config:
def __init__(self, configdict):
self.prefix: str = configdict["prefix"]
self.management_guild_id: int = configdict["management_guild"]
self.management_guild_id: int = configdict["management_guild"]
self.db_path: str = configdict["db_path"]

+ 21
- 5
models/management_guild.py View File

@ -1,7 +1,15 @@
# stdlib
import io
# external libraries
import discord
# internal files
from client import KirigiriBot
class ManagementGuild:
def __init__(self, guild: discord.Guild):
def __init__(self, bot: KirigiriBot, guild: discord.Guild):
self.bot = bot
self.guild = guild
# functionality attributes
@ -19,11 +27,19 @@ class ManagementGuild:
# animated
}
async def store_asset(self, id: int, filename: str, url: str):
async def store_asset(self, id: int, filename: str, asset: discord.Asset):
"""
Creates a copy of an URL to an image in the image_log channel.
Creates a copy of a discord Asset using the image-log channel.
ID is a unique way to identify the Asset (use a related discord snowflake)
Returns the new URL.
"""
return str(url)
# TODO: Convert URL to filestream, identify filetype, upload to discord
result = self.bot.db_manager.asset_store.get_asset_url(id)
if not result:
filestream = io.BytesIO()
await asset.save(filestream)
message = await self.image_log.send(file=discord.File(filestream, filename))
url = message.attachments[0].url
result = self.bot.db_manager.asset_store.store_asset(id, url)
return result

+ 10
- 4
modules/admin/guild_log.py View File

@ -19,6 +19,8 @@ class GuildWatcher(commands.Cog):
@commands.Cog.listener()
async def on_guild_join(self, guild: discord.Guild):
guild_log = (await self.bot.get_management_guild()).guild_log
join_time = datetime.datetime.now()
embed = discord.Embed()
icon_url = await self.get_guild_icon(guild)
@ -41,10 +43,12 @@ class GuildWatcher(commands.Cog):
embed.timestamp = join_time
# Send embed
await self.bot.management_guild.guild_log.send(embed=embed)
await guild_log.send(embed=embed)
@commands.Cog.listener()
async def on_guild_remove(self, guild: discord.Guild):
guild_log = (await self.bot.get_management_guild()).guild_log
left_time = datetime.datetime.now()
embed = discord.Embed()
icon_url = await self.get_guild_icon(guild)
@ -59,12 +63,14 @@ class GuildWatcher(commands.Cog):
embed.timestamp = left_time
# Send embed
await self.bot.management_guild.guild_log.send(embed=embed)
await guild_log.send(embed=embed)
async def get_guild_icon(self, guild: discord.Guild):
management_guild = await self.bot.get_management_guild()
if guild.is_icon_animated():
icon_url = await self.bot.management_guild.store_asset(guild.id, f"{guild.id}.gif", guild.icon_url_as(format='gif'))
icon_url = await management_guild.store_asset(guild.id, f"{guild.id}.gif", guild.icon_url_as(format='gif'))
else:
icon_url = await self.bot.management_guild.store_asset(guild.id, f"{guild.id}.png", guild.icon_url_as(format='png'))
icon_url = await management_guild.store_asset(guild.id, f"{guild.id}.png", guild.icon_url_as(format='png'))
return icon_url

+ 2
- 1
requirements.txt View File

@ -1,3 +1,4 @@
discord.py
pyyaml
docopt
docopt
sqlalchemy

+ 0
- 0
storage/.gitkeep View File


Loading…
Cancel
Save