diff --git a/sqlite_utils/db.py b/sqlite_utils/db.py index 32587ce..4e7fbf1 100644 --- a/sqlite_utils/db.py +++ b/sqlite_utils/db.py @@ -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) diff --git a/tests/test_create.py b/tests/test_create.py index 5b95eec..3ab62a1 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -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(