0
\$\begingroup\$

I am (re)writing a project management CLI and want to put it on a solid foundation concerning the database connections and testing. My focus was to build a DB class which is easy to test and could be reused for any tables I would maybe add at some point. However I am not sure if the duplicate creation of the engine is an issue.

This is my first project which also includes a database, and also my first time trying a TDD approach. So if you would use a wildly different approach to set up the DB connections, I would be also thankful for those alternative ways. Please also feel free to review the testing part, since that's something I am not that familiar with yet too.

db.py consists of two models (they inherit from SQLModel, which is specific to sqlmodel), Project and Account, which define the tables. The class DB is meant as a wrapper around the sqlmodel ORM-functionality, where I reduce the interaction needed for each CRUD-operation to one method.

The tests in test_db.py test all CRUD-operations of the class for both tables, and test one case were both instances of the class are used in the same test. All tests pass (with deprecation warnings from sqlalchemy which are caused by the implementation used by sqlmodel).

For executing the code in your own environment you might need to replace from projects.db import ... with from db import ... in test_db.py (I used a src-layout with an editable pip install for running pytest).

db.py

from typing import Optional

from sqlmodel import Field, Session, SQLModel, create_engine, select


class Account(SQLModel, table=True):
    """Model for the accounts table. Accounts belong to a project."""

    id: Optional[int] | None = Field(default=None, primary_key=True)
    account_number: str
    description: Optional[str] = None

    project_id: Optional[int] = Field(default=None, foreign_key="project.id")


class Project(SQLModel, table=True):
    """Model for the projects table."""

    id: Optional[int] | None = Field(default=None, primary_key=True)
    name: str
    project_number: str


class DB:
    """Database wrapper specific to the supplied database table.

    url: URL of the database file.
    table: Model of the table to which the operations should refer.
            Must be a Subclass of SQLModel.
    """

    def __init__(self, url: str, table: SQLModel, *, echo=False):
        self.url = url
        self.table = table
        self.engine = create_engine(url, echo=echo)

    def create_metadata(self):
        """Creates metadata, call only once per database connection."""
        SQLModel.metadata.create_all(self.engine)

    def read_all(self):
        """Returns all rows of the table."""
        with Session(self.engine) as session:
            projects = session.exec(select(self.table)).all()
            return projects

    def read(self, _id):
        """Returns a row of the table."""
        with Session(self.engine) as session:
            project = session.get(self.table, _id)
            return project

    def add(self, **fields):
        """Adds a row to the table. Fields must map to the table definition."""
        with Session(self.engine) as session:
            entry = self.table(**fields)
            session.add(entry)
            session.commit()

    def update(self, _id, **updates):
        """Updates a row of the table. Updates must map to the table definition."""
        with Session(self.engine) as session:
            entry = self.read(_id)
            for key, val in updates.items():
                setattr(entry, key, val)
            session.add(entry)
            session.commit()

    def delete(self, _id):
        """Delete a row of the table."""
        with Session(self.engine) as session:
            entry = self.read(_id)
            session.delete(entry)
            session.commit()

test_db.py

import sqlite3

import pytest

from projects.db import DB, Account, Project


@pytest.fixture
def account_db(tmp_path):
    db = DB(url=f"sqlite:///{tmp_path}/database_account.db", table=Account, echo=True)
    db.create_metadata()

    conn = sqlite3.Connection(f"{tmp_path}/database_account.db")
    cur = conn.cursor()
    cur.execute(
        """
            INSERT INTO "project" ("name", "project_number") VALUES
            ("Project", "00-001"),
            ("Project2", "01-174"),
            ("Project3", "20-005")
        """
    )
    cur.execute(
        """
            INSERT INTO "account" ("account_number", "description", "project_id") VALUES
            ("10-0450-001", "Primary account for this project.", 1),
            ("10-0450-002", "Account specific for department XY", 1)
        """
    )
    conn.commit()

    return db


@pytest.fixture
def project_db(tmp_path):
    db = DB(url=f"sqlite:///{tmp_path}/database_project.db", table=Project, echo=True)
    db.create_metadata()

    conn = sqlite3.Connection(f"{tmp_path}/database_project.db")
    cur = conn.cursor()
    cur.execute(
        """
                INSERT INTO "project" ("name", "project_number") VALUES
                ("Project", "00-001"),
                ("Project2", "01-174"),
                ("Project3", "20-005")
                """
    )
    conn.commit()

    return db


def test_db_read_all(project_db):
    results = project_db.read_all()
    assert results == [
        {"id": 1, "name": "Project", "project_number": "00-001"},
        {"id": 2, "name": "Project2", "project_number": "01-174"},
        {"id": 3, "name": "Project3", "project_number": "20-005"},
    ]


def test_db_read_one(project_db):
    result = project_db.read(_id=1)
    assert result == {"id": 1, "name": "Project", "project_number": "00-001"}


def test_db_add(project_db):
    project_db.add(name="NewProject", project_number="27-asdf-23487")
    entries = project_db.read_all()
    assert entries[-1].dict()["name"] == "NewProject"
    assert entries[-1].dict()["project_number"] == "27-asdf-23487"


def test_db_update(project_db):
    project_db.update(1, name="ModifiedProject")
    entry = project_db.read(1)
    assert entry == {"id": 1, "name": "ModifiedProject", "project_number": "00-001"}


def test_db_delete(project_db):
    project_db.delete(1)
    deleted_item = project_db.read(1)
    assert deleted_item is None


def test_account_read_all(account_db):
    results = account_db.read_all()
    assert results == [
        {
            "id": 1,
            "account_number": "10-0450-001",
            "description": "Primary account for this project.",
            "project_id": 1,
        },
        {
            "id": 2,
            "account_number": "10-0450-002",
            "description": "Account specific for department XY",
            "project_id": 1,
        },
    ]


def test_account_read(account_db):
    result = account_db.read(_id=1)
    assert result == {
        "id": 1,
        "account_number": "10-0450-001",
        "description": "Primary account for this project.",
        "project_id": 1,
    }


def test_account_add(account_db):
    account_db.add(account_number="20-5000.1-20", project_id=2)
    entries = account_db.read_all()
    assert entries[-1].dict()["account_number"] == "20-5000.1-20"
    assert entries[-1].dict()["project_id"] == 2


def test_account_update(account_db):
    account_db.update(1, account_number="20-000", description="")
    entry = account_db.read(1)
    assert entry == {"id": 1, "account_number": "20-000", "description": "", "project_id": 1}


def test_account_delete(account_db):
    account_db.delete(1)
    deleted_item = account_db.read(1)
    assert deleted_item is None


def test_read_from_both(account_db, tmp_path):
    """Test if the duplicate engine connecting to the db makes issues."""
    project_db = DB(url=f"sqlite:///{tmp_path}/database_account.db", table=Project, echo=True)
    project_result = project_db.read(_id=1)
    account_result = account_db.read(_id=1)
    assert project_result == {"id": 1, "name": "Project", "project_number": "00-001"}
    assert account_result == {
        "id": 1,
        "account_number": "10-0450-001",
        "description": "Primary account for this project.",
        "project_id": 1,
    }

Dependencies are pytest and sqlmodel.

\$\endgroup\$

1 Answer 1

1
\$\begingroup\$

Optional means None

class Account(SQLModel, table=True):
    ...
    id: Optional[int] | None = ...

Tell us that id is Optional[int] (old way), or preferably that it is int | None (new way), but not both.

Also, the notion of a PK allowing a NULL value seems pretty weird. Maybe we can just insist that it is always an int?

Similarly for Project.


class docstring

In DB, move that docstring down so it's describing the __init__ ctor rather than the whole class. Feel free to write a new docstring for the class if you like.

I like the "Must be a subclass of SQLModel" advice. But we're already saying that with the optional type annotation. Consider decorating the class with @beartype for runtime type checking enforcement. Now you can delete that sentence, since caller will get a helpful diagnostic message if he ignores the rules.

    ... , *, echo=False):

Nice signature, I like it! Now caller can't offer an ambiguous ... , False), he is required to describe what it means.


predictable order

My only caveat about the read_all method is it lacks an explicit ORDER BY. Unit tests love predictable row order.

Now I guess this is predictable if we require all tables to have a PK. (I don't let a new table definition sneak past PR code review if it lacks a PK, but I have seen some codebases that are a bit more lax.) The natural order might plausibly be in PK order, though I'm not sure the DB vendor guarantees that behavior across versions.

Consider requiring that the table shall have an ID column, and then you could explicitly ORDER BY that.

Choosing the identifier projects seems a bit odd, given that this method works for all tables.

(BTW in {add, update, delete} I like the explicit .commit().)


execute()

In the account_db fixture, I am a little bit sad that DB doesn't offer an .execute() method for these ad hoc INSERTs. Or at least a method to get a .cursor().

A hardcoded echo=True in the fixtures seems undesirable, now that debugging is done.


multiple databases

not sure if the duplicate creation of the engine is an issue.

When I first read your remark, I thought you had asked sqlmodel twice for engine objects that point at the same underlying database, perhaps due to some calling detail where existing engine was not visible to another layer that needed it.

What we see instead is "one table per database", which is worse.

def account_db(tmp_path):
    db = DB(url=f"sqlite:///{tmp_path}/database_account.db", table=Account, ...)
    ...
def project_db(tmp_path):
    db = DB(url=f"sqlite:///{tmp_path}/database_project.db", table=Project, ...)

Sometimes we need to put tables in separate databases on separate servers. But to the greatest extent possible, strive to keep tables together within a single RDBMS instance. Why? So we can JOIN them.

It's unclear to me how the Account project_id FK reference to project.id even works. Are you maybe using a non-sql92 vendor specific command like ATTACH?

Your code is built atop some very nice DB layers. You happen to be using sqlite now, but this code should work just the same if a connect string is changed to mention a MariaDB or Postgres database. And we would want FKs and JOINs to still work across tables, which won't be feasible if tables are in e.g. a pair of different postgres databases.


This code substantially achieves its design goals, and is accompanied by a very nice test suite.

I would be willing to delegate or accept maintenance tasks on it.

\$\endgroup\$
5
  • \$\begingroup\$ Thank you first of all. Oops, yeah that Optional[int] | None was an oversight. The PK is optional as sqlite sets the PK, that's an implementation detail from sqlmodel. That orderby is a good hint, I'll try that! \$\endgroup\$
    – Koala
    Dec 4, 2023 at 9:31
  • \$\begingroup\$ now for the multiple DBs: You are right, I shouldn't use different DBs for testing. In the CLI app I am using a single db for everything, but for testing I thought it was a good idea to not have the tables interefere? But I guess thats exactly what I'm trying to detect with tests . Just learned about autouse in fixtures though, so thats something I can clean up more easily now (seperate db setup from instantiating the table-specific db instances). \$\endgroup\$
    – Koala
    Dec 4, 2023 at 9:41
  • \$\begingroup\$ The issue with multiple engines however remains: is it a problem if I call create_engine twice? Or should I somehow seperate the engine creation process from the DB class instantiation? Maybe compose with a second Engine class which only has one instance? \$\endgroup\$
    – Koala
    Dec 4, 2023 at 9:45
  • \$\begingroup\$ I tried to tackle the multiple engine issue with some more description, follow up is at stackoverflow.com/q/77598964/18904265 \$\endgroup\$
    – Koala
    Dec 4, 2023 at 11:13
  • 1
    \$\begingroup\$ I have written lots of sqlalchemy DB client apps and honestly, managing an engine is just not a problem. It is an attribute of the app, or of the unit test. Store it as a an object attribute self._engine, or as a function's local variable, or even as a module global. Pass it in to functions that will need it (that's why making it a self attribute is so handy, it comes along for free). Engine lifetime is same duration as your object's lifetime. \$\endgroup\$
    – J_H
    Dec 4, 2023 at 17:51

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.