mirror of
https://github.com/NoiTheCat/WorldTime.git
synced 2024-10-16 08:29:57 +00:00
119 lines
No EOL
3.8 KiB
Python
119 lines
No EOL
3.8 KiB
Python
# User database abstractions
|
|
|
|
import psycopg2
|
|
from common import tzPrint
|
|
|
|
class UserDatabase:
|
|
def __init__(self, connstr):
|
|
'''
|
|
Sets up the PostgreSQL connection to be used by this instance.
|
|
'''
|
|
self.db = psycopg2.connect(connstr)
|
|
cur = self.db.cursor()
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS userdata (
|
|
guild_id BIGINT,
|
|
user_id BIGINT,
|
|
zone TEXT NOT NULL,
|
|
last_active TIMESTAMPTZ NOT NULL DEFAULT now(),
|
|
PRIMARY KEY (guild_id, user_id)
|
|
)""")
|
|
self.db.commit()
|
|
cur.close()
|
|
|
|
def update_activity(self, serverid : str, authorid : str):
|
|
'''
|
|
If a user exists in the database, updates their last activity timestamp.
|
|
'''
|
|
c = self.db.cursor()
|
|
c.execute("""
|
|
UPDATE userdata SET last_active = now()
|
|
WHERE guild_id = %s AND user_id = %s
|
|
""", (serverid, authorid))
|
|
self.db.commit()
|
|
c.close()
|
|
|
|
def delete_user(self, serverid : str, authorid : str):
|
|
'''
|
|
Deletes existing user from the database.
|
|
'''
|
|
c = self.db.cursor()
|
|
c.execute("""
|
|
DELETE FROM userdata
|
|
WHERE guild_id = %s AND user_id = %s
|
|
""", (serverid, authorid))
|
|
self.db.commit()
|
|
c.close()
|
|
|
|
def update_user(self, serverid : str, authorid : str, zone : str):
|
|
'''
|
|
Insert or update user in the database.
|
|
Does not do any sanitizing of incoming values, as only a small set of
|
|
values are allowed anyway. This is enforced by the caller.
|
|
'''
|
|
self.delete_user(serverid, authorid)
|
|
c = self.db.cursor()
|
|
c.execute("""
|
|
INSERT INTO userdata (guild_id, user_id, zone) VALUES
|
|
(%s, %s, %s)
|
|
ON CONFLICT (guild_id, user_id)
|
|
DO UPDATE SET zone = EXCLUDED.zone
|
|
""", (serverid, authorid, zone))
|
|
self.db.commit()
|
|
c.close()
|
|
|
|
def get_user(self, serverid, userid):
|
|
'''
|
|
Retrieves the time zone name of a single user.
|
|
'''
|
|
c = self.db.cursor()
|
|
c.execute("""
|
|
SELECT zone FROM userdata
|
|
WHERE guild_id = %s and user_id = %s
|
|
""", (serverid, userid))
|
|
result = c.fetchone()
|
|
c.close()
|
|
if result is None: return None
|
|
return result[0]
|
|
|
|
def get_users(self, serverid):
|
|
"""
|
|
Retrieves all user time zones for all recently active members.
|
|
Users not present are not filtered here. Must be handled by the caller.
|
|
Returns a dictionary of lists - Key is formatted zone, value is list of users represented.
|
|
"""
|
|
c = self.db.cursor()
|
|
c.execute("""
|
|
SELECT zone, user_id
|
|
FROM userdata
|
|
WHERE
|
|
last_active >= now() - INTERVAL '30 DAYS' -- only users active in the last 30 days
|
|
AND guild_id = %(guild)s
|
|
AND zone in (SELECT zone from (
|
|
SELECT zone, count(*) as ct
|
|
FROM userdata
|
|
WHERE
|
|
guild_id = %(guild)s
|
|
AND last_active >= now() - INTERVAL '30 DAYS'
|
|
GROUP BY zone
|
|
LIMIT 20
|
|
) as pop_zones)
|
|
ORDER BY RANDOM() -- Randomize display order (expected by consumer)
|
|
""", {'guild': serverid})
|
|
result = {}
|
|
for row in c:
|
|
resultrow = tzPrint(row[0])
|
|
result[resultrow] = result.get(resultrow, list())
|
|
result[resultrow].append(row[1])
|
|
c.close()
|
|
return result
|
|
|
|
def get_unique_tz_count(self):
|
|
'''
|
|
Gets the number of unique time zones in the database.
|
|
'''
|
|
c = self.db.cursor()
|
|
c.execute('SELECT COUNT(DISTINCT zone) FROM userdata')
|
|
result = c.fetchall()
|
|
c.close()
|
|
return result[0][0] |