table.add_column(name, type) method, refs #15

This commit is contained in:
Simon Willison 2019-02-24 11:40:26 -08:00
commit dc2560879e
3 changed files with 108 additions and 30 deletions

View file

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

View file

@ -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))

View file

@ -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)