.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

@@ -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::
 
     [

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

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)

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}
+{columns_sql}{extra_pk}
-){extra};
+);
         """.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(
+                    row = list(self.rows_where("rowid = ?", [self.last_rowid]))[0]
-                        "select [{}] from [{}] where rowid = ?".format(
+                    if hash_id:
-                            hash_id or pk, self.name
+                        self.last_pk = row[hash_id]
-                        ),
+                    elif isinstance(pk, str):
-                        (self.last_rowid,),
+                        self.last_pk = row[pk]
-                    ).fetchone()[0]
+                    else:
+                        self.last_pk = tuple(row[p] for p in pk)
         return self
 
     def upsert(

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 = [

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(

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(