class IndexTestModel(model.Model):
__table__ = 'IndexTestModel'
key = field.Field(field.String, primary_key=True)
value = field.Field(field.String)
value_index = index.Index(['value'])
class OriginalSmallTestModelTable(spanner_orm.model.Model): """ORM Model with the original schema for the SmallTestModel table.""" __table__ = 'SmallTestModel' key = field.Field(field.String, primary_key=True) value_1 = field.Field(field.String) value_2 = field.Field(field.String, nullable=True)
class ChildTestModel(model.Model):
"""Model class for testing interleaved tables."""
__table__ = 'ChildTestModel'
__interleaved__ = 'SmallTestModel'
key = field.Field(field.String, primary_key=True)
child_key = field.Field(field.String, primary_key=True)
class SmallTestModel(model.Model):
"""Model class used for testing."""
__table__ = 'SmallTestModel'
key = field.Field(field.String, primary_key=True)
value_1 = field.Field(field.String)
value_2 = field.Field(field.String, nullable=True)
index_1 = index.Index(['value_1'])
class MigrationStatus(model.Model):
@classmethod
def spanner_api(cls) -> api.SpannerAdminApi:
return api.spanner_admin_api()
__table__ = 'spanner_orm_migrations'
id = typing.cast(str, field.Field(field.String, primary_key=True))
migrated = typing.cast(bool, field.Field(field.Boolean))
update_time = typing.cast(datetime.datetime, field.Field(field.Timestamp))
class MigrationStatus(model.Model):
@classmethod
def spanner_api(cls) -> api.SpannerAdminApi:
return api.spanner_admin_api()
__table__ = "spanner_orm_migrations"
id = field.Field(field.String, primary_key=True)
migrated = field.Field(field.Boolean)
update_time = field.Field(field.Timestamp)
class TableSchema(schema.InformationSchema):
"""Model for interacting with Spanner column schema table."""
__table__ = "information_schema.tables"
table_catalog = field.Field(field.String, primary_key=True)
table_schema = field.Field(field.String, primary_key=True)
table_name = field.Field(field.String, primary_key=True)
parent_table_name = field.Field(field.String, nullable=True)
on_delete_action = field.Field(field.String, nullable=True)
class IndexTestModel(model.Model):
__table__ = "IndexTestModel"
key = field.Field(field.String, primary_key=True)
value = field.Field(field.String)
value_idx = index.Index(["value"], name="value")
value_idx2 = index.Index(
["value"], name="value_desc", column_ordering={"value": False}
)
class RelationshipTestModel(model.Model):
"""Model class for testing relationships."""
__table__ = 'RelationshipTestModel'
parent_key = field.Field(field.String, primary_key=True)
child_key = field.Field(field.String, primary_key=True)
parent = relationship.Relationship(
'spanner_orm.tests.models.SmallTestModel', {'parent_key': 'key'},
single=True)
parents = relationship.Relationship(
'spanner_orm.tests.models.SmallTestModel', {'parent_key': 'key'})
class SmallTestParentModel(model.Model):
"""Model class used for testing."""
__table__ = "SmallTestParentModel"
key = field.Field(field.String, primary_key=True)
value_1 = field.Field(field.String)
value_2 = field.Field(field.String, nullable=True)
index_1 = index.Index(["value_1"])
children = relationship.Relationship(
"spanner_orm.tests.models.ChildTestModel", {"key": "key"}
)
class TableSchema(schema.InformationSchema):
"""Model for interacting with Spanner column schema table."""
__table__ = 'information_schema.tables'
table_catalog = typing.cast(str, field.Field(field.String,
primary_key=True))
table_schema = typing.cast(str, field.Field(field.String,
primary_key=True))
table_name = typing.cast(str, field.Field(field.String, primary_key=True))
parent_table_name = typing.cast(Optional[str],
field.Field(field.String, nullable=True))
on_delete_action = typing.cast(Optional[str],
field.Field(field.String, nullable=True))
def tables(cls) -> Dict[str, Dict[str, Any]]:
"""Compiles table information from column schema."""
column_data = collections.defaultdict(dict)
columns = column.ColumnSchema.where(
None,
condition.equal_to("table_catalog", ""),
condition.equal_to("table_schema", ""),
)
for column_row in columns:
new_field = field.Field(
column_row.field_type,
nullable=column_row.nullable,
size=column_row.size,
)
new_field.name = column_row.column_name
new_field.position = column_row.ordinal_position
column_data[column_row.table_name][column_row.column_name] = new_field
table_data = collections.defaultdict(dict)
tables = table.TableSchema.where(
None,
condition.equal_to("table_catalog", ""),
condition.equal_to("table_schema", ""),
)
for table_row in tables:
name = table_row.table_name
table_data[name]["parent_table"] = table_row.parent_table_name
table_data[name]["fields"] = column_data[name]
return table_data
class RelationshipTestModel(model.Model):
"""Model class for testing relationships."""
__table__ = "RelationshipTestModel"
parent_key = field.Field(field.String, primary_key=True)
child_key = field.Field(field.String, primary_key=True)
parent = relationship.Relationship(
"spanner_orm.tests.models.SmallTestModel", {"parent_key": "key"}, single=True
)
parents = relationship.Relationship(
"spanner_orm.tests.models.SmallTestModel", {"parent_key": "key"}
)
fk_multicolumn = relationship.Relationship(
"spanner_orm.tests.models.SmallTestModel",
{"parent_key": "key", "parent_key2": "key2"},
)
class IndexColumnSchema(schema.InformationSchema):
"""Model for interacting with Spanner index column schema table."""
__table__ = "information_schema.index_columns"
table_catalog = field.Field(field.String, primary_key=True)
table_schema = field.Field(field.String, primary_key=True)
table_name = field.Field(field.String, primary_key=True)
index_name = field.Field(field.String, primary_key=True)
column_name = field.Field(field.String, primary_key=True)
ordinal_position = field.Field(field.Integer, nullable=True)
column_ordering = field.Field(field.String, nullable=True)
is_nullable = field.Field(field.String)
spanner_type = field.Field(field.String)
class IndexSchema(schema.InformationSchema): """Model for interacting with Spanner index schema table.""" __table__ = 'information_schema.indexes' table_catalog = field.Field(field.String, primary_key=True) table_schema = field.Field(field.String, primary_key=True) table_name = field.Field(field.String, primary_key=True) index_name = field.Field(field.String, primary_key=True) index_type = field.Field(field.String) parent_table_name = field.Field(field.String, nullable=True) is_unique = field.Field(field.Boolean) is_null_filtered = field.Field(field.Boolean) index_state = field.Field(field.String)
class IndexSchema(schema.InformationSchema):
"""Model for interacting with Spanner index schema table."""
__table__ = 'information_schema.indexes'
table_catalog = typing.cast(str, field.Field(field.String, primary_key=True))
table_schema = typing.cast(str, field.Field(field.String, primary_key=True))
table_name = typing.cast(str, field.Field(field.String, primary_key=True))
index_name = typing.cast(str, field.Field(field.String, primary_key=True))
index_type = typing.cast(str, field.Field(field.String))
parent_table_name = typing.cast(Optional[str],
field.Field(field.String, nullable=True))
is_unique = typing.cast(bool, field.Field(field.Boolean))
is_null_filtered = typing.cast(bool, field.Field(field.Boolean))
index_state = typing.cast(str, field.Field(field.String))
class OriginalForeignKeyTestModelTable(spanner_orm.model.Model):
"""ORM Model with the original schema for the ForeignKeyTestModel table."""
__table__ = 'ForeignKeyTestModel'
referencing_key_1 = field.Field(field.String, primary_key=True)
referencing_key_2 = field.Field(field.String, primary_key=True)
referencing_key_3 = field.Field(field.Integer, primary_key=True)
self_referencing_key = field.Field(field.String, nullable=True)
foreign_key_1 = foreign_key_relationship.ForeignKeyRelationship(
'SmallTestModel', {'referencing_key_1': 'key'})
foreign_key_2 = foreign_key_relationship.ForeignKeyRelationship(
'UnittestModel',
{
'referencing_key_2': 'string',
'referencing_key_3': 'int_'
},
)
foreign_key_3 = foreign_key_relationship.ForeignKeyRelationship(
'ForeignKeyTestModel', {'self_referencing_key': 'referencing_key_1'})
class OriginalUnittestModelTable(spanner_orm.model.Model):
"""ORM Model with the original schema for the UnittestModel table."""
__table__ = 'table'
int_ = field.Field(field.Integer, primary_key=True)
int_2 = field.Field(field.Integer, nullable=True)
float_ = field.Field(field.Float, primary_key=True)
float_2 = field.Field(field.Float, nullable=True)
string = field.Field(field.String, primary_key=True)
string_2 = field.Field(field.String, nullable=True)
timestamp = field.Field(field.Timestamp)
string_array = field.Field(field.StringArray, nullable=True)
def test_create_table_foreign_keys_no_model(self, get_model):
get_model.return_value = None
test_update = update.CreateTable(
table_name="RelationshipTestModel",
primary_keys=["parent_key", "child_key"],
fields={
"parent_key":
field.Field(field.String, primary_key=True, name="parent_key"),
"child_key":
field.Field(field.String, primary_key=True, name="child_key"),
},
relations={
"parent":
relationship.Relationship(
"spanner_orm.tests.models.SmallTestModel",
{"parent_key": "key"},
single=True,
),
"parents":
relationship.Relationship(
"spanner_orm.tests.models.SmallTestModel",
{"parent_key": "key"}),
"fk_multicolumn":
relationship.Relationship(
"spanner_orm.tests.models.SmallTestModel",
{
"parent_key": "key",
"parent_key2": "key2"
},
),
},
)
test_update.validate()
test_model_ddl = (
"CREATE TABLE RelationshipTestModel ("
"parent_key STRING(MAX) NOT NULL, "
"child_key STRING(MAX) NOT NULL, "
"CONSTRAINT parent FOREIGN KEY (parent_key) REFERENCES spanner_orm.tests.models.SmallTestModel (key), "
"CONSTRAINT parents FOREIGN KEY (parent_key) REFERENCES spanner_orm.tests.models.SmallTestModel (key), "
"CONSTRAINT fk_multicolumn FOREIGN KEY (parent_key, parent_key2) REFERENCES spanner_orm.tests.models.SmallTestModel (key, key2)) "
"PRIMARY KEY (parent_key, child_key)")
self.assertEqual(test_update.ddl(), test_model_ddl)
def test_add_column(self, get_model):
table_name = models.SmallTestModel.table
get_model.return_value = models.SmallTestModel
new_field = field.Field(field.String, nullable=True)
test_update = update.AddColumn(table_name, 'bar', new_field)
test_update.validate()
self.assertEqual(
test_update.ddl(),
'ALTER TABLE {} ADD COLUMN bar STRING(MAX)'.format(table_name))
class UnittestModelWithoutSecondaryIndexes(model.Model):
"""Same as UnittestModel, but with no secondary indexes."""
__table__ = 'table'
int_ = field.Field(field.Integer, primary_key=True)
int_2 = field.Field(field.Integer, nullable=True)
float_ = field.Field(field.Float, primary_key=True)
float_2 = field.Field(field.Float, nullable=True)
string = field.Field(field.String, primary_key=True)
string_2 = field.Field(field.String, nullable=True)
timestamp = field.Field(field.Timestamp)
string_array = field.Field(field.StringArray, nullable=True)
class IndexColumnSchema(schema.InformationSchema):
"""Model for interacting with Spanner index column schema table."""
__table__ = 'information_schema.index_columns'
table_catalog = typing.cast(str, field.Field(field.String,
primary_key=True))
table_schema = typing.cast(str, field.Field(field.String,
primary_key=True))
table_name = typing.cast(str, field.Field(field.String, primary_key=True))
index_name = typing.cast(str, field.Field(field.String, primary_key=True))
column_name = typing.cast(str, field.Field(field.String, primary_key=True))
ordinal_position = typing.cast(Optional[int],
field.Field(field.Integer, nullable=True))
column_ordering = typing.cast(Optional[str],
field.Field(field.String, nullable=True))
is_nullable = typing.cast(str, field.Field(field.String))
spanner_type = typing.cast(str, field.Field(field.String))
class ForeignKeyTestModel(model.Model):
"""Model class for testing foreign keys."""
__table__ = 'ForeignKeyTestModel'
referencing_key_1 = field.Field(field.String, primary_key=True)
referencing_key_2 = field.Field(field.String, primary_key=True)
referencing_key_3 = field.Field(field.Integer, primary_key=True)
self_referencing_key = field.Field(field.String, nullable=True)
foreign_key_1 = foreign_key_relationship.ForeignKeyRelationship(
'spanner_orm.tests.models.SmallTestModel',
{'referencing_key_1': 'key'})
foreign_key_2 = foreign_key_relationship.ForeignKeyRelationship(
'spanner_orm.tests.models.UnittestModel',
{
'referencing_key_2': 'string',
'referencing_key_3': 'int_'
},
)
foreign_key_3 = foreign_key_relationship.ForeignKeyRelationship(
'spanner_orm.tests.models.ForeignKeyTestModel',
{'self_referencing_key': 'referencing_key_1'})
class UnittestModel(model.Model):
"""Model class used for model testing."""
__table__ = 'table'
int_ = field.Field(field.Integer, primary_key=True)
int_2 = field.Field(field.Integer, nullable=True)
float_ = field.Field(field.Float, primary_key=True)
float_2 = field.Field(field.Float, nullable=True)
string = field.Field(field.String, primary_key=True)
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'])
def test_create_table_interleaved_no_model(self, get_model):
get_model.return_value = None
test_update = update.CreateTable(
table_name="ChildTestModel",
primary_keys=["key", "child_key"],
fields={
"key":
field.Field(field.String, primary_key=True, name="key"),
"child_key":
field.Field(field.String, primary_key=True, name="child_key"),
},
interleaved=models.SmallTestModel,
)
test_update.validate()
test_model_ddl = (
"CREATE TABLE ChildTestModel ("
"key STRING(MAX) NOT NULL, "
"child_key STRING(MAX) NOT NULL) "
"PRIMARY KEY (key, child_key), "
"INTERLEAVE IN PARENT SmallTestModel ON DELETE CASCADE")
self.assertEqual(test_update.ddl(), test_model_ddl)
class ColumnSchema(schema.InformationSchema):
"""Model for interacting with Spanner column schema table."""
__table__ = "information_schema.columns"
table_catalog = field.Field(field.String, primary_key=True)
table_schema = field.Field(field.String, primary_key=True)
table_name = field.Field(field.String, primary_key=True)
column_name = field.Field(field.String, primary_key=True)
ordinal_position = field.Field(field.Integer)
is_nullable = field.Field(field.String)
spanner_type = field.Field(field.String)
@property
def nullable(self) -> bool:
return self.is_nullable == "YES"
@property
def field_type(self) -> Type[field.FieldType]:
for field_type in field.ALL_TYPES:
if self.spanner_type == field_type.ddl():
return field_type
if bool(_string_pattern.match(self.spanner_type)):
return field.String
elif bool(_string_array_pattern.match(self.spanner_type)):
return field.StringArray
raise error.SpannerError(
"No corresponding Type for {}".format(self.spanner_type)
)
@property
def size(self) -> Union[None, int]:
if bool(_string_pattern.match(self.spanner_type)) or bool(
_string_array_pattern.match(self.spanner_type)
):
# Extract digits from string (i.e. STRING(50) -> 50)
return int("".join(filter(str.isdigit, self.spanner_type)))
else:
return None
class ColumnSchema(schema.InformationSchema):
"""Model for interacting with Spanner column schema table."""
__table__ = 'information_schema.columns'
table_catalog = typing.cast(str, field.Field(field.String, primary_key=True))
table_schema = typing.cast(str, field.Field(field.String, primary_key=True))
table_name = typing.cast(str, field.Field(field.String, primary_key=True))
column_name = typing.cast(str, field.Field(field.String, primary_key=True))
ordinal_position = typing.cast(int, field.Field(field.Integer))
is_nullable = typing.cast(str, field.Field(field.String))
spanner_type = typing.cast(str, field.Field(field.String))
def nullable(self) -> bool:
return self.is_nullable == 'YES'
def field_type(self) -> Type[field.FieldType]:
for field_type in field.ALL_TYPES:
if self.spanner_type == field_type.ddl():
return field_type
raise error.SpannerError('No corresponding Type for {}'.format(
self.spanner_type))
class ColumnSchema(schema.InformationSchema):
"""Model for interacting with Spanner column schema table."""
__table__ = "information_schema.columns"
table_catalog = field.Field(field.String, primary_key=True)
table_schema = field.Field(field.String, primary_key=True)
table_name = field.Field(field.String, primary_key=True)
column_name = field.Field(field.String, primary_key=True)
ordinal_position = field.Field(field.Integer)
is_nullable = field.Field(field.String)
spanner_type = field.Field(field.String)
@property
def nullable(self) -> bool:
return self.is_nullable == "YES"
@property
def field_type(self) -> Type[field.FieldType]:
for field_type in field.ALL_TYPES:
if self.spanner_type == field_type.ddl():
return field_type
raise error.SpannerError("No corresponding Type for {}".format(
self.spanner_type))
class InheritanceTestModel(SmallTestModel):
"""Model class used for testing model inheritance."""
value_3 = field.Field(field.String, nullable=True)
def test_create_table_no_model(self, get_model):
get_model.return_value = None
test_update = update.CreateTable(
table_name="table",
primary_keys=["int_", "float_", "string"],
fields={
"int_":
field.Field(field.Integer, primary_key=True, name="int_"),
"int_2":
field.Field(field.Integer, nullable=True, name="int_2"),
"float_":
field.Field(field.Float, primary_key=True, name="float_"),
"float_2":
field.Field(field.Float, nullable=True, name="float_2"),
"string":
field.Field(field.String, primary_key=True, name="string"),
"string_2":
field.Field(field.String, nullable=True, name="string_2"),
"string_3":
field.Field(field.String,
nullable=True,
size=10,
name="string_3"),
"timestamp":
field.Field(field.Timestamp, name="timestamp"),
"timestamp_2":
field.Field(
field.Timestamp,
nullable=True,
allow_commit_timestamp=True,
name="timestamp_2",
),
"date":
field.Field(field.Date, nullable=True, name="date"),
"bool_array":
field.Field(field.BoolArray, nullable=True, name="bool_array"),
"int_array":
field.Field(field.IntegerArray,
nullable=True,
name="int_array"),
"float_array":
field.Field(field.FloatArray,
nullable=True,
name="float_array"),
"date_array":
field.Field(field.DateArray, nullable=True, name="date_array"),
"string_array":
field.Field(field.StringArray,
nullable=True,
name="string_array"),
"string_array_2":
field.Field(field.StringArray,
nullable=True,
size=50,
name="string_array_2"),
},
)
test_update.validate()
test_model_ddl = (
"CREATE TABLE table (int_ INT64 NOT NULL, int_2 INT64,"
" float_ FLOAT64 NOT NULL, float_2 FLOAT64,"
" string STRING(MAX) NOT NULL, string_2 STRING(MAX),"
" string_3 STRING(10),"
" timestamp TIMESTAMP NOT NULL,"
" timestamp_2 TIMESTAMP OPTIONS (allow_commit_timestamp=true),"
" date DATE,"
" bool_array ARRAY<BOOL>,"
" int_array ARRAY<INT64>,"
" float_array ARRAY<FLOAT64>,"
" date_array ARRAY<DATE>,"
" string_array ARRAY<STRING(MAX)>,"
" string_array_2 ARRAY<STRING(50)>) PRIMARY KEY (int_, float_, string)"
)
self.assertEqual(test_update.ddl(), test_model_ddl)