Added db_connection hook as an example (MarkusRost's request)

This commit is contained in:
Frisk 2024-09-08 00:08:02 +02:00
parent 59452b4a5a
commit ac26ca1b96
2 changed files with 81 additions and 0 deletions

View file

@ -16,3 +16,4 @@
#import extensions.hooks.example_hook
#import extensions.hooks.usertalk
#import extensions.hooks.edit_alerts
#import extensions.hooks.db_connection

View file

@ -0,0 +1,80 @@
# This file is part of Recent changes Goat compatible Discord webhook (RcGcDw).
#
# RcGcDw is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# RcGcDw is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with RcGcDw. If not, see <http://www.gnu.org/licenses/>.
import psycopg2
# from src.api.context import Context
# from src.discord.message import DiscordMessage, DiscordMessageMetadata
from src.api.hook import post_hook
from src.configloader import settings
from contextlib import contextmanager
# {
# "hooks": {
# "dbconn": {
# "database": "database",
# "host": "localhost",
# "port": 5432,
# "user": "rcgcdw",
# "password": "secret",
# "socket": null
# }
# }
# }
dbconn = settings.get("hooks", {}).get("dbconn", [])
class DBConnector:
def __init__(self):
self.session = psycopg2.connect(dbname=dbconn.get("database", "rcgcdw"), user=dbconn.get("user", "rcgcdw"), password=dbconn.get("password", "rcgcdw"),
host=dbconn.get("host", "localhost"), port=dbconn.get("port", "5432"))
@contextmanager
def obtain_cursor(self) -> psycopg2._psycopg.cursor:
"""Obtain cursor for database connection"""
cur = self.session.cursor()
try:
yield cur
finally:
self.session.commit()
cur.close()
def shutdown(self): # placeholder, rcgcdw doesn't signal to its extensions when shutdown occurs so we cannot cleanup
self.session.close()
db_singleton = DBConnector()
@post_hook
def db_interaction_hook(message, metadata, context, change):
# Check if user is in the database
if not dbconn:
return
with db_singleton.obtain_cursor() as cursor:
cursor.execute(
"INSERT INTO rcgcdw (username, edit_amount) VALUES (%s, 1) ON conflict (username) DO UPDATE SET edit_amount = rcgcdw.edit_amount + 1;",
(change["user"],))
cursor.execute("SELECT edit_amount FROM rcgcdw WHERE username = %s;", (change["user"],))
record = cursor.fetchone()
print(record)
if record is None or record[0] == 1:
if isinstance(message["description"], str):
message["description"] = message["description"] + "\n**This user is new! Let them know they are awesome! :)**"
else:
message["description"] = "\n**This user is new! Let them know they are awesome! :)**"
# if not record:
# cursor.execute("INSERT INTO rcgcdw (user, edit_amount) VALUES (%s, %s) ON CONFLICT DO UPDATE SET edit_amount = edit_amount + 1", (change["user"], 1))
# else:
# cursor.execute("UPDATE rcgcdw SET edit_amount = edit_amount + 1 WHERE user = %s;", (change["user"],))