Add functionality to add/remove foreign key from existing table

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."""

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)

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