diff --git a/spanner_orm/field.py b/spanner_orm/field.py index 64f40c6..733130a 100644 --- a/spanner_orm/field.py +++ b/spanner_orm/field.py @@ -48,16 +48,27 @@ class Field(object): def __init__(self, field_type: Type[FieldType], nullable: bool = False, - primary_key: bool = False): + primary_key: bool = False, + allow_commit_timestamp: bool = False): self.name = None self._type = field_type self._nullable = nullable self._primary_key = primary_key + self._allow_commit_timestamp = allow_commit_timestamp + + if self._type.ddl() != "TIMESTAMP" and self._allow_commit_timestamp: + raise error.ValidationError('allow_commit_timestamp can not be set on field {}'.format(self._type)) def ddl(self) -> str: if self._nullable: - return self._type.ddl() - return '{field_type} NOT NULL'.format(field_type=self._type.ddl()) + nullable = '' + else: + nullable = ' NOT NULL' + if self._allow_commit_timestamp: + options = ' OPTIONS (allow_commit_timestamp=true)' + else: + options = '' + return '{field_type}{nullable}{options}'.format(field_type=self._type.ddl(), nullable=nullable, options=options) def field_type(self) -> Type[FieldType]: return self._type diff --git a/spanner_orm/tests/models.py b/spanner_orm/tests/models.py index 2a8218a..19f8ef4 100644 --- a/spanner_orm/tests/models.py +++ b/spanner_orm/tests/models.py @@ -76,6 +76,7 @@ class UnittestModel(model.Model): string = field.Field(field.String, primary_key=True) string_2 = field.Field(field.String, nullable=True) timestamp = field.Field(field.Timestamp) + timestamp_2 = field.Field(field.Timestamp, nullable=True, allow_commit_timestamp=True) string_array = field.Field(field.StringArray, nullable=True) test_index = index.Index(['string_2']) diff --git a/spanner_orm/tests/update_test.py b/spanner_orm/tests/update_test.py index 7b55965..f12bae5 100644 --- a/spanner_orm/tests/update_test.py +++ b/spanner_orm/tests/update_test.py @@ -67,8 +67,9 @@ def test_create_table(self, get_model): test_model_ddl = ('CREATE TABLE table (int_ INT64 NOT NULL, int_2 INT64,' ' string STRING(MAX) NOT NULL, string_2 STRING(MAX),' - ' timestamp TIMESTAMP NOT NULL, string_array' - ' ARRAY) PRIMARY KEY (int_, string)') + ' timestamp TIMESTAMP NOT NULL,' + ' timestamp_2 TIMESTAMP OPTIONS (allow_commit_timestamp=true),' + ' string_array ARRAY) PRIMARY KEY (int_, string)') self.assertEqual(test_update.ddl(), test_model_ddl) @mock.patch('spanner_orm.admin.metadata.SpannerMetadata.model')