python - 手动插入连接表继承模型的测试数据

Question

我在 SQLAlchemy 的帮助器类中有一个 get_or_create 方法，但 session 永远不会从查询中返回任何数据，因此总是尝试创建新行 - 即使数据库中存在相同的数据(然后抛出一个 UNIQUE CONSTRANT 错误)。

通过 create 方法单独添加数据确实可以成功工作，并且数据会持久保存到数据库中。在 get_or_create() 方法中，实例始终为 None，因此会尝试再次将数据添加到数据库。

下面的代码是完整应用程序的片段，应该足以重现问题。

models.py:

from sqlalchemy import Column, Integer, String, ForeignKey, Float, Enum
from sqlalchemy.ext.declarative import declarative_base
import enum

Base = declarative_base()


class DepositStatus(enum.Enum):
    not_yet_paid = 1
    paid = 2
    to_be_refunded = 3
    deposit_kept = 4
    refunded = 5


class RiderStatus(enum.Enum):
    not_yet_started = 1
    active = 2
    finished = 3
    scratched = 4


class RiderCategories(enum.Enum):
    solo = 1
    pair = 2


class RiderGenders(enum.Enum):
    male = 1
    female = 2


class TrackerLocations(Base):
    __tablename__ = 'tracker_locations'
    id = Column('id', Integer, primary_key=True)
    types = Column(String)
    __mapper_args__ = {
        'polymorphic_on': types
    }


class Riders(TrackerLocations):

    __tablename__ = 'riders'
    id = Column(Integer, ForeignKey('tracker_locations.id'), primary_key=True,
                autoincrement=True)
    first_name = Column('first_name', String)
    last_name = Column('last_name', String)
    cap_number = Column('cap_number', String)
    category = Column('category', Enum(RiderCategories))
    status = Column('status', Enum(RiderStatus))
    gender = Column('gender', Enum(RiderGenders))
    deposit_status = Column('deposit_status', Enum(DepositStatus))
    deposit_amount = Column('deposit_amount', Float)
    email = Column('email', String)

    __mapper_args__ = {
        'polymorphic_identity': 'riders'
    }

db.py:

from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy import create_engine
from models import Base

class DB(object):
    """
    To be used to interact with SQLAlchemy's ORM.
    """

    def __init__(self, db_uri):
        self.engine = create_engine(db_uri)
        Base.metadata.create_all(self.engine)
        self._session = scoped_session(sessionmaker(bind=self.engine))

    def close(self):
        self.session.close()

    def rollback(self):
        self.session.rollback()

    @property
    def session(self):
        return self._session()

    def get_or_create(self, model, commit=True, **kwargs):
        instance = self.session.query(model).filter_by(**kwargs).first()
        if instance:
            return instance, False
        else:
            instance = self.create(model, commit=commit, **kwargs)
            return instance, True

    def create(self, model, commit=True, **kwargs):
        instance = model(**kwargs)
        self.session.add(instance)
        if commit:
            self.session.commit()
        return instance

测试.py:

import os
import sqlite3
from tempfile import mkdtemp
from unittest import TestCase

from models import Riders
from db import DB

class TestDBInteractions(TestCase):

    @classmethod
    def setUpClass(cls):
        cls.temp_ = mkdtemp()
        cls.temp_db = os.path.join(cls.temp_, 'test_db.db')

    def setUp(self):
        self.db_uri = 'sqlite:///' + self.temp_db
        self.db = DB(self.db_uri)
        self.conn = sqlite3.connect(self.temp_db)
        self.cur = self.conn.cursor()

    def tearDown(self):
        self.db.close()

    def test_get_or_create(self):
        test_data = {
            'id': 1,
            'first_name': 'Bobby',
            'last_name': 'Hill',
            'cap_number': '123',
            'category': 'solo',
            'status': 'finished',
            'gender': 'male',
            'deposit_status': 'paid',
            'deposit_amount': 123.1,
            'email': 'hello@bob.com'
        }
        params = tuple(test_data.values())
        query = "INSERT INTO riders VALUES (?,?,?,?,?,?,?,?,?,?)"
        self.cur.execute(query, params)
        # # sqlite doesn't allow multiple write connections to the same DB, so close this connection
        self.conn.commit()
        self.conn.close()
        self.db1 = DB(db_uri=self.db_uri)
        response = self.db1.get_or_create(Riders, **{'id': 1})
        print(response)

Answer 1

与 Joined Table Inheritance ，SQLA 做得非常出色，让您忘记甚至有两个表正在使用。

您的测试设置不会将任何数据插入到 tracker_locations 中，就像您实例化 Riders 实例并通过 ORM 插入数据一样。

这是 SQLA 生成的查询的连接部分:

FROM tracker_locations JOIN riders ON tracker_locations.id = riders.id

在您的测试数据集中，tracker_locations 中没有行，因此连接条件 tracker_locations.id = riders.id 的计算结果不为 true。

因此，如果您还在 tracker_locations 表中包含一些测试数据，它将起作用:

tracker_locations_test_data = {'id': 1, 'type': 'riders'}
self.cur.execute(
    'INSERT INTO tracker_locations VALUES (?,?)', 
    tuple(tracker_locations_test_data.values()))