Added table.foreign_keys property, fixed bug in foreign key creation

This commit is contained in:
Simon Willison 2018-07-28 15:41:18 -07:00
commit e51f36d3eb
2 changed files with 39 additions and 1 deletions

View file

@ -5,6 +5,9 @@ import json
Column = namedtuple(
"Column", ("cid", "name", "type", "notnull", "default_value", "is_pk")
)
ForeignKey = namedtuple(
"ForeignKey", ("table", "column", "other_table", "other_column")
)
class Database:
@ -42,7 +45,7 @@ class Database:
}[col_type],
primary_key=" PRIMARY KEY" if (pk == col_name) else "",
references=(
" REFERENCES [{other_table}({other_column})]".format(
" REFERENCES [{other_table}]([{other_column}])".format(
other_table=foreign_keys_by_name[col_name][2],
other_column=foreign_keys_by_name[col_name][3],
)
@ -77,6 +80,24 @@ class Table:
).fetchall()
return [Column(*row) for row in rows]
@property
def foreign_keys(self):
fks = []
for row in self.db.conn.execute(
"PRAGMA foreign_key_list([{}])".format(self.name)
).fetchall():
if row is not None:
id, seq, table_name, from_, to_, on_update, on_delete, match = row
fks.append(
ForeignKey(
table=self.name,
column=from_,
other_table=table_name,
other_column=to_,
)
)
return fks
def create(self, columns, pk=None, foreign_keys=None):
columns = {name: value for (name, value) in columns.items()}
self.db.create_table(self.name, columns, pk=pk, foreign_keys=foreign_keys)

View file

@ -55,6 +55,23 @@ def test_create_table_works_for_m2m_with_only_foreign_keys(fresh_db):
{"name": "one_id", "type": "INTEGER"},
{"name": "two_id", "type": "INTEGER"},
] == [{"name": col.name, "type": col.type} for col in fresh_db["m2m"].columns]
assert sorted(
[
{"column": "one_id", "other_table": "one", "other_column": "id"},
{"column": "two_id", "other_table": "two", "other_column": "id"},
],
key=lambda s: repr(s),
) == sorted(
[
{
"column": fk.column,
"other_table": fk.other_table,
"other_column": fk.other_column,
}
for fk in fresh_db["m2m"].foreign_keys
],
key=lambda s: repr(s),
)
@pytest.mark.parametrize(