From 82f3f1bc04da52f5e08d839b357a91272132c3aa Mon Sep 17 00:00:00 2001
From: Daniel Gorelik <dgorelik@google.com>
Date: Sun, 17 Jan 2021 11:05:49 -0500
Subject: [PATCH] Add method to add foreign key

---
 spanner_orm/admin/update.py      | 24 +++++++++++++++++++++++-
 spanner_orm/tests/update_test.py | 25 +++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/spanner_orm/admin/update.py b/spanner_orm/admin/update.py
index 54f7ec8..6f78a57 100644
--- a/spanner_orm/admin/update.py
+++ b/spanner_orm/admin/update.py
@@ -318,7 +318,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/tests/update_test.py b/spanner_orm/tests/update_test.py
index a96190f..f348b82 100644
--- a/spanner_orm/tests/update_test.py
+++ b/spanner_orm/tests/update_test.py
@@ -18,6 +18,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
 
@@ -142,6 +143,30 @@ def test_add_index(self, get_model):
     self.assertEqual(test_update.ddl(),
                      'CREATE INDEX foo ON {} (value_1)'.format(table_name))
 
+  @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()