diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index e03ef63..39a7416 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -18,6 +18,7 @@ ForeignKey = namedtuple( "ForeignKey", ("table", "column", "other_table", "other_column") ) Index = namedtuple("Index", ("seq", "name", "unique", "origin", "partial", "columns")) +DEFAULT = object() COLUMN_TYPE_MAPPING = { float: "FLOAT", @@ -93,11 +94,14 @@ class Database: self.conn = filename_or_conn def __getitem__(self, table_name): - return Table(self, table_name) + return self.table(table_name) def __repr__(self): return "".format(self.conn) + def table(self, table_name, **kwargs): + return Table(self, table_name, **kwargs) + def escape(self, value): # Normally we would use .execute(sql, [params]) for escaping, but # occasionally that isn't available - most notable when we need @@ -343,10 +347,36 @@ class Database: class Table: - def __init__(self, db, name): + def __init__( + self, + db, + name, + pk=None, + foreign_keys=None, + column_order=None, + not_null=None, + defaults=None, + upsert=False, + batch_size=100, + hash_id=None, + alter=False, + ignore=False, + ): self.db = db self.name = name self.exists = self.name in self.db.table_names() + self._defaults = dict( + pk=pk, + foreign_keys=foreign_keys, + column_order=column_order, + not_null=not_null, + defaults=defaults, + upsert=upsert, + batch_size=batch_size, + hash_id=hash_id, + alter=alter, + ignore=ignore, + ) def __repr__(self): return "".format( @@ -716,18 +746,21 @@ class Table: ) return self.db.conn.execute(sql, (q,)).fetchall() + def value_or_default(self, key, value): + return self._defaults[key] if value is DEFAULT else value + def insert( self, record, - pk=None, - foreign_keys=None, - column_order=None, - not_null=None, - defaults=None, - upsert=False, - hash_id=None, - alter=False, - ignore=False, + pk=DEFAULT, + foreign_keys=DEFAULT, + column_order=DEFAULT, + not_null=DEFAULT, + defaults=DEFAULT, + upsert=DEFAULT, + hash_id=DEFAULT, + alter=DEFAULT, + ignore=DEFAULT, ): return self.insert_all( [record], @@ -745,22 +778,33 @@ class Table: def insert_all( self, records, - pk=None, - foreign_keys=None, - column_order=None, - not_null=None, - defaults=None, - upsert=False, - batch_size=100, - hash_id=None, - alter=False, - ignore=False, + pk=DEFAULT, + foreign_keys=DEFAULT, + column_order=DEFAULT, + not_null=DEFAULT, + defaults=DEFAULT, + upsert=DEFAULT, + batch_size=DEFAULT, + hash_id=DEFAULT, + alter=DEFAULT, + ignore=DEFAULT, ): """ Like .insert() but takes a list of records and ensures that the table that it creates (if table does not exist) has columns for ALL of that data """ + pk = self.value_or_default("pk", pk) + foreign_keys = self.value_or_default("foreign_keys", foreign_keys) + column_order = self.value_or_default("column_order", column_order) + not_null = self.value_or_default("not_null", not_null) + defaults = self.value_or_default("defaults", defaults) + upsert = self.value_or_default("upsert", upsert) + batch_size = self.value_or_default("batch_size", batch_size) + hash_id = self.value_or_default("hash_id", hash_id) + alter = self.value_or_default("alter", alter) + ignore = self.value_or_default("ignore", ignore) + assert not (hash_id and pk), "Use either pk= or hash_id=" assert not ( ignore and upsert @@ -842,13 +886,13 @@ class Table: def upsert( self, record, - pk=None, - foreign_keys=None, - column_order=None, - not_null=None, - defaults=None, - hash_id=None, - alter=False, + pk=DEFAULT, + foreign_keys=DEFAULT, + column_order=DEFAULT, + not_null=DEFAULT, + defaults=DEFAULT, + hash_id=DEFAULT, + alter=DEFAULT, ): return self.insert( record, @@ -865,14 +909,14 @@ class Table: def upsert_all( self, records, - pk=None, - foreign_keys=None, - column_order=None, - not_null=None, - defaults=None, - batch_size=100, - hash_id=None, - alter=False, + pk=DEFAULT, + foreign_keys=DEFAULT, + column_order=DEFAULT, + not_null=DEFAULT, + defaults=DEFAULT, + batch_size=DEFAULT, + hash_id=DEFAULT, + alter=DEFAULT, ): return self.insert_all( records, diff --git a/tests/test_create.py b/tests/test_create.py index b6519c8..30131fc 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -144,19 +144,22 @@ def test_create_table_from_example_with_compound_primary_keys(fresh_db): assert record == table.get(("staff", 2)) -def test_create_table_column_order(fresh_db): - fresh_db["table"].insert( - collections.OrderedDict( - ( - ("zzz", "third"), - ("abc", "first"), - ("ccc", "second"), - ("bbb", "second-to-last"), - ("aaa", "last"), - ) - ), - column_order=("abc", "ccc", "zzz"), +@pytest.mark.parametrize("use_class_constructor", [True, False]) +def test_create_table_column_order(fresh_db, use_class_constructor): + row = collections.OrderedDict( + ( + ("zzz", "third"), + ("abc", "first"), + ("ccc", "second"), + ("bbb", "second-to-last"), + ("aaa", "last"), + ) ) + column_order = ("abc", "ccc", "zzz") + if use_class_constructor: + fresh_db.table("table", column_order=column_order).insert(row) + else: + fresh_db["table"].insert(row, column_order=column_order) assert [ {"name": "abc", "type": "TEXT"}, {"name": "ccc", "type": "TEXT"}, @@ -191,21 +194,31 @@ def test_create_table_column_order(fresh_db): (({"one_id": "one"},), AssertionError), ), ) +@pytest.mark.parametrize("use_class_constructor", [True, False]) def test_create_table_works_for_m2m_with_only_foreign_keys( - fresh_db, foreign_key_specification, expected_exception + fresh_db, foreign_key_specification, expected_exception, use_class_constructor ): - fresh_db["one"].insert({"id": 1}, pk="id") - fresh_db["two"].insert({"id": 1}, pk="id") + if use_class_constructor: + fresh_db.table("one", pk="id").insert({"id": 1}) + fresh_db.table("two", pk="id").insert({"id": 1}) + else: + fresh_db["one"].insert({"id": 1}, pk="id") + fresh_db["two"].insert({"id": 1}, pk="id") + + row = {"one_id": 1, "two_id": 1} + + def do_it(): + if use_class_constructor: + fresh_db.table("m2m", foreign_keys=foreign_key_specification).insert(row) + else: + fresh_db["m2m"].insert(row, foreign_keys=foreign_key_specification) + if expected_exception: with pytest.raises(expected_exception): - fresh_db["m2m"].insert( - {"one_id": 1, "two_id": 1}, foreign_keys=foreign_key_specification - ) + do_it() return else: - fresh_db["m2m"].insert( - {"one_id": 1, "two_id": 1}, foreign_keys=foreign_key_specification - ) + do_it() assert [ {"name": "one_id", "type": "INTEGER"}, {"name": "two_id", "type": "INTEGER"}, @@ -407,7 +420,10 @@ def test_index_foreign_keys(fresh_db): ), ], ) -def test_insert_row_alter_table(fresh_db, extra_data, expected_new_columns): +@pytest.mark.parametrize("use_class_constructor", [True, False]) +def test_insert_row_alter_table( + fresh_db, extra_data, expected_new_columns, use_class_constructor +): table = fresh_db["books"] table.insert({"title": "Hedgehogs of the world", "author_id": 1}) assert [ @@ -416,7 +432,10 @@ def test_insert_row_alter_table(fresh_db, extra_data, expected_new_columns): ] == [{"name": col.name, "type": col.type} for col in table.columns] record = {"title": "Squirrels of the world", "author_id": 2} record.update(extra_data) - fresh_db["books"].insert(record, alter=True) + if use_class_constructor: + fresh_db.table("books", alter=True).insert(record) + else: + fresh_db["books"].insert(record, alter=True) assert [ {"name": "title", "type": "TEXT"}, {"name": "author_id", "type": "INTEGER"}, @@ -425,21 +444,26 @@ def test_insert_row_alter_table(fresh_db, extra_data, expected_new_columns): ] -def test_upsert_rows_alter_table(fresh_db): - table = fresh_db["books"] - table.insert({"id": 1, "title": "Hedgehogs of the world", "author_id": 1}, pk="id") - table.upsert_all( - [ - {"id": 1, "title": "Hedgehogs of the World", "species": "hedgehogs"}, - {"id": 2, "title": "Squirrels of the World", "num_species": 200}, - { - "id": 3, - "title": "Badgers of the World", - "significant_continents": ["Europe", "North America"], - }, - ], - alter=True, - ) +@pytest.mark.parametrize("use_class_constructor", [True, False]) +def test_upsert_rows_alter_table(fresh_db, use_class_constructor): + first_row = {"id": 1, "title": "Hedgehogs of the world", "author_id": 1} + next_rows = [ + {"id": 1, "title": "Hedgehogs of the World", "species": "hedgehogs"}, + {"id": 2, "title": "Squirrels of the World", "num_species": 200}, + { + "id": 3, + "title": "Badgers of the World", + "significant_continents": ["Europe", "North America"], + }, + ] + if use_class_constructor: + table = fresh_db.table("books", pk="id", alter=True) + table.insert(first_row) + table.upsert_all(next_rows) + else: + table = fresh_db["books"] + table.insert(first_row, pk="id") + table.upsert_all(next_rows, alter=True) assert { "author_id": int, "id": int,