diff --git a/spanner_orm/admin/update.py b/spanner_orm/admin/update.py index ad0c961..c8e91d6 100644 --- a/spanner_orm/admin/update.py +++ b/spanner_orm/admin/update.py @@ -334,7 +334,29 @@ def validate(self) -> None: if db_index.primary: raise error.SpannerError('Index {} is the primary index'.format( self._index)) - + +class AddForeignKey(SchemaUpdate): + """...""" + + def __init__(self, + table_name: str, + constraint_name: str, + foreign_key_relationship: foreign_key_relationship.ForeignKeyRelationship): + self._table = table_name + # self._constraint_name = constraint_name + self._foreign_key_relationship = foreign_key_relationship + self._foreign_key_relationship.name = constraint_name + + def ddl(self) -> str: + return 'ALTER TABLE {} ADD {}'.format( + self._table, + self._foreign_key_relationship.ddl) + + def validate(self) -> None: + model_ = metadata.SpannerMetadata.model(self._table) + if not model_: + raise error.SpannerError('Table {} does not exist'.format(self._table)) + # TODO: add more validations class NoUpdate(SchemaUpdate): """Update that does nothing, for migrations that don't update db schemas.""" diff --git a/spanner_orm/model.py b/spanner_orm/model.py index c235aa0..3f074cf 100644 --- a/spanner_orm/model.py +++ b/spanner_orm/model.py @@ -603,7 +603,6 @@ def _execute_write( else: return cls.spanner_api().run_write(db_api, *args) - def __setattr__(self, name: str, value: Any) -> None: if name in self._relations: raise AttributeError(name) diff --git a/spanner_orm/tests/update_test.py b/spanner_orm/tests/update_test.py index 4bb4ac1..0b73fd0 100644 --- a/spanner_orm/tests/update_test.py +++ b/spanner_orm/tests/update_test.py @@ -20,6 +20,7 @@ from spanner_orm import error from spanner_orm import field +from spanner_orm import foreign_key_relationship from spanner_orm.admin import update from spanner_orm.tests import models @@ -175,6 +176,30 @@ def test_add_index(self, test_update, expected_ddl, get_model): test_update.validate() self.assertEqual(test_update.ddl(), expected_ddl) + @mock.patch('spanner_orm.admin.metadata.SpannerMetadata.model') + def test_add_foreign_key(self, get_model): + table_name = models.SmallTestModel.table + get_model.return_value = models.SmallTestModel + + # new_field = field.Field(field.String, nullable=True) + new_foreign_key = foreign_key_relationship.ForeignKeyRelationship( + 'SmallTestModel', {'referencing_key_1': 'key'}) + test_update = update.AddForeignKey( + table_name, 'constraint_name', new_foreign_key) + test_update.validate() + # ALTER TABLE Orders + # ADD CONSTRAINT FK_ProductOrder FOREIGN KEY (ProductID) REFERENCES Products (ProductID); + self.assertEqual( + test_update.ddl(), + 'ALTER TABLE {} ADD CONSTRAINT {} FOREIGN KEY ({}) ' + 'REFERENCES {} ({})'.format( + table_name, + 'constraint_name', + 'referencing_key_1',#referencing_column, + 'SmallTestModel',#referenced_table_name, + 'key',#referenced_column, + )) + if __name__ == '__main__': logging.basicConfig()