diff --git a/docs/python-api.rst b/docs/python-api.rst index 68ef7e2..fb060b4 100644 --- a/docs/python-api.rst +++ b/docs/python-api.rst @@ -182,6 +182,40 @@ Note that the ``pk`` and ``column_order`` parameters here are optional if you ar An ``upsert_all()`` method is also available, which behaves like ``insert_all()`` but performs upserts instead. + +Adding columns +============== + +You can add a new column to a table using the ``.add_column(col_name, col_type)`` method: + +.. code-block:: python + + db["dogs"].add_column("dob", datetime.date) + +You can pass any of the following Python types as the second argument - ``sqlite-utils`` will create the column using the appropriate SQLite type from this list:: + + float: "FLOAT" + int: "INTEGER" + bool: "INTEGER" + str: "TEXT" + bytes: "BLOB" + datetime.datetime: "TEXT" + datetime.date: "TEXT" + datetime.time: "TEXT" + + # If numpy is installed + np.int8: "INTEGER" + np.int16: "INTEGER" + np.int32: "INTEGER" + np.int64: "INTEGER" + np.uint8: "INTEGER" + np.uint16: "INTEGER" + np.uint32: "INTEGER" + np.uint64: "INTEGER" + np.float16: "FLOAT" + np.float32: "FLOAT" + np.float64: "FLOAT" + .. _python_api_hash: Setting an ID based on the hash of the row contents diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 996a501..3a41434 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -19,6 +19,36 @@ ForeignKey = namedtuple( ) Index = namedtuple("Index", ("seq", "name", "unique", "origin", "partial", "columns")) +COLUMN_TYPE_MAPPING = { + float: "FLOAT", + int: "INTEGER", + bool: "INTEGER", + str: "TEXT", + bytes.__class__: "BLOB", + bytes: "BLOB", + datetime.datetime: "TEXT", + datetime.date: "TEXT", + datetime.time: "TEXT", + None.__class__: "TEXT", +} +# If numpy is available, add more types +if np: + COLUMN_TYPE_MAPPING.update( + { + np.int8: "INTEGER", + np.int16: "INTEGER", + np.int32: "INTEGER", + np.int64: "INTEGER", + np.uint8: "INTEGER", + np.uint16: "INTEGER", + np.uint32: "INTEGER", + np.uint64: "INTEGER", + np.float16: "FLOAT", + np.float32: "FLOAT", + np.float64: "FLOAT", + } + ) + class Database: def __init__(self, filename_or_conn): @@ -67,39 +97,10 @@ class Database: column_items.insert(0, (hash_id, str)) pk = hash_id extra = "" - col_type_mapping = { - float: "FLOAT", - int: "INTEGER", - bool: "INTEGER", - str: "TEXT", - bytes.__class__: "BLOB", - bytes: "BLOB", - datetime.datetime: "TEXT", - datetime.date: "TEXT", - datetime.time: "TEXT", - None.__class__: "TEXT", - } - # If numpy is available, add more types - if np: - col_type_mapping.update( - { - np.int8: "INTEGER", - np.int16: "INTEGER", - np.int32: "INTEGER", - np.int64: "INTEGER", - np.uint8: "INTEGER", - np.uint16: "INTEGER", - np.uint32: "INTEGER", - np.uint64: "INTEGER", - np.float16: "FLOAT", - np.float32: "FLOAT", - np.float64: "FLOAT", - } - ) columns_sql = ",\n".join( " [{col_name}] {col_type}{primary_key}{references}".format( col_name=col_name, - col_type=col_type_mapping[col_type], + col_type=COLUMN_TYPE_MAPPING[col_type], primary_key=" PRIMARY KEY" if (pk == col_name) else "", references=( " REFERENCES [{other_table}]([{other_column}])".format( @@ -254,6 +255,15 @@ class Table: self.db.conn.execute(sql) return self + def add_column(self, col_name, col_type): + sql = "ALTER TABLE [{table}] ADD COLUMN [{col_name}] {col_type};".format( + table=self.name, + col_name=col_name, + col_type=COLUMN_TYPE_MAPPING[col_type], + ) + self.db.conn.execute(sql) + return self + def drop(self): return self.db.conn.execute("DROP TABLE {}".format(self.name)) diff --git a/tests/test_create.py b/tests/test_create.py index 4c7b550..6528a23 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -4,6 +4,7 @@ import datetime import json import pathlib import pytest +import re import sqlite3 try: @@ -122,6 +123,32 @@ def test_create_table_works_for_m2m_with_only_foreign_keys(fresh_db): ) +@pytest.mark.parametrize( + "col_name,col_type,expected_schema", + ( + ( + "nickname", str, "CREATE TABLE [dogs] ( [name] TEXT , [nickname] TEXT)" + ), + ( + "dob", datetime.date, "CREATE TABLE [dogs] ( [name] TEXT , [dob] TEXT)" + ), + ( + "age", int, "CREATE TABLE [dogs] ( [name] TEXT , [age] INTEGER)" + ), + ( + "weight", float, "CREATE TABLE [dogs] ( [name] TEXT , [weight] FLOAT)" + ), + ) +) +def test_add_column(fresh_db, col_name, col_type, expected_schema): + fresh_db.create_table("dogs", {"name": str}) + assert "CREATE TABLE [dogs] ( [name] TEXT )" == collapse_whitespace( + fresh_db["dogs"].schema + ) + fresh_db["dogs"].add_column(col_name, col_type) + assert expected_schema == collapse_whitespace(fresh_db["dogs"].schema) + + @pytest.mark.parametrize( "columns,index_name,expected_index", ( @@ -337,3 +364,10 @@ def test_create_table_numpy(fresh_db): "np.uint8": 8, } ] == list(fresh_db["types"].rows) + + +COLLAPSE_RE = re.compile("\s+") + + +def collapse_whitespace(s): + return COLLAPSE_RE.sub(" ", s)