heartbeats/common/db.py

import os

import redis
import uuid
import requests as rq
import base64

class Key:
    @staticmethod
    def users():
        return f"users"
    @staticmethod
    def spotify_token(spotify_id):
        return f"user:{spotify_id}:token"

    @staticmethod
    def spotify_token_refresh(spotify_id):
        return f"user:{spotify_id}:token_refresh"

    @staticmethod
    def client_secret(spotify_id):
        return f"user:{spotify_id}:client_secret"

    @staticmethod
    def email(spotify_id):
        return f"user:{spotify_id}:email"

    @staticmethod
    def user_playlists(spotify_id):
            return f"user:{spotify_id}:playlists"

    @staticmethod
    def playlist_count(playlist_id):
        return f"playlist:{playlist_id}:count"

    @staticmethod
    def playlist_counttype(playlist_id):
        return f"playlist:{playlist_id}:counttype"

    @staticmethod
    def playlist_snapshot(playlist_id):
        return f"playlist:{playlist_id}:snapshot"


class Db:
    def __init__(self, host, port, db, client_id, client_secret):
        self.rdb = redis.Redis(host=host, port=port, db=db, decode_responses=True)
        self.client_id = client_id
        self.client_secret = client_secret

    def get_all_users(self):
        return self.rdb.smembers(Key.users())

    # user secret is forgotten after a year by default, requiring them to log in again
    def add_user(self, client_secret, spotify_token, spotify_id, token_expires, refresh_token, client_secret_expires=31557600):
        self.rdb.setex(Key.client_secret(spotify_id), client_secret_expires, client_secret)

        self.rdb.setex(Key.spotify_token(spotify_id), token_expires, spotify_token)
        self.rdb.set(Key.spotify_token_refresh(spotify_id), refresh_token)
        self.rdb.sadd(Key.users(), spotify_id)

    def user_add_playlist(self, spotify_id, playlist_id, count, counttype):
        self.rdb.sadd(Key.user_playlists(spotify_id), playlist_id)
        self.rdb.set(Key.playlist_count(playlist_id), count)
        self.rdb.set(Key.playlist_counttype(playlist_id), counttype)

    def user_del_playlist(self, spotify_id, playlist_id, delete_on_spotify = True):
        self.rdb.srem(Key.user_playlists(spotify_id), playlist_id)
        self.rdb.delete(Key.playlist_count(playlist_id))
        self.rdb.delete(Key.playlist_counttype(playlist_id))


    def del_user(self, spotify_id):
        self.rdb.delete(Key.spotify_token(spotify_id))
        self.rdb.delete(Key.email(spotify_id))
        self.rdb.srem(Key.users(), spotify_id)

        playlists = self.rdb.smembers(Key.user_playlists(spotify_id))
        for pid in playlists:
            self.rdb.delete(Key.playlist_count(pid))
            self.rdb.delete(Key.playlist_counttype)


    def get_user_auth_token(self, spotify_id):
        auth_token = self.rdb.get(Key.spotify_token(spotify_id))
        if auth_token:
            return auth_token

        refresh_token = self.rdb.get(Key.spotify_token_refresh(spotify_id))
        if not refresh_token:
            return None

        print(f"{refresh_token=}")
        authstring = base64.b64encode(bytes(f"{self.client_id}:{self.client_secret}", 'utf-8')).decode('utf-8')
        res = rq.post('https://accounts.spotify.com/api/token', headers = { 'Authorization': f'Basic {authstring}', 'Content-Type': 'application/x-www-form-urlencoded' }, data = {
            'grant_type': 'refresh_token',
            'refresh_token': refresh_token,
            })
        if res.status_code != 200:
            return None

        print(res.status_code)
        print(res.content)
        res = res.json()
        self.rdb.setex(Key.spotify_token(spotify_id), res['expires_in'], res['access_token'])

        return res['access_token']
        

    def get_user_client_secret(self, spotify_id):
        return self.rdb.get(Key.client_secret(spotify_id))

    def get_user_playlists(self, spotify_id):
        return self.rdb.smembers(Key.user_playlists(spotify_id))

    def get_user_playlist_count(self, playlist_id):
        return int(self.rdb.get(Key.playlist_count(playlist_id)))

    def get_user_playlist_counttype(self, playlist_id):
        return self.rdb.get(Key.playlist_counttype(playlist_id))