diff --git a/.gitignore b/.gitignore index 4d248a3..4a13b25 100644 --- a/.gitignore +++ b/.gitignore @@ -159,9 +159,8 @@ cython_debug/ # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ -# some configuration files -tracked_accounts.json -cookies.json +# data +*.db # automatically generated files session.tw_session diff --git a/README.md b/README.md index dacc799..fb78211 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ Tweetcord is a discord bot that uses the tweety module to let you receive tweet 👇When the followed user posts a new tweet, your server will also receive a notification. -![](https://i.imgur.com/lavcfOz.png) +![](https://i.imgur.com/SXITM0a.png) diff --git a/bot.py b/bot.py index bc53c58..ed215dd 100644 --- a/bot.py +++ b/bot.py @@ -3,9 +3,9 @@ from discord import app_commands from dotenv import load_dotenv import os -import json from src.log import setup_logger +from src.init_db import init_db from configs.load_configs import configs log = setup_logger(__name__) @@ -18,8 +18,7 @@ @bot.event async def on_ready(): await bot.change_presence(activity=discord.Activity(type=discord.ActivityType.watching, name=configs['activity_name'])) - if not(os.path.isfile(f"{os.getenv('DATA_PATH')}tracked_accounts.json")): - with open(f"{os.getenv('DATA_PATH')}tracked_accounts.json", 'w', encoding='utf8') as jfile: json.dump(dict(), jfile, sort_keys=True, indent=4) + if not(os.path.isfile(f"{os.getenv('DATA_PATH')}tracked_accounts.db")): init_db() bot.tree.on_error = on_tree_error for filename in os.listdir('./cogs'): if filename.endswith('.py'): @@ -53,17 +52,16 @@ async def reload(ctx, extension): @bot.command() @commands.is_owner() async def download_data(ctx : commands.context.Context): - message = await ctx.send(file=discord.File(f"{os.getenv('DATA_PATH')}tracked_accounts.json")) + message = await ctx.send(file=discord.File(f"{os.getenv('DATA_PATH')}tracked_accounts.db")) await message.delete(delay=15) @bot.command() @commands.is_owner() async def upload_data(ctx : commands.context.Context): - raw = await [attachment for attachment in ctx.message.attachments if attachment.filename[-4:] == '.txt'][0].read() - data = json.loads(raw) - with open(f"{os.getenv('DATA_PATH')}tracked_accounts.json", 'w', encoding='utf8') as jfile: - json.dump(data, jfile, sort_keys=True, indent=4) + raw = await [attachment for attachment in ctx.message.attachments if attachment.filename[-3:] == '.db'][0].read() + with open(f"{os.getenv('DATA_PATH')}tracked_accounts.db", 'wb') as wbf: + wbf.write(raw) message = await ctx.send('successfully uploaded data') await message.delete(delay=5) diff --git a/cogs/notification.py b/cogs/notification.py index 8a10540..40eda17 100644 --- a/cogs/notification.py +++ b/cogs/notification.py @@ -5,7 +5,7 @@ from datetime import datetime, timezone from dotenv import load_dotenv import os -import json +import sqlite3 from src.log import setup_logger from src.notification.account_tracker import AccountTracker @@ -39,10 +39,15 @@ async def notifier(self, itn : discord.Interaction, username: str, channel: disc """ await itn.response.defer(ephemeral=True) - with open(f"{os.getenv('DATA_PATH')}tracked_accounts.json", 'r', encoding='utf8') as jfile: - users = json.load(jfile) - match_user = list(filter(lambda item: item[1]["username"] == username, users.items())) - if match_user == []: + + conn = sqlite3.connect(f"{os.getenv('DATA_PATH')}tracked_accounts.db") + cursor = conn.cursor() + + cursor.execute(f"SELECT * FROM user WHERE username='{username}'") + match_user = cursor.fetchone() + + roleID = str(mention.id) if mention != None else '' + if match_user == None: app = Twitter("session") app.load_auth_token(os.getenv('TWITTER_TOKEN')) try: @@ -50,21 +55,23 @@ async def notifier(self, itn : discord.Interaction, username: str, channel: disc except: await itn.followup.send(f'user {username} not found', ephemeral=True) return - roleID = str(mention.id) if mention != None else '' - users[str(new_user.id)] = {'username': username, 'channels': {str(channel.id): roleID}, 'lastest_tweet': datetime.utcnow().replace(tzinfo=timezone.utc).strftime('%Y-%m-%d %H:%M:%S%z')} + + cursor.execute('INSERT INTO user VALUES (?, ?, ?)', (str(new_user.id), username, datetime.utcnow().replace(tzinfo=timezone.utc).strftime('%Y-%m-%d %H:%M:%S%z'))) + cursor.execute('INSERT OR IGNORE INTO channel VALUES (?)', (str(channel.id),)) + cursor.execute('INSERT INTO notification VALUES (?, ?, ?)', (str(new_user.id), str(channel.id), roleID)) app.follow_user(new_user) if app.enable_user_notification(new_user): log.info(f'successfully opened notification for {username}') else: log.warning(f'unable to turn on notifications for {username}') else: - user = match_user[0][1] - user['channels'][str(channel.id)] = str(mention.id) if mention != None else '' - - with open(f"{os.getenv('DATA_PATH')}tracked_accounts.json", 'w', encoding='utf8') as jfile: - json.dump(users, jfile, sort_keys=True, indent=4) + cursor.execute('INSERT OR IGNORE INTO channel VALUES (?)', (str(channel.id),)) + cursor.execute('REPLACE INTO notification VALUES (?, ?, ?)', (match_user[0], str(channel.id), roleID)) + + conn.commit() + conn.close() - if match_user == []: await self.account_tracker.addTask(username) + if match_user == None: await self.account_tracker.addTask(username) await itn.followup.send(f'successfully add notifier of {username}!', ephemeral=True) diff --git a/src/init_db.py b/src/init_db.py new file mode 100644 index 0000000..d7d990c --- /dev/null +++ b/src/init_db.py @@ -0,0 +1,13 @@ +import os +import sqlite3 + +def init_db(): + conn = sqlite3.connect(f"{os.getenv('DATA_PATH')}tracked_accounts.db") + cursor = conn.cursor() + cursor.executescript(""" + CREATE TABLE IF NOT EXISTS user (id TEXT PRIMARY KEY, username TEXT, lastest_tweet TEXT); + CREATE TABLE IF NOT EXISTS channel (id TEXT PRIMARY KEY); + CREATE TABLE IF NOT EXISTS notification (user_id TEXT, channel_id TEXT, role_id TEXT, FOREIGN KEY (user_id) REFERENCES user (id), FOREIGN KEY (channel_id) REFERENCES channel (id), PRIMARY KEY(user_id, channel_id)); + """) + conn.commit() + conn.close() \ No newline at end of file diff --git a/src/notification/account_tracker.py b/src/notification/account_tracker.py index 9e7109f..28bc6ec 100644 --- a/src/notification/account_tracker.py +++ b/src/notification/account_tracker.py @@ -3,7 +3,7 @@ from dotenv import load_dotenv from datetime import datetime, timedelta import os -import json +import sqlite3 import asyncio from src.log import setup_logger @@ -22,19 +22,23 @@ def __init__(self, bot): self.tasksMonitorLogAt = datetime.utcnow() - timedelta(seconds=configs['tasks_monitor_log_period']) bot.loop.create_task(self.setup_tasks()) - async def setup_tasks(self): + async def setup_tasks(self): app = Twitter("session") app.load_auth_token(os.getenv('TWITTER_TOKEN')) - - with open(f"{os.getenv('DATA_PATH')}tracked_accounts.json", 'r', encoding='utf8') as jfile: - users = json.load(jfile) + + conn = sqlite3.connect(f"{os.getenv('DATA_PATH')}tracked_accounts.db") + cursor = conn.cursor() + self.bot.loop.create_task(self.tweetsUpdater(app)).set_name('TweetsUpdater') + cursor.execute('SELECT username FROM user') usernames = [] - for user in users.values(): - username = user['username'] + for user in cursor: + username = user[0] usernames.append(username) self.bot.loop.create_task(self.notification(username)).set_name(username) self.bot.loop.create_task(self.tasksMonitor(set(usernames))).set_name('TasksMonitor') + + conn.close() async def notification(self, username): @@ -45,20 +49,22 @@ async def notification(self, username): await task lastest_tweet = task.result() if lastest_tweet == None: continue + + conn = sqlite3.connect(f"{os.getenv('DATA_PATH')}tracked_accounts.db") + conn.row_factory = sqlite3.Row + cursor = conn.cursor() - with open(f"{os.getenv('DATA_PATH')}tracked_accounts.json", 'r', encoding='utf8') as jfile: - users = json.load(jfile) - - user = list(filter(lambda item: item[1]["username"] == username, users.items()))[0][1] + user = cursor.execute('SELECT * FROM user WHERE username = ?', (username,)).fetchone() if date_comparator(lastest_tweet.created_on, user['lastest_tweet']) == 1: - user['lastest_tweet'] = str(lastest_tweet.created_on) + cursor.execute('UPDATE user SET lastest_tweet = ? WHERE username = ?', (str(lastest_tweet.created_on), username)) log.info(f'find a new tweet from {username}') - with open(f"{os.getenv('DATA_PATH')}tracked_accounts.json", 'w', encoding='utf8') as jfile: - json.dump(users, jfile, sort_keys=True, indent=4) - for chnl in user['channels'].keys(): - channel = self.bot.get_channel(int(chnl)) - mention = f"{channel.guild.get_role(int(user['channels'][chnl])).mention} " if user['channels'][chnl] != '' else '' + for data in cursor.execute('SELECT * FROM notification WHERE user_id = ?', (user['id'],)): + channel = self.bot.get_channel(int(data['channel_id'])) + mention = f"{channel.guild.get_role(int(data['role_id'])).mention} " if data['role_id'] != '' else '' await channel.send(f"{mention}**{lastest_tweet.author.name}** just {get_action(lastest_tweet)} here: \n{lastest_tweet.url}", file = discord.File('images/twitter.png', filename='twitter.png'), embeds = gen_embed(lastest_tweet)) + + conn.commit() + conn.close() async def tweetsUpdater(self, app): @@ -95,8 +101,9 @@ async def tasksMonitor(self, users : set): async def addTask(self, username : str): - with open(f"{os.getenv('DATA_PATH')}tracked_accounts.json", 'r', encoding='utf8') as jfile: - users = json.load(jfile) + conn = sqlite3.connect(f"{os.getenv('DATA_PATH')}tracked_accounts.db") + cursor = conn.cursor() + self.bot.loop.create_task(self.notification(username)).set_name(username) log.info(f'new task {username} added successfully') @@ -104,5 +111,7 @@ async def addTask(self, username : str): if task.get_name() == 'TasksMonitor': try: log.info(f'existing TasksMonitor has been closed') if task.cancel() else log.info('existing TasksMonitor failed to close') except Exception as e: log.warning(f'addTask : {e}') - self.bot.loop.create_task(self.tasksMonitor(set([user['username'] for user in users.values()]))).set_name('TasksMonitor') - log.info(f'new TasksMonitor has been started') \ No newline at end of file + self.bot.loop.create_task(self.tasksMonitor({user[0] for user in cursor.execute('SELECT username FROM user').fetchall()})).set_name('TasksMonitor') + log.info(f'new TasksMonitor has been started') + + conn.close() \ No newline at end of file