From 9eacd30b1d6e3f1bd138fb330cfea4830197cb1d Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Wed, 8 Aug 2018 16:06:49 -0700 Subject: [PATCH] New column_order= parameter for setting column order --- docs/table.rst | 16 ++++++++++ sqlite_utils/db.py | 72 +++++++++++++++++++++++++++++++++++--------- tests/test_create.py | 23 ++++++++++++++ 3 files changed, 96 insertions(+), 15 deletions(-) diff --git a/docs/table.rst b/docs/table.rst index 988501d..bc69dd2 100644 --- a/docs/table.rst +++ b/docs/table.rst @@ -74,6 +74,22 @@ You can also specify a primary key by passing the ``pk=`` parameter to the ``.in "is_good_dog": True, }, pk="id") +The order of the columns in the table will be derived from the order of the keys in the dictionary, provided you are using Python 3.6 or later. + +If you want to explicitly set the order of the columns you can do so using the ``column_order=`` parameter: + +.. code-block:: python + + dogs.insert({ + "id": 1, + "name": "Cleo", + "twitter": "cleopaws", + "age": 3, + "is_good_dog": True, + }, pk="id", column_order=("id", "twitter", "name")) + +You don't need to pass all of the columns to the ``column_order`` parameter. If you only pass a subset of the columns the remaining columns will be ordered based on the key order of the dictionary. + Creating views ============== diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 88a676a..8b9f597 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -42,12 +42,19 @@ class Database: keys = [d[0] for d in cursor.description] return [dict(zip(keys, row)) for row in cursor.fetchall()] - def create_table(self, name, columns, pk=None, foreign_keys=None): + def create_table( + self, name, columns, pk=None, foreign_keys=None, column_order=None + ): foreign_keys = foreign_keys or [] foreign_keys_by_name = {fk[0]: fk for fk in foreign_keys} + column_items = list(columns.items()) + if column_order is not None: + column_items.sort( + key=lambda p: column_order.index(p[0]) if p[0] in column_order else 999 + ) extra = "" columns_sql = ",\n".join( - " {col_name} {col_type} {primary_key} {references}".format( + " [{col_name}] {col_type} {primary_key} {references}".format( col_name=col_name, col_type={ float: "FLOAT", @@ -66,9 +73,9 @@ class Database: else "" ), ) - for col_name, col_type in columns.items() + for col_name, col_type in column_items ) - sql = """CREATE TABLE {table} ( + sql = """CREATE TABLE [{table}] ( {columns_sql} ){extra}; """.format( @@ -160,9 +167,15 @@ class Table: indexes.append(Index(**row)) return indexes - def create(self, columns, pk=None, foreign_keys=None): + def create(self, columns, pk=None, foreign_keys=None, column_order=None): columns = {name: value for (name, value) in columns.items()} - self.db.create_table(self.name, columns, pk=pk, foreign_keys=foreign_keys) + self.db.create_table( + self.name, + columns, + pk=pk, + foreign_keys=foreign_keys, + column_order=column_order, + ) self.exists = True return self @@ -256,13 +269,25 @@ class Table: ) return self.db.conn.execute(sql, (q,)).fetchall() - def insert(self, record, pk=None, foreign_keys=None, upsert=False): + def insert( + self, record, pk=None, foreign_keys=None, upsert=False, column_order=None + ): return self.insert_all( - [record], pk=pk, foreign_keys=foreign_keys, upsert=upsert + [record], + pk=pk, + foreign_keys=foreign_keys, + upsert=upsert, + column_order=column_order, ) def insert_all( - self, records, pk=None, foreign_keys=None, upsert=False, batch_size=100 + self, + records, + pk=None, + foreign_keys=None, + upsert=False, + batch_size=100, + column_order=None, ): """ Like .insert() but takes a list of records and ensures that the table @@ -270,14 +295,19 @@ class Table: data """ if not self.exists: - self.create(self.detect_column_types(records), pk, foreign_keys) + self.create( + self.detect_column_types(records), + pk, + foreign_keys, + column_order=column_order, + ) all_columns = set() for record in records: all_columns.update(record.keys()) all_columns = list(sorted(all_columns)) for chunk in chunks(records, batch_size): sql = """ - INSERT {upsert} INTO {table} ({columns}) VALUES {rows}; + INSERT {upsert} INTO [{table}] ({columns}) VALUES {rows}; """.format( upsert="OR REPLACE" if upsert else "", table=self.name, @@ -301,11 +331,23 @@ class Table: self.last_id = result.lastrowid return self - def upsert(self, record, pk=None, foreign_keys=None): - return self.insert(record, pk=pk, foreign_keys=foreign_keys, upsert=True) + def upsert(self, record, pk=None, foreign_keys=None, column_order=None): + return self.insert( + record, + pk=pk, + foreign_keys=foreign_keys, + upsert=True, + column_order=column_order, + ) - def upsert_all(self, records, pk=None, foreign_keys=None): - return self.insert_all(records, pk=pk, foreign_keys=foreign_keys, upsert=True) + def upsert_all(self, records, pk=None, foreign_keys=None, column_order=None): + return self.insert_all( + records, + pk=pk, + foreign_keys=foreign_keys, + upsert=True, + column_order=column_order, + ) def chunks(sequence, size): diff --git a/tests/test_create.py b/tests/test_create.py index 3a38284..090aeb5 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -1,5 +1,6 @@ from .fixtures import fresh_db from sqlite_utils.db import Index +import collections import pytest import json @@ -36,6 +37,28 @@ def test_create_table_from_example(fresh_db, example, expected_columns): ] +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"), + ) + assert [ + {"name": "abc", "type": "TEXT"}, + {"name": "ccc", "type": "TEXT"}, + {"name": "zzz", "type": "TEXT"}, + {"name": "bbb", "type": "TEXT"}, + {"name": "aaa", "type": "TEXT"}, + ] == [{"name": col.name, "type": col.type} for col in fresh_db["table"].columns] + + def test_create_table_works_for_m2m_with_only_foreign_keys(fresh_db): fresh_db["one"].insert({"id": 1}, pk="id") fresh_db["two"].insert({"id": 1}, pk="id")