From c65b67ca46f70e2da46a5b945f4ed358173262e9 Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Sun, 14 Jul 2019 21:28:51 -0700 Subject: [PATCH] .get() method plus support for compound primary keys (#40) * create_table now handles compound primary keys * CLI now accepts multiple --pk for compound primary keys * Docs for compound primary keys with CLI and Python library * New .get() method plus documentation Closes #36, closes #39 --- docs/cli.rst | 2 ++ docs/python-api.rst | 33 ++++++++++++++++++++++ sqlite_utils/cli.py | 6 +++- sqlite_utils/db.py | 67 ++++++++++++++++++++++++++++++++++++-------- tests/test_cli.py | 25 +++++++++++++++++ tests/test_create.py | 22 +++++++++++++++ tests/test_get.py | 30 ++++++++++++++++++++ 7 files changed, 173 insertions(+), 12 deletions(-) diff --git a/docs/cli.rst b/docs/cli.rst index db1954c..cb9574c 100644 --- a/docs/cli.rst +++ b/docs/cli.rst @@ -200,6 +200,8 @@ Here's the simplest possible example:: To specify a column as the primary key, use ``--pk=column_name``. +To create a compound primary key across more than one column, use ``--pk`` multiple times. + If you feed it a JSON list it will insert multiple records. For example, if ``dogs.json`` looks like this:: [ diff --git a/docs/python-api.rst b/docs/python-api.rst index 0cc5007..287bdcb 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -72,6 +72,21 @@ You can filter rows by a WHERE clause using ``.rows_where(where, where_args)``:: ... print(row) {'id': 1, 'age': 4, 'name': 'Cleo'} +.. _python_api_get: + +Retrieving a specific record +============================ + +You can retrieve a record by its primary key using ``table.get()``:: + + >>> db = sqlite_utils.Database("dogs.db") + >>> print(db["dogs"].get(1)) + {'id': 1, 'age': 4, 'name': 'Cleo'} + +If the table has a compound primary key you can pass in the primary key values as a tuple:: + + >>> db["compound_dogs"].get(("mixed", 3)) + Creating tables =============== @@ -147,6 +162,24 @@ The first argument here is a dictionary specifying the columns you would like to This method takes optional arguments ``pk=``, ``column_order=``, ``foreign_keys=``, ``not_null=set()`` and ``defaults=dict()`` - explained below. +.. _python_api_compound_primary_keys: + +Compound primary keys +--------------------- + +If you want to create a table with a compound primary key that spans multiple columns, you can do so by passing a tuple of column names to any of the methods that accept a ``pk=`` parameter. For example: + +.. code-block:: python + + db["cats"].create({ + "id": int, + "breed": str, + "name": str, + "weight": float, + }, pk=("breed", "id")) + +This also works for the ``.insert()``, ``.insert_all()``, ``.upsert()`` and ``.upsert_all()`` methods. + .. _python_api_foreign_keys: Specifying foreign keys diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index 1c8661f..d53a77c 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -302,7 +302,9 @@ def insert_upsert_options(fn): ), click.argument("table"), click.argument("json_file", type=click.File(), required=True), - click.option("--pk", help="Column to use as the primary key, e.g. id"), + click.option( + "--pk", help="Columns to use as the primary key, e.g. id", multiple=True + ), click.option("--nl", is_flag=True, help="Expect newline-delimited JSON"), click.option("-c", "--csv", is_flag=True, help="Expect CSV"), click.option( @@ -348,6 +350,8 @@ def insert_upsert_implementation( if nl and csv: click.echo("Use just one of --nl and --csv", err=True) return + if pk and len(pk) == 1: + pk = pk[0] if csv: reader = csv_std.reader(json_file) headers = next(reader) diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 9b179ac..d40951e 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -79,6 +79,10 @@ class BadPrimaryKey(Exception): pass +class NotFoundError(Exception): + pass + + class Database: def __init__(self, filename_or_conn): if isinstance(filename_or_conn, str): @@ -205,11 +209,15 @@ class Database: raise AlterError( "No such column: {}.{}".format(fk.other_table, fk.other_column) ) - extra = "" + column_defs = [] + # ensure pk is a tuple + single_pk = None + if isinstance(pk, str): + single_pk = pk for column_name, column_type in column_items: column_extras = [] - if pk == column_name: + if column_name == single_pk: column_extras.append("PRIMARY KEY") if column_name in not_null: column_extras.append("NOT NULL") @@ -233,12 +241,17 @@ class Database: else "", ) ) + extra_pk = "" + if single_pk is None and pk and len(pk) > 1: + extra_pk = ",\n PRIMARY KEY ({pks})".format( + pks=", ".join(["[{}]".format(p) for p in pk]) + ) columns_sql = ",\n".join(column_defs) sql = """CREATE TABLE [{table}] ( -{columns_sql} -){extra}; +{columns_sql}{extra_pk} +); """.format( - table=name, columns_sql=columns_sql, extra=extra + table=name, columns_sql=columns_sql, extra_pk=extra_pk ) self.conn.execute(sql) return self[name] @@ -382,6 +395,37 @@ class Table: def pks(self): return [column.name for column in self.columns if column.is_pk] + def get(self, pk_values): + if not isinstance(pk_values, (list, tuple)): + pk_values = [pk_values] + pks = self.pks + pk_names = [] + if len(pks) == 0: + # rowid table + pk_names = ["rowid"] + last_pk = pk_values[0] + elif len(pks) == 1: + pk_names = [pks[0]] + last_pk = pk_values[0] + elif len(pks) > 1: + pk_names = pks + last_pk = pk_values + if len(pk_names) != len(pk_values): + raise NotFoundError( + "Need {} primary key value{}".format( + len(pk_names), "" if len(pk_names) == 1 else "s" + ) + ) + + wheres = ["[{}] = ?".format(pk_name) for pk_name in pk_names] + rows = self.rows_where(" and ".join(wheres), pk_values) + try: + row = list(rows)[0] + self.last_pk = last_pk + return row + except IndexError: + raise NotFoundError + @property def foreign_keys(self): fks = [] @@ -786,12 +830,13 @@ class Table: self.last_pk = None # self.last_rowid will be 0 if a "INSERT OR IGNORE" happened if (hash_id or pk) and self.last_rowid: - self.last_pk = self.db.conn.execute( - "select [{}] from [{}] where rowid = ?".format( - hash_id or pk, self.name - ), - (self.last_rowid,), - ).fetchone()[0] + row = list(self.rows_where("rowid = ?", [self.last_rowid]))[0] + if hash_id: + self.last_pk = row[hash_id] + elif isinstance(pk, str): + self.last_pk = row[pk] + else: + self.last_pk = tuple(row[p] for p in pk) return self def upsert( diff --git a/tests/test_cli.py b/tests/test_cli.py index 47afbb0..0e01814 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -427,6 +427,31 @@ def test_insert_multiple_with_primary_key(db_path, tmpdir): assert ["id"] == db["dogs"].pks +def test_insert_multiple_with_compound_primary_key(db_path, tmpdir): + json_path = str(tmpdir / "dogs.json") + dogs = [ + {"breed": "mixed", "id": i, "name": "Cleo {}".format(i), "age": i + 3} + for i in range(1, 21) + ] + open(json_path, "w").write(json.dumps(dogs)) + result = CliRunner().invoke( + cli.cli, ["insert", db_path, "dogs", json_path, "--pk", "id", "--pk", "breed"] + ) + assert 0 == result.exit_code + db = Database(db_path) + assert dogs == db.execute_returning_dicts("select * from dogs order by breed, id") + assert {"breed", "id"} == set(db["dogs"].pks) + assert ( + "CREATE TABLE [dogs] (\n" + " [breed] TEXT,\n" + " [id] INTEGER,\n" + " [name] TEXT,\n" + " [age] INTEGER,\n" + " PRIMARY KEY ([id], [breed])\n" + ")" + ) == db["dogs"].schema + + def test_insert_not_null_default(db_path, tmpdir): json_path = str(tmpdir / "dogs.json") dogs = [ diff --git a/tests/test_create.py b/tests/test_create.py index 5e76d0c..91ad169 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -55,6 +55,21 @@ def test_create_table(fresh_db): ) == table.schema +def test_create_table_compound_primary_key(fresh_db): + table = fresh_db.create_table( + "test_table", {"id1": str, "id2": str, "value": int}, pk=("id1", "id2") + ) + assert ( + "CREATE TABLE [test_table] (\n" + " [id1] TEXT,\n" + " [id2] TEXT,\n" + " [value] INTEGER,\n" + " PRIMARY KEY ([id1], [id2])\n" + ")" + ) == table.schema + assert ["id1", "id2"] == table.pks + + def test_create_table_with_bad_defaults(fresh_db): with pytest.raises(AssertionError): fresh_db.create_table( @@ -122,6 +137,13 @@ def test_create_table_from_example(fresh_db, example, expected_columns): ] +def test_create_table_from_example_with_compound_primary_keys(fresh_db): + record = {"name": "Zhang", "group": "staff", "employee_id": 2} + table = fresh_db["people"].insert(record, pk=("group", "employee_id")) + assert ["group", "employee_id"] == table.pks + assert record == table.get(("staff", 2)) + + def test_create_table_column_order(fresh_db): fresh_db["table"].insert( collections.OrderedDict( diff --git a/tests/test_get.py b/tests/test_get.py index 75ddaef..ae5ffdf 100644 --- a/tests/test_get.py +++ b/tests/test_get.py @@ -1,4 +1,34 @@ import pytest +from sqlite_utils.db import NotFoundError + + +def test_get_rowid(fresh_db): + dogs = fresh_db["dogs"] + cleo = {"name": "Cleo", "age": 4} + row_id = dogs.insert(cleo).last_rowid + assert cleo == dogs.get(row_id) + + +def test_get_primary_key(fresh_db): + dogs = fresh_db["dogs"] + cleo = {"name": "Cleo", "age": 4, "id": 5} + last_pk = dogs.insert(cleo, pk="id").last_pk + assert 5 == last_pk + assert cleo == dogs.get(5) + + +@pytest.mark.parametrize( + "argument,expected_msg", + [(100, None), (None, None), ((1, 2), "Need 1 primary key value"), ("2", None)], +) +def test_get_not_found(argument, expected_msg, fresh_db): + fresh_db["dogs"].insert( + {"id": 1, "name": "Cleo", "age": 4, "is_good": True}, pk="id" + ) + with pytest.raises(NotFoundError) as excinfo: + fresh_db["dogs"].get(argument) + if expected_msg is not None: + assert expected_msg == excinfo.value.args[0] @pytest.mark.parametrize(