Merge pull request google#41 from google/force-index

dcbrandao · web-flow · commit 4cdc8aa90ad3 · 2019-03-08T15:53:47.000-05:00
Add force_index condition
diff --git a/spanner_orm/__init__.py b/spanner_orm/__init__.py
@@ -44,6 +44,7 @@
 Timestamp = field.Timestamp
 
 equal_to = condition.equal_to
+force_index = condition.force_index
 greater_than = condition.greater_than
 greater_than_or_equal_to = condition.greater_than_or_equal_to
 includes = condition.includes
diff --git a/spanner_orm/admin/metadata.py b/spanner_orm/admin/metadata.py
@@ -113,11 +113,12 @@ def indexes(cls):
     indexes = collections.defaultdict(dict)
     for schema in index_schemas:
       key = (schema.table_name, schema.index_name)
-      indexes[schema.table_name][schema.index_name] = index.Index(
+      new_index = index.Index(
           index_columns[key],
-          schema.index_name,
           parent=schema.parent_table_name,
           null_filtered=schema.is_null_filtered,
           unique=schema.is_unique,
           storing_columns=storing_columns[key])
+      new_index.name = schema.index_name
+      indexes[schema.table_name][schema.index_name] = new_index
     return indexes
diff --git a/spanner_orm/admin/update.py b/spanner_orm/admin/update.py
@@ -289,9 +289,10 @@ def validate(self):
     if not self._model:
       raise error.SpannerError('Table {} does not exist'.format(self._table))
 
-    if self._index not in self._model.indexes:
+    db_index = self._model.indexes.get(self._index)
+    if not db_index:
       raise error.SpannerError('Index {} does not exist'.format(self._index))
-    if self._index == index.Index.PRIMARY_INDEX:
+    if db_index.primary_index:
       raise error.SpannerError('Index {} is the primary index'.format(
           self._index))
 
diff --git a/spanner_orm/condition.py b/spanner_orm/condition.py
@@ -24,10 +24,11 @@
 
 class Segment(enum.Enum):
   """The segment of the SQL query that a Condition belongs to."""
-  WHERE = 1
-  ORDER_BY = 2
-  LIMIT = 3
-  JOIN = 4
+  FROM = 1
+  JOIN = 2
+  WHERE = 3
+  ORDER_BY = 4
+  LIMIT = 5
 
 
 class Condition(abc.ABC):
@@ -118,6 +119,35 @@ def _validate(self, model):
     assert (origin.field_type() == dest.field_type() and
             origin.nullable() == dest.nullable())
 
+class ForceIndexCondition(Condition):
+  """Used to indicate which index should be used in a Spanner query."""
+
+  def __init__(self, name):
+    super().__init__()
+    self.name = name
+    self.index = None
+
+  def bind(self, model):
+    super().bind(model)
+    self.index = self.model.indexes[self.name]
+
+  def _params(self):
+    return {}
+
+  @staticmethod
+  def segment():
+    return Segment.FROM
+
+  def _sql(self):
+    return '@{{FORCE_INDEX={}}}'.format(self.index.name)
+
+  def _types(self):
+    return {}
+
+  def _validate(self, model):
+    assert self.name in model.indexes
+    assert not model.indexes[self.name].primary
+
 
 class IncludesCondition(Condition):
   """Used to include related models via a relation in a Spanner query."""
@@ -439,6 +469,10 @@ def equal_to(column, value):
   return EqualityCondition(column, value)
 
 
+def force_index(index):
+  return ForceIndexCondition(index)
+
+
 def greater_than(column, value):
   return ComparisonCondition('>', column, value)
 
diff --git a/spanner_orm/index.py b/spanner_orm/index.py
@@ -20,15 +20,18 @@ class Index(object):
 
   def __init__(self,
                columns,
-               name,
                parent=None,
                null_filtered=False,
                unique=False,
                storing_columns=None):
     assert len(columns) > 0, 'An index must have at least one column'
     self.columns = columns
-    self.name = name
+    self.name = None
     self.parent = parent
     self.null_filtered = null_filtered
     self.unique = unique
     self.storing_columns = storing_columns or []
+
+  @property
+  def primary(self):
+    return self.name == self.PRIMARY_INDEX
diff --git a/spanner_orm/model.py b/spanner_orm/model.py
@@ -60,8 +60,9 @@ def finalize(self):
 
     if index.Index.PRIMARY_INDEX not in self.indexes:
       primary_keys = [f.name for f in sorted_fields if f.primary_key()]
-      primary_index = index.Index(primary_keys, index.Index.PRIMARY_INDEX)
-      self.indexes[primary_index.name] = primary_index
+      primary_index = index.Index(primary_keys)
+      primary_index.name = index.Index.PRIMARY_INDEX
+      self.indexes[index.Index.PRIMARY_INDEX] = primary_index
     self.primary_keys = self.indexes[index.Index.PRIMARY_INDEX].columns
 
     self.columns = [f.name for f in sorted_fields]
@@ -87,7 +88,7 @@ def add_relation(self, name, new_relation):
 
   def add_index(self, name, new_index):
     new_index.name = name
-    self.relations[name] = new_index
+    self.indexes[name] = new_index
 
 
 class ModelBase(type):
diff --git a/spanner_orm/query.py b/spanner_orm/query.py
@@ -77,7 +77,17 @@ def _select(self):
 
   def _from(self):
     """Processes the FROM segment of the SQL query."""
-    return (' FROM {}'.format(self._model.table), {}, {})
+    froms = self._segments(condition.Segment.FROM)
+    index_sql = ''
+    if froms:
+      if len(froms) != 1:
+        raise error.SpannerError('Only one index can be specified')
+      force_index = froms[0]
+      index_sql = force_index.sql()
+
+    sql = ' FROM {}{}'.format(self._model.table, index_sql)
+
+    return (sql, {}, {})
 
   def _where(self):
     """Processes the WHERE segment of the SQL query."""
diff --git a/spanner_orm/tests/models.py b/spanner_orm/tests/models.py
@@ -15,6 +15,7 @@
 """Models used by unit tests."""
 
 from spanner_orm import field
+from spanner_orm import index
 from spanner_orm import model
 from spanner_orm import relationship
 
@@ -57,3 +58,5 @@ class UnittestModel(model.Model):
   string_2 = field.Field(field.String, nullable=True)
   timestamp = field.Field(field.Timestamp)
   string_array = field.Field(field.StringArray, nullable=True)
+
+  test_index = index.Index(['string_2'])
diff --git a/spanner_orm/tests/query_test.py b/spanner_orm/tests/query_test.py
@@ -149,6 +149,11 @@ def test_only_one_limit_allowed(self):
     with self.assertRaises(error.SpannerError):
       self.select(condition.limit(2), condition.limit(2))
 
+  def test_force_index(self):
+    select_query = self.select(condition.force_index('test_index'))
+    expected_sql = 'FROM table@{FORCE_INDEX=test_index}'
+    self.assertEndsWith(select_query.sql(), expected_sql)
+
   def includes(self, relation, *conditions):
     include_condition = condition.includes(relation, list(conditions))
     return query.SelectQuery(models.ChildTestModel, [include_condition])
