MetaOfX
diff --git a/‎spanner_orm/__init__.py‎
Lines changed: 1 addition & 0 deletions b/‎spanner_orm/__init__.py‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎spanner_orm/condition.py‎
Lines changed: 74 additions & 12 deletions b/‎spanner_orm/condition.py‎
Lines changed: 74 additions & 12 deletions
diff --git a/‎spanner_orm/tests/condition_test.py‎
Lines changed: 10 additions & 0 deletions b/‎spanner_orm/tests/condition_test.py‎
Lines changed: 10 additions & 0 deletions
diff --git a/‎spanner_orm/tests/migrations_for_emulator_test/create_null_filtered_index_model_760ec5fae5da.py‎
Lines changed: 40 additions & 0 deletions b/‎spanner_orm/tests/migrations_for_emulator_test/create_null_filtered_index_model_760ec5fae5da.py‎
Lines changed: 40 additions & 0 deletions
diff --git a/‎spanner_orm/tests/migrations_for_emulator_test/create_null_filtered_index_model_value_index_69a8f072dacf.py‎
Lines changed: 28 additions & 0 deletions b/‎spanner_orm/tests/migrations_for_emulator_test/create_null_filtered_index_model_value_index_69a8f072dacf.py‎
Lines changed: 28 additions & 0 deletions
diff --git a/‎spanner_orm/tests/models.py‎
Lines changed: 10 additions & 0 deletions b/‎spanner_orm/tests/models.py‎
Lines changed: 10 additions & 0 deletions
@@ -81,6 +81,7 @@
 contains = condition.contains
 equal_to = condition.equal_to
 force_index = condition.force_index
+force_null_filtered_index = condition.force_null_filtered_index
 greater_than = condition.greater_than
 greater_than_or_equal_to = condition.greater_than_or_equal_to
 in_list = condition.in_list
 
@@ -21,7 +21,7 @@
 import decimal
 import enum
 import string
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple, Type, TypeVar, Union
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union
 
 from spanner_orm import error
 from spanner_orm import field
@@ -364,10 +364,10 @@ def _validate(self, model_class: Type[Any]) -> None:
           origin.name, dest.name))
 
 
-class ForceIndexCondition(Condition):
-  """Used to indicate which index should be used in a Spanner query."""
+class _IndexCondition(Condition):
+  """Base class for conditions based on an Index."""
 
-  def __init__(self, index_or_name: Union[Type[index.Index], str]):
+  def __init__(self, index_or_name: Union[index.Index, str]):
     super().__init__()
     if isinstance(index_or_name, index.Index):
       self.name = index_or_name.name
@@ -380,30 +380,63 @@ def bind(self, model_class: Type[Any]) -> None:
     super().bind(model_class)
     self.index = self.model_class.indexes[self.name]
 
+  def _validate(self, model_class: Type[Any]) -> None:
+    if self.name not in model_class.indexes:
+      raise error.ValidationError('{} is not an index on {}'.format(
+          self.name, model_class.table))
+    if self.index and self.index != model_class.indexes[self.name]:
+      raise error.ValidationError('{} does not belong to {}'.format(
+          self.index.name, model_class.table))
+
+
+class ForceIndexCondition(_IndexCondition):
+  """Used to indicate which index should be used in a Spanner query."""
+
+  def __init__(
+      self,
+      index_or_name: Union[index.Index, str],
+      *,
+      extra_hints: Sequence[str] = (),
+  ):
+    super().__init__(index_or_name)
+    self._extra_hints = extra_hints
+
   def _params(self) -> Dict[str, Any]:
     return {}
 
   def segment(self) -> Segment:
     return Segment.FROM
 
   def _sql(self) -> str:
-    return '@{{FORCE_INDEX={}}}'.format(self.name)
+    hints = (f'FORCE_INDEX={self.name}', *self._extra_hints)
+    return f'@{{{",".join(hints)}}}'
 
   def _types(self) -> Dict[str, type_pb2.Type]:
     return {}
 
   def _validate(self, model_class: Type[Any]) -> None:
-    if self.name not in model_class.indexes:
-      raise error.ValidationError('{} is not an index on {}'.format(
-          self.name, model_class.table))
-    if self.index and self.index != model_class.indexes[self.name]:
-      raise error.ValidationError('{} does not belong to {}'.format(
-          self.index.name, model_class.table))
-
+    super()._validate(model_class)
     if model_class.indexes[self.name].primary:
       raise error.ValidationError('Cannot force query using primary index')
 
 
+class _IndexIgnoreNullsCondition(_IndexCondition):
+  """Condition to filter NULL values in any column of an index."""
+
+  def _params(self) -> Dict[str, Any]:
+    return {}
+
+  def segment(self) -> Segment:
+    return Segment.WHERE
+
+  def _sql(self) -> str:
+    return '({})'.format(' AND '.join(
+        f'{column} IS NOT NULL' for column in self.index.columns))
+
+  def _types(self) -> Dict[str, type_pb2.Type]:
+    return {}
+
+
 class IncludesCondition(Condition):
   """Used to include related model_classs via a relation in a Spanner query."""
 
@@ -888,6 +921,35 @@ def force_index(forced_index: Union[index.Index, str]) -> ForceIndexCondition:
   return ForceIndexCondition(forced_index)
 
 
+def force_null_filtered_index(
+    forced_index: Union[index.Index, str]) -> Sequence[Condition]:
+  """Returns conditions to force the query to use the given NULL_FILTERED index.
+
+  In Cloud Spanner, a query against a NULL_FILTERED index is tested to see if it
+  can use safely use that index. If using the index would result in incorrect
+  results (e.g., by ignoring NULL values that would be in the same query without
+  using the index), it's an error. However, the Cloud Spanner Emulator
+  doesn't support that check:
+  https://github.com/GoogleCloudPlatform/cloud-spanner-emulator/blob/e887ff5569684e6e45ce7c90d0fdfb7b1faa1491/common/errors.cc#L1790-L1800
+
+  For queries that can safely ignore any  NULL values covered by the index, this
+  function returns conditions that both filter out all relevant NULLs (avoiding
+  the potential error in Cloud Spanner) and disable the check in Cloud Spanner
+  Emulator.
+
+  Args:
+    forced_index: NULL_FILTERED index to use.
+  """
+  return (
+      ForceIndexCondition(
+          forced_index,
+          extra_hints=(
+              'spanner_emulator.disable_query_null_filtered_index_check=true',
+          )),
+      _IndexIgnoreNullsCondition(forced_index),
+  )
+
+
 def greater_than(column: Union[field.Field, str],
                  value: Any) -> ComparisonCondition:
   """Condition where the specified column is greater than the given value.
 
@@ -333,6 +333,16 @@ def test_contains(
             )),
     )
 
+  def test_force_null_filtered_index(self):
+    non_null_model = models.NullFilteredIndexModel(
+        dict(key='a', value_1='a', value_2=1))
+    non_null_model.save()
+    models.NullFilteredIndexModel(dict(key='b', value_1=None, value_2=2)).save()
+    self.assertCountEqual((non_null_model,),
+                          models.NullFilteredIndexModel.where(
+                              *spanner_orm.force_null_filtered_index(
+                                  models.NullFilteredIndexModel.value_index)))
+
 
 if __name__ == '__main__':
   logging.basicConfig()
 
@@ -0,0 +1,40 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Spanner ORM migration: create_null_filtered_index_model.
+
+Migration ID: '760ec5fae5da'
+Created: 2022-03-01 16:50:32-05:00
+"""
+
+import spanner_orm
+
+migration_id = '760ec5fae5da'
+prev_migration_id = 'f735d6b706d4'
+
+
+class _NullFilteredIndexModel(spanner_orm.Model):
+  __table__ = 'NullFilteredIndexModel'
+  key = spanner_orm.Field(spanner_orm.String, primary_key=True)
+  value_1 = spanner_orm.Field(spanner_orm.String, nullable=True)
+  value_2 = spanner_orm.Field(spanner_orm.Integer)
+
+
+def upgrade() -> spanner_orm.MigrationUpdate:
+  """See spanner_orm migrations interface."""
+  return spanner_orm.CreateTable(_NullFilteredIndexModel)
+
+
+def downgrade() -> spanner_orm.MigrationUpdate:
+  """See spanner_orm migrations interface."""
+  return spanner_orm.DropTable(_NullFilteredIndexModel.__table__)
@@ -0,0 +1,28 @@
+"""Spanner ORM migration: create_null_filtered_index_model_value_index.
+
+Migration ID: '69a8f072dacf'
+Created: 2022-03-01 16:53:59-05:00
+"""
+
+import spanner_orm
+
+migration_id = '69a8f072dacf'
+prev_migration_id = '760ec5fae5da'
+
+
+def upgrade() -> spanner_orm.MigrationUpdate:
+  """See spanner_orm migrations interface."""
+  return spanner_orm.CreateIndex(
+      table_name='NullFilteredIndexModel',
+      index_name='value_index',
+      columns=['value_1', 'value_2'],
+      null_filtered=True,
+  )
+
+
+def downgrade() -> spanner_orm.MigrationUpdate:
+  """See spanner_orm migrations interface."""
+  return spanner_orm.DropIndex(
+      table_name='NullFilteredIndexModel',
+      index_name='value_index',
+  )
@@ -130,3 +130,13 @@ class UnittestModelWithoutSecondaryIndexes(model.Model):
   bytes_2 = field.Field(field.BytesBase64, nullable=True)
   timestamp = field.Field(field.Timestamp)
   string_array = field.Field(field.StringArray, nullable=True)
+
+
+class NullFilteredIndexModel(model.Model):
+  """Model class for testing NULL_FILTERED indexes."""
+
+  __table__ = 'NullFilteredIndexModel'
+  key = field.Field(field.String, primary_key=True)
+  value_1 = field.Field(field.String, nullable=True)
+  value_2 = field.Field(field.Integer)
+  value_index = index.Index(['value_1', 'value_2'], null_filtered=True)