Search code examples
pythondiscord.pyrefactoringdry

How to stay dry while connecting to database with every fucntion?


I am working on discordpy bot, I connect and disconnect to database in every single function, I don't think this is a good method but I tried using decorators to connect and disconnect to the database and sending the cursor as a parameter to the function.
here is my code:

def connectToDB(func):
    async def wrapper(*args, **kwargs):
        db = sqlite3.connect(DB_DIR)
        cursor = db.cursor()
        print(*args, **kwargs)

        res = await func(*args, **kwargs, cursor=cursor)

        db.commit()
        cursor.close()
        db.close()
        return res

    return wrapper


class MusicBot(commands.Cog):
    def __init__(self, bot):
        self.bot = bot

    @commands.command(name="myplaylist", aliases=['mypl'])
    @connectToDB
    async def myPlaylist(self, ctx, name=None, currentPage=0, cursor=None):
        currentPage = int(currentPage)
        if name:
            await self.getPlaylistByName(ctx, name, cursor)
            return

        result = cursor.execute(f"SELECT playlist_name, playlist_items, playlist_length, date"
                                f" FROM PLAYLIST WHERE user = ?", (str(ctx.author),))
        playlists = result.fetchall()

        embed = getPlaylistsEmbed(ctx, playlists, currentPage)
        await ctx.send(embed=embed)

def setup(bot):
    bot.add_cog(MusicBot(bot))

But this does not optmize the code, because discordpy must know what parameter do you want for the command, and so in wrapper I have said def wrapper(*args, **kwargs): which doesn't specify any parameter so discordpy assumes I only need the default parameters.

To solve this I have tried:

  1. Using exec to change the needed parameters with each run, and inspect to get the parameters, like so:
def connectToDB(f):
    global func
    func = f
    params = inspect.signature(func)
    exec(
       f"async def wrapper{str(params)}:"
        "    db = sqlite3.connect(DB_DIR)"
        "    cursor = db.cursor()"
        "    print(*args, **kwargs)"
       f"    res = await func{str(params)}"
        "    db.commit()"
        "    cursor.close()"
        "    db.close()"
        "    return res", globals()
    )

    return wrapper

this dosen't work, because after running for each command in the end it will stay with its last change.
   2. asigning the decorators like this

@connectToDB
@commands.command(name="myplaylist", aliases=['mypl'])

   3. this one is quite stupid. Making a decorator factory to get the params

def connectToDB(*args, **kwargs):
    print(args, kwargs)

    def decorator(func):
        async def wrapper(*_args, **_kwargs):
            db = sqlite3.connect(DB_DIR)
            cursor = db.cursor()
            print(*_args, **_kwargs)

            res = await func(*_args, **_kwargs, cursor=cursor)

            db.commit()
            cursor.close()
            db.close()
            return res

        return wrapper

    return decorator


@commands.command(name="myplaylist", aliases=['mypl'])
@connectToDB(self=None, ctx=None, name=None, page=0, msg=None, cursor=None)
async def myPlaylist(self, ctx, name=None, page=0, msg=None, cursor=None):
   ...

Solution

  • This is an issue that occurs quite often when trying to use decorators with libraries that depend on introspecting your functions. The solution is functools.wraps().

    def connectToDB(func):
        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            # setup database...
            res = await func(*args, **kwargs, cursor=cursor)
            # teardown database...
            return res
        return wrapper
    

    functools.wraps calls functools.update_wrapper, which will update "a wrapper function to look like the wrapped function", copying metadata so that introspection tools will report the properties of the original function, not of the wrapper.