Merge pull request #117 from google/create-index-options

dseomn · web-flow · commit 3bbac233e7fe · 2020-12-23T19:40:08.000-05:00
Add support for NULL_FILTERED and UNIQUE in CreateIndex.
diff --git a/spanner_orm/admin/update.py b/spanner_orm/admin/update.py
@@ -240,16 +240,25 @@ def __init__(self,
                index_name: str,
                columns: Iterable[str],
                interleaved: Optional[str] = None,
+               null_filtered: bool = False,
+               unique: bool = False,
                storing_columns: Optional[Iterable[str]] = None):
     self._table = table_name
     self._index = index_name
     self._columns = columns
     self._parent_table = interleaved
+    self._null_filtered = null_filtered
+    self._unique = unique
     self._storing_columns = storing_columns or []
 
   def ddl(self) -> str:
-    statement = 'CREATE INDEX {} ON {} ({})'.format(self._index, self._table,
-                                                    ', '.join(self._columns))
+    statement = 'CREATE'
+    if self._unique:
+      statement += ' UNIQUE'
+    if self._null_filtered:
+      statement += ' NULL_FILTERED'
+    statement += (f' INDEX {self._index} '
+                  f'ON {self._table} ({", ".join(self._columns)})')
     if self._storing_columns:
       statement += 'STORING ({})'.format(', '.join(self._storing_columns))
     if self._parent_table:
diff --git a/spanner_orm/tests/update_test.py b/spanner_orm/tests/update_test.py
@@ -16,13 +16,15 @@
 import unittest
 from unittest import mock
 
+from absl.testing import parameterized
+
 from spanner_orm import error
 from spanner_orm import field
 from spanner_orm.admin import update
 from spanner_orm.tests import models
 
 
-class UpdateTest(unittest.TestCase):
+class UpdateTest(parameterized.TestCase):
 
   @mock.patch('spanner_orm.admin.metadata.SpannerMetadata.model')
   def test_add_column(self, get_model):
@@ -118,15 +120,35 @@ def test_drop_table(self, get_model, tables, indexes):
     test_update.validate()
     self.assertEqual(test_update.ddl(), 'DROP TABLE {}'.format(table_name))
 
+  @parameterized.named_parameters(
+      (
+          'basic',
+          update.CreateIndex(
+              table_name=models.SmallTestModel.table,
+              index_name='foo',
+              columns=['value_1'],
+          ),
+          f'CREATE INDEX foo ON {models.SmallTestModel.table} (value_1)',
+      ),
+      (
+          'with_options',
+          update.CreateIndex(
+              table_name=models.SmallTestModel.table,
+              index_name='foo',
+              columns=['value_1'],
+              null_filtered=True,
+              unique=True,
+          ),
+          (f'CREATE UNIQUE NULL_FILTERED INDEX foo '
+           f'ON {models.SmallTestModel.table} (value_1)'),
+      ),
+  )
   @mock.patch('spanner_orm.admin.metadata.SpannerMetadata.model')
-  def test_add_index(self, get_model):
-    table_name = models.SmallTestModel.table
+  def test_add_index(self, test_update, expected_ddl, get_model):
     get_model.return_value = models.SmallTestModel
 
-    test_update = update.CreateIndex(table_name, 'foo', ['value_1'])
     test_update.validate()
-    self.assertEqual(test_update.ddl(),
-                     'CREATE INDEX foo ON {} (value_1)'.format(table_name))
+    self.assertEqual(test_update.ddl(), expected_ddl)
 
 
 if __name__ == '__main__':
