How to manage a peewee database in a separate module?

Question

I want to have my database implementation in a separate module or class. But I am struggling with a few details. A simple example:

from peewee import *

db = SqliteDatabase(':memory:')

class BaseModel(Model):
    class Meta:
        database = db

class User(BaseModel):
    name = CharField()

db.connect()
db.create_tables([User,])
db.commit()

@db.atomic()
def add_user(name):
    User.create(name=name).save()

@db.atomic()
def get_user(name):
    return User.get(User.name == name)

So far this is working fine. I can implement my interface to the database here and import this as a module.

Now I want to be able to choose the database file at runtime. So I need a way to define the Model classes without defining SqliteDatabase('somefile') before. I tried to encapsulate everything in a new Database class, which I can later import and create an instance from:

from peewee import *

class Database:

    def __init__(self, dbfile):
        self.db = SqliteDatabase(dbfile)

        class BaseModel(Model):
            class Meta:
                database = self.db

        class User(BaseModel):
            name = CharField()

        self.User = User

        self.db.connect()
        self.db.create_tables([User,])
        self.db.commit()

    @self.db.atomic()    # Error as self is not known on this level
    def add_user(self, name):
        self.User.create(name=name).save()

    @self.db.atomic()    # Error as self is not known on this level
    def get_user(self, name):
        return self.User.get(self.User.name == name)

Now I can call for example database = Database('database.db') or choose any other file name. I can even use multiple database instance in the same program, each with its own file.

However, there are two problems with this approach:

• I still need to specify the database driver (SqliteDatabase) before defining the Model classes. To solve this I define the Model classes within the __init__() method, and then create an alias to with self.User = User. I don't really like this approach (it just doesn't feel like neat code), but at least it works.
• I cannot use the @db.atomic() decorator since self is not known at class level, I would an instance here.

So this class approach does not seem to work very well. Is there some better way to define the Model classes without having to choose where you want to store your database first?

Answer 1

If you need to change database driver at the runtime, then Proxy is a way to go

# database.py
import peewee as pw

proxy = pw.Proxy()

class BaseModel(pw.Model):
  class Meta:
    database = proxy

class User(BaseModel):
  name = pw.CharField()

def add_user(name):
  with proxy.atomic() as txn:
    User.create(name=name).save()

def get_user(name):
  with proxy.atomic() as txn:
    return User.get(User.name == name)

From now on each time you load the module, it won't need a database to be initialized. Instead, you can initialize it at the runtime and switch between multiple as follows

# main.py

import peewee as pw
import database as db

sqlite_1 = pw.SqliteDatabase('sqlite_1.db')
sqlite_2 = pw.PostgresqlDatabase('sqlite_2.db')

db.proxy.initialize(sqlite_1)
sqlite_1.create_tables([db.User], safe=True)
db.add_user(name="Tom")

db.proxy.initialize(sqlite_2)
sqlite_2.create_tables([db.User], safe=True)
db.add_user(name="Jerry")

But if the connection is the only thing that matters, then init() method will be enough.

Answer 2

Now I want to be able to choose the database file at runtime. So I need a way to define the Model classes without defining SqliteDatabase('somefile') before. I tried to encapsulate everything in a new Database class, which I can later import and create an instance from

Peewee uses the meta class to define the name of the table (Model.Meta.db_table) and database( Model.Meta.database)

set these attribute before calling a Model specific code ( either to create a table or to DML statements)

'Allow to define database dynamically'

Answer 3

Question: I cannot use the @db.atomic() decorator since self is not known at class level

Do it, as you do it with self.User.
I wonder about atomic() instead of atomic, but you tell is working fine.

class Database:

    def __init__(self, dbfile):
        self.db = SqliteDatabase(dbfile)

        ...
        @self.db.atomic()
        def __add_user(self, name):
            self.User.create(name=name).save()
        self.add_user = __add_user

        @self.db.atomic()
        def __get_user(self, name):
            return self.User.get(self.User.name == name)
        self.get_user = __get_user

Related: Define models separately from Database() initialization