diff --git a/docs/python-api.rst b/docs/python-api.rst index fa814b3..e8c4bda 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -633,7 +633,7 @@ You can set default values for these methods by accessing the table through the # Now you can call .insert() like so: table.insert({"id": 1, "name": "Tracy", "score": 5}) -The configuration options that can be specified in this way are ``pk``, ``foreign_keys``, ``column_order``, ``not_null``, ``defaults``, ``batch_size``, ``hash_id``, ``alter``, ``ignore``, ``replace``, ``extracts``, ``conversions``, ``columns``. These are all documented below. +The configuration options that can be specified in this way are ``pk``, ``foreign_keys``, ``column_order``, ``not_null``, ``defaults``, ``batch_size``, ``hash_id``, ``hash_id_columns``, ``alter``, ``ignore``, ``replace``, ``extracts``, ``conversions``, ``columns``. These are all documented below. .. _python_api_defaults_not_null: @@ -1594,6 +1594,17 @@ If you are going to use that ID straight away, you can access it using ``last_pk }, hash_id="id").last_pk # dog_id is now "f501265970505d9825d8d9f590bfab3519fb20b1" +The hash will be created using all of the column values. To create a hash using a subset of the columns, pass the ``hash_id_columns=`` parameter:: + + db["dogs"].upsert( + {"name": "Cleo", "twitter": "cleopaws", "age": 7}, + hash_id_columns=("name", "twitter") + ) + +The `hash_id=` parameter is optional if you specify ``hash_id_columns=`` - it will default to putting the hash in a column called ``id``. + +You can manually calculate these hashes using the :ref:`hash_record(record, keys=...) ` utility function. + .. _python_api_create_view: Creating views diff --git a/docs/reference.rst b/docs/reference.rst index 4638026..5bd609f 100644 --- a/docs/reference.rst +++ b/docs/reference.rst @@ -68,3 +68,13 @@ sqlite_utils.db.ColumnDetails ----------------------------- .. autoclass:: sqlite_utils.db.ColumnDetails + +sqlite_utils.utils +================== + +.. _reference_utils_hash_record: + +sqlite_utils.utils.hash_record +------------------------------ + +.. autofunction:: sqlite_utils.utils.hash_record diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 1bae5ad..70b452c 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -1,5 +1,6 @@ from .utils import ( chunks, + hash_record, sqlite3, OperationalError, suggest_column_types, @@ -13,7 +14,6 @@ from collections.abc import Mapping import contextlib import datetime import decimal -import hashlib import inspect import itertools import json @@ -663,13 +663,16 @@ class Database: pk: Optional[Any] = None, foreign_keys: Optional[ForeignKeysType] = None, column_order: Optional[List[str]] = None, - not_null: Iterable[str] = None, + not_null: Optional[Iterable[str]] = None, defaults: Optional[Dict[str, Any]] = None, hash_id: Optional[str] = None, + hash_id_columns: Optional[Iterable[str]] = None, extracts: Optional[Union[Dict[str, str], List[str]]] = None, if_not_exists: bool = False, ) -> str: "Returns the SQL ``CREATE TABLE`` statement for creating the specified table." + if hash_id_columns and (hash_id is None): + hash_id = "id" foreign_keys = self.resolve_foreign_keys(name, foreign_keys or []) foreign_keys_by_column = {fk.column: fk for fk in foreign_keys} # any extracts will be treated as integer columns with a foreign key @@ -779,6 +782,7 @@ class Database: not_null: Iterable[str] = None, defaults: Optional[Dict[str, Any]] = None, hash_id: Optional[str] = None, + hash_id_columns: Optional[Iterable[str]] = None, extracts: Optional[Union[Dict[str, str], List[str]]] = None, if_not_exists: bool = False, ) -> "Table": @@ -796,6 +800,7 @@ class Database: not_null=not_null, defaults=defaults, hash_id=hash_id, + hash_id_columns=hash_id_columns, extracts=extracts, if_not_exists=if_not_exists, ) @@ -808,6 +813,7 @@ class Database: not_null=not_null, defaults=defaults, hash_id=hash_id, + hash_id_columns=hash_id_columns, ) return cast(Table, table) @@ -1126,12 +1132,13 @@ class Table(Queryable): defaults: Optional[Dict[str, Any]] = None, batch_size: int = 100, hash_id: Optional[str] = None, + hash_id_columns: Optional[Iterable[str]] = None, alter: bool = False, ignore: bool = False, replace: bool = False, extracts: Optional[Union[Dict[str, str], List[str]]] = None, conversions: Optional[dict] = None, - columns: Optional[Union[Dict[str, Any]]] = None, + columns: Optional[Dict[str, Any]] = None, ): super().__init__(db, name) self._defaults = dict( @@ -1142,6 +1149,7 @@ class Table(Queryable): defaults=defaults, batch_size=batch_size, hash_id=hash_id, + hash_id_columns=hash_id_columns, alter=alter, ignore=ignore, replace=replace, @@ -1314,6 +1322,7 @@ class Table(Queryable): not_null: Iterable[str] = None, defaults: Optional[Dict[str, Any]] = None, hash_id: Optional[str] = None, + hash_id_columns: Optional[Iterable[str]] = None, extracts: Optional[Union[Dict[str, str], List[str]]] = None, if_not_exists: bool = False, ) -> "Table": @@ -1333,6 +1342,7 @@ class Table(Queryable): not_null=not_null, defaults=defaults, hash_id=hash_id, + hash_id_columns=hash_id_columns, extracts=extracts, if_not_exists=if_not_exists, ) @@ -2389,6 +2399,7 @@ class Table(Queryable): chunk, all_columns, hash_id, + hash_id_columns, upsert, pk, conversions, @@ -2400,12 +2411,19 @@ class Table(Queryable): # .execute() method - but some of them may be replaced by # new primary keys if we are extracting any columns. values = [] + if hash_id_columns and hash_id is None: + hash_id = "id" extracts = resolve_extracts(extracts) for record in chunk: record_values = [] for key in all_columns: value = jsonify_if_needed( - record.get(key, None if key != hash_id else _hash(record)) + record.get( + key, + None + if key != hash_id + else hash_record(record, hash_id_columns), + ) ) if key in extracts: extract_table = extracts[key] @@ -2486,6 +2504,7 @@ class Table(Queryable): chunk, all_columns, hash_id, + hash_id_columns, upsert, pk, conversions, @@ -2498,6 +2517,7 @@ class Table(Queryable): chunk, all_columns, hash_id, + hash_id_columns, upsert, pk, conversions, @@ -2527,6 +2547,7 @@ class Table(Queryable): first_half, all_columns, hash_id, + hash_id_columns, upsert, pk, conversions, @@ -2541,6 +2562,7 @@ class Table(Queryable): second_half, all_columns, hash_id, + hash_id_columns, upsert, pk, conversions, @@ -2575,6 +2597,7 @@ class Table(Queryable): not_null: Optional[Union[Set[str], Default]] = DEFAULT, defaults: Optional[Union[Dict[str, Any], Default]] = DEFAULT, hash_id: Optional[Union[str, Default]] = DEFAULT, + hash_id_columns: Optional[Iterable[str]] = DEFAULT, alter: Optional[Union[bool, Default]] = DEFAULT, ignore: Optional[Union[bool, Default]] = DEFAULT, replace: Optional[Union[bool, Default]] = DEFAULT, @@ -2621,6 +2644,7 @@ class Table(Queryable): not_null=not_null, defaults=defaults, hash_id=hash_id, + hash_id_columns=hash_id_columns, alter=alter, ignore=ignore, replace=replace, @@ -2639,6 +2663,7 @@ class Table(Queryable): defaults=DEFAULT, batch_size=DEFAULT, hash_id=DEFAULT, + hash_id_columns=DEFAULT, alter=DEFAULT, ignore=DEFAULT, replace=DEFAULT, @@ -2662,6 +2687,7 @@ class Table(Queryable): defaults = self.value_or_default("defaults", defaults) batch_size = self.value_or_default("batch_size", batch_size) hash_id = self.value_or_default("hash_id", hash_id) + hash_id_columns = self.value_or_default("hash_id_columns", hash_id_columns) alter = self.value_or_default("alter", alter) ignore = self.value_or_default("ignore", ignore) replace = self.value_or_default("replace", replace) @@ -2669,9 +2695,14 @@ class Table(Queryable): conversions = self.value_or_default("conversions", conversions) or {} columns = self.value_or_default("columns", columns) + if hash_id_columns and hash_id is None: + hash_id = "id" + if upsert and (not pk and not hash_id): raise PrimaryKeyRequired("upsert() requires a pk") assert not (hash_id and pk), "Use either pk= or hash_id=" + if hash_id_columns and (hash_id is None): + hash_id = "id" if hash_id: pk = hash_id @@ -2716,6 +2747,7 @@ class Table(Queryable): not_null=not_null, defaults=defaults, hash_id=hash_id, + hash_id_columns=hash_id_columns, extracts=extracts, ) all_columns_set = set() @@ -2738,6 +2770,7 @@ class Table(Queryable): chunk, all_columns, hash_id, + hash_id_columns, upsert, pk, conversions, @@ -2760,6 +2793,7 @@ class Table(Queryable): not_null=DEFAULT, defaults=DEFAULT, hash_id=DEFAULT, + hash_id_columns=DEFAULT, alter=DEFAULT, extracts=DEFAULT, conversions=DEFAULT, @@ -2779,6 +2813,7 @@ class Table(Queryable): not_null=not_null, defaults=defaults, hash_id=hash_id, + hash_id_columns=hash_id_columns, alter=alter, extracts=extracts, conversions=conversions, @@ -2795,6 +2830,7 @@ class Table(Queryable): defaults=DEFAULT, batch_size=DEFAULT, hash_id=DEFAULT, + hash_id_columns=DEFAULT, alter=DEFAULT, extracts=DEFAULT, conversions=DEFAULT, @@ -2813,6 +2849,7 @@ class Table(Queryable): defaults=defaults, batch_size=batch_size, hash_id=hash_id, + hash_id_columns=hash_id_columns, alter=alter, extracts=extracts, conversions=conversions, @@ -3184,14 +3221,6 @@ def jsonify_if_needed(value): return value -def _hash(record): - return hashlib.sha1( - json.dumps(record, separators=(",", ":"), sort_keys=True, default=repr).encode( - "utf8" - ) - ).hexdigest() - - def resolve_extracts( extracts: Optional[Union[Dict[str, str], List[str], Tuple[str]]] ) -> dict: diff --git a/sqlite_utils/utils.py b/sqlite_utils/utils.py index e3386d8..4ec8b9f 100644 --- a/sqlite_utils/utils.py +++ b/sqlite_utils/utils.py @@ -2,12 +2,13 @@ import base64 import contextlib import csv import enum +import hashlib import io import itertools import json import os from . import recipes -from typing import cast, BinaryIO, Iterable, Optional, Tuple, Type +from typing import Dict, cast, BinaryIO, Iterable, Optional, Tuple, Type import click @@ -345,3 +346,31 @@ def chunks(sequence, size): iterator = iter(sequence) for item in iterator: yield itertools.chain([item], itertools.islice(iterator, size - 1)) + + +def hash_record(record: Dict, keys: Optional[Iterable[str]] = None): + """ + ``record`` should be a Python dictionary. Returns a sha1 hash of the + keys and values in that record. + + If ``keys=`` is provided, uses just those keys to generate the hash. + + Example usage:: + + from sqlite_utils.utils import hash_record + + hashed = hash_record({"name": "Cleo", "twitter": "CleoPaws"}) + # Or with the keys= option: + hashed = hash_record( + {"name": "Cleo", "twitter": "CleoPaws", "age": 7}, + keys=("name", "twitter") + ) + """ + to_hash = record + if keys is not None: + to_hash = {key: record[key] for key in keys} + return hashlib.sha1( + json.dumps(to_hash, separators=(",", ":"), sort_keys=True, default=repr).encode( + "utf8" + ) + ).hexdigest() diff --git a/tests/test_create.py b/tests/test_create.py index 35844a1..c3871e1 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -9,7 +9,7 @@ from sqlite_utils.db import ( Table, View, ) -from sqlite_utils.utils import sqlite3 +from sqlite_utils.utils import hash_record, sqlite3 import collections import datetime import decimal @@ -888,6 +888,32 @@ def test_insert_hash_id(fresh_db): assert 1 == dogs.count +@pytest.mark.parametrize("use_table_factory", [True, False]) +def test_insert_hash_id_columns(fresh_db, use_table_factory): + if use_table_factory: + dogs = fresh_db.table("dogs", hash_id_columns=("name", "twitter")) + insert_kwargs = {} + else: + dogs = fresh_db["dogs"] + insert_kwargs = dict(hash_id_columns=("name", "twitter")) + + id = dogs.insert( + {"name": "Cleo", "twitter": "cleopaws", "age": 5}, + **insert_kwargs, + ).last_pk + expected_hash = hash_record({"name": "Cleo", "twitter": "cleopaws"}) + assert id == expected_hash + assert dogs.count == 1 + # Insert replacing a second time should not create a new row + id2 = dogs.insert( + {"name": "Cleo", "twitter": "cleopaws", "age": 6}, + **insert_kwargs, + replace=True, + ).last_pk + assert id2 == expected_hash + assert dogs.count == 1 + + def test_vacuum(fresh_db): fresh_db["data"].insert({"foo": "foo", "bar": "bar"}) fresh_db.vacuum() diff --git a/tests/test_upsert.py b/tests/test_upsert.py index 09bdacc..cef8af0 100644 --- a/tests/test_upsert.py +++ b/tests/test_upsert.py @@ -45,6 +45,30 @@ def test_upsert_with_hash_id(fresh_db): assert "a5e744d0164540d33b1d7ea616c28f2fa97e754a" == table.last_pk +@pytest.mark.parametrize("hash_id", (None, "custom_id")) +def test_upsert_with_hash_id_columns(fresh_db, hash_id): + table = fresh_db["table"] + table.upsert({"a": 1, "b": 2, "c": 3}, hash_id=hash_id, hash_id_columns=("a", "b")) + assert list(table.rows) == [ + { + hash_id or "id": "4acc71e0547112eb432f0a36fb1924c4a738cb49", + "a": 1, + "b": 2, + "c": 3, + } + ] + assert table.last_pk == "4acc71e0547112eb432f0a36fb1924c4a738cb49" + table.upsert({"a": 1, "b": 2, "c": 4}, hash_id=hash_id, hash_id_columns=("a", "b")) + assert list(table.rows) == [ + { + hash_id or "id": "4acc71e0547112eb432f0a36fb1924c4a738cb49", + "a": 1, + "b": 2, + "c": 4, + } + ] + + def test_upsert_compound_primary_key(fresh_db): table = fresh_db["table"] table.upsert_all( diff --git a/tests/test_utils.py b/tests/test_utils.py index 66ce413..e76e97a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -35,3 +35,17 @@ def test_chunks(size, expected): input = ["a", "b", "c", "d"] chunks = list(map(list, utils.chunks(input, size))) assert chunks == expected + + +def test_hash_record(): + expected = "d383e7c0ba88f5ffcdd09be660de164b3847401a" + assert utils.hash_record({"name": "Cleo", "twitter": "CleoPaws"}) == expected + assert ( + utils.hash_record( + {"name": "Cleo", "twitter": "CleoPaws", "age": 7}, keys=("name", "twitter") + ) + == expected + ) + assert ( + utils.hash_record({"name": "Cleo", "twitter": "CleoPaws", "age": 7}) != expected + )