diff --git a/fastlite/kw.py b/fastlite/kw.py index 109d878..3fc7d4e 100644 --- a/fastlite/kw.py +++ b/fastlite/kw.py @@ -4,6 +4,8 @@ from sqlite_minutils.db import Database,Table,DEFAULT,ForeignKeysType,Default,Queryable,NotFoundError from enum import Enum +class MissingPrimaryKey(Exception): pass + opt_bool = Union[bool, Default, None] def database(path, wal=True)->Any: @@ -203,13 +205,14 @@ def upsert( columns: Union[Dict[str, Any], Default, None]=DEFAULT, strict: Union[bool, Default]|None=DEFAULT, **kwargs) -> Table: + record = _process_row(record) + record = {**record, **kwargs} + if not record: return {} if pk==DEFAULT: assert len(self.pks)==1 pk = self.pks[0] - if not record: record={} - record = _process_row(record) - record = {**record, **kwargs} - last_pk = record[pk] + try: last_pk = record[pk] + except KeyError as e: raise MissingPrimaryKey(e.args[0]) self._orig_upsert( record=record, pk=pk, foreign_keys=foreign_keys, column_order=column_order, not_null=not_null, defaults=defaults, hash_id=hash_id, hash_id_columns=hash_id_columns, alter=alter, diff --git a/nbs/test_upsert.ipynb b/nbs/test_upsert.ipynb new file mode 100644 index 0000000..e1d2326 --- /dev/null +++ b/nbs/test_upsert.ipynb @@ -0,0 +1,467 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "fd325418", + "metadata": {}, + "source": [ + "# Test Upsert Operations" + ] + }, + { + "cell_type": "markdown", + "id": "417f2c4e", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ad470f25", + "metadata": {}, + "outputs": [], + "source": [ + "from fastlite import *" + ] + }, + { + "cell_type": "markdown", + "id": "e4788661", + "metadata": {}, + "source": [ + "Note: Make sure to use fastlite's `database()` here" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "97dd1b48", + "metadata": {}, + "outputs": [], + "source": [ + "db = database(':memory:')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5102a3ac", + "metadata": {}, + "outputs": [], + "source": [ + "class People: id: int; name: str" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9188c149", + "metadata": {}, + "outputs": [], + "source": [ + "people = db.create(People, pk='id')" + ] + }, + { + "cell_type": "markdown", + "id": "6c99cbae", + "metadata": {}, + "source": [ + "## Test Single Upserts" + ] + }, + { + "cell_type": "markdown", + "id": "dbc67ac6", + "metadata": {}, + "source": [ + "Here we test `upsert()`" + ] + }, + { + "cell_type": "markdown", + "id": "a0673d88", + "metadata": {}, + "source": [ + "### Test Cases for `upsert()` Where Nothing Is Inserted" + ] + }, + { + "cell_type": "markdown", + "id": "eb45e038", + "metadata": {}, + "source": [ + "Test that calling `upsert()` without any parameters doesn't change anything, and returns nothing" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fba0c4f7", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{}" + ] + }, + "execution_count": null, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "people.upsert()" + ] + }, + { + "cell_type": "markdown", + "id": "0355fe0a", + "metadata": {}, + "source": [ + "Test None doesn't change anything." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ace59c88", + "metadata": {}, + "outputs": [], + "source": [ + "count = people.count\n", + "assert people.upsert(None) == {}\n", + "assert people.count == count" + ] + }, + { + "cell_type": "markdown", + "id": "2ab1795b", + "metadata": {}, + "source": [ + "Test empty dict doesn't change anything " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a93ec70a", + "metadata": {}, + "outputs": [], + "source": [ + "count = people.count\n", + "assert people.upsert({}) == {}\n", + "assert people.count == count" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "79cd5186", + "metadata": {}, + "outputs": [], + "source": [ + "# Test empty dataclass doesn't change anything\n", + "PersonDC = people.dataclass()\n", + "count = people.count\n", + "assert people.upsert(PersonDC()) == {}\n", + "assert people.count == count" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aa988175", + "metadata": {}, + "outputs": [], + "source": [ + "# Test empty class instance doesn't change anything\n", + "class EmptyPerson: pass\n", + "count = people.count\n", + "assert people.upsert(EmptyPerson()) == {}\n", + "assert people.count == count" + ] + }, + { + "cell_type": "markdown", + "id": "811bc666", + "metadata": {}, + "source": [ + "### Single Insert Types" + ] + }, + { + "cell_type": "markdown", + "id": "157baebb", + "metadata": {}, + "source": [ + "Test upsert with keyword argument without id. Result should be a MissingPrimaryKey error" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1fdd0aaf", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Correct throwing of MissingPrimaryError\n" + ] + } + ], + "source": [ + "try: people.upsert(name='Alice')\n", + "except MissingPrimaryKey: print('Correct throwing of MissingPrimaryError')" + ] + }, + { + "cell_type": "markdown", + "id": "e1300c26", + "metadata": {}, + "source": [ + "Use upsert to insert a new record via a dataclass. Since it can't find the id, it adds the record" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "de73d39a", + "metadata": {}, + "outputs": [], + "source": [ + "person = people.upsert(People(name='Alice', id=people.count+1))" + ] + }, + { + "cell_type": "markdown", + "id": "447e13c9", + "metadata": {}, + "source": [ + "Test upsert that updates with dataclass. Since it can find the id, it updates the record." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c736aa0f", + "metadata": {}, + "outputs": [], + "source": [ + "assert people.upsert(People(name='Bobba', id=person.id)).name == 'Bobba'" + ] + }, + { + "cell_type": "markdown", + "id": "77e6e4c0", + "metadata": {}, + "source": [ + "Use upsert to insert a new record via a class. Since it can't find the id, it adds the record" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dd80748f", + "metadata": {}, + "outputs": [], + "source": [ + "count = people.count\n", + "class Student: pass\n", + "student = Student()\n", + "student.name = 'Daniel Greenfeld'\n", + "student.id = people.count+1\n", + "\n", + "assert people.upsert(student).name == 'Daniel Greenfeld'\n", + "assert people.count == count+1" + ] + }, + { + "cell_type": "markdown", + "id": "0b4eb6df", + "metadata": {}, + "source": [ + "Test upsert that updates with class. Since it can find the id, it updates the record." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cfd90ab0", + "metadata": {}, + "outputs": [], + "source": [ + "count = people.count\n", + "student = Student()\n", + "student.name = 'Daniel Roy Greenfeld'\n", + "student.id = person.id\n", + "\n", + "assert people.upsert(student).name == 'Daniel Roy Greenfeld'\n", + "assert people.count == count" + ] + }, + { + "cell_type": "markdown", + "id": "26a9c38a", + "metadata": {}, + "source": [ + "### None and Empty String Handling" + ] + }, + { + "cell_type": "markdown", + "id": "37ad998d", + "metadata": {}, + "source": [ + "Test upserting a record with name set to None. First assert checks the method result, the second assert tests that the database was altered correctly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5a968d13", + "metadata": {}, + "outputs": [], + "source": [ + "result = people.upsert(People(name=None, id=person.id))\n", + "assert result.name is None\n", + "assert people[person.id].name is None" + ] + }, + { + "cell_type": "markdown", + "id": "dd0c180d", + "metadata": {}, + "source": [ + "Test with empty string." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "92d53608", + "metadata": {}, + "outputs": [], + "source": [ + "result = people.upsert(People(name='', id=person.id))\n", + "assert result.name == ''\n", + "assert people[person.id].name == ''" + ] + }, + { + "cell_type": "markdown", + "id": "d855c6a8", + "metadata": {}, + "source": [ + "### Other Cases" + ] + }, + { + "cell_type": "markdown", + "id": "1ee61d32", + "metadata": {}, + "source": [ + "Test upserts with special characters. Let's do updates first" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "972bab86", + "metadata": {}, + "outputs": [], + "source": [ + "assert people.upsert(People(name='O\\'Connor', id=person.id)).name == \"O'Connor\"\n", + "assert people[person.id].name == \"O'Connor\"\n", + "assert people.upsert(People(name='José', id=person.id)).name == 'José'\n", + "assert people[person.id].name == \"José\"" + ] + }, + { + "cell_type": "markdown", + "id": "b1069ca8", + "metadata": {}, + "source": [ + "Now test special characters with upserts that insert." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2b702435", + "metadata": {}, + "outputs": [], + "source": [ + "person = people.upsert(People(name='O\\'Connor', id=people.count+1))\n", + "assert person.name == \"O'Connor\"\n", + "assert people[person.id].name == \"O'Connor\"\n", + "person = people.upsert(People(name='José', id=people.count+1))\n", + "assert person.name == \"José\"\n", + "assert people[person.id].name == \"José\"" + ] + }, + { + "cell_type": "markdown", + "id": "f27e986a", + "metadata": {}, + "source": [ + "Test dict upsert" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "45a4c2aa", + "metadata": {}, + "outputs": [], + "source": [ + "assert people.upsert({'name': 'Dict Test', 'id': person.id}).name == 'Dict Test'" + ] + }, + { + "cell_type": "markdown", + "id": "f1209e4b", + "metadata": {}, + "source": [ + "Test that extra fields raise `sqlite3.OperationalError`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "07c034e9", + "metadata": {}, + "outputs": [], + "source": [ + "from sqlite3 import OperationalError" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "963008b6", + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " p = people.upsert(dict(name='Extra', age=25, title='Dr', id=person.id))\n", + "except OperationalError as e:\n", + " assert e.args[0] == 'no such column: age'" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "python3", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}