feat: proto column feature tests

tests/_fixtures.py

@@ -68,6 +68,29 @@
      email STRING(MAX),
      deleted BOOL NOT NULL )
      PRIMARY KEY(id, commit_ts DESC);
+CREATE PROTO BUNDLE (
+    spanner.examples.music.SingerInfo,
+    spanner.examples.music.Genre,
+    );
+CREATE TABLE Types (
+    RowID		INT64 NOT NULL,
+    Int64a		INT64,
+    Bytes		BYTES(MAX),
+    Int64Array	ARRAY<INT64>,
+    BytesArray	ARRAY<BYTES(MAX)>,
+    ProtoMessage    spanner.examples.music.SingerInfo,
+    ProtoEnum   spanner.examples.music.Genre,
+    ProtoMessageArray   ARRAY<spanner.examples.music.SingerInfo>,
+    ProtoEnumArray  ARRAY<spanner.examples.music.Genre>, ) 
+    PRIMARY KEY (RowID);
+CREATE TABLE Singers (
+     SingerId   INT64 NOT NULL,
+     FirstName  STRING(1024),
+     LastName   STRING(1024),
+     SingerInfo spanner.examples.music.SingerInfo,
+     SingerGenre spanner.examples.music.Genre, )
+     PRIMARY KEY (SingerId);
+CREATE INDEX SingerByGenre ON Singers(SingerGenre) STORING (FirstName, LastName);
 """
 
 EMULATOR_DDL = """\
@@ -159,8 +182,25 @@
 CREATE INDEX name ON contacts(first_name, last_name);
 """
 
+PROTO_COLUMNS_DDL = """\
+CREATE PROTO BUNDLE (
+    spanner.examples.music.SingerInfo,
+    spanner.examples.music.Genre,
+    );
+CREATE TABLE Singers (
+     SingerId   INT64 NOT NULL,
+     FirstName  STRING(1024),
+     LastName   STRING(1024),
+     SingerInfo spanner.examples.music.SingerInfo,
+     SingerGenre spanner.examples.music.Genre, )
+     PRIMARY KEY (SingerId);
+"""
+
 DDL_STATEMENTS = [stmt.strip() for stmt in DDL.split(";") if stmt.strip()]
 EMULATOR_DDL_STATEMENTS = [
     stmt.strip() for stmt in EMULATOR_DDL.split(";") if stmt.strip()
 ]
 PG_DDL_STATEMENTS = [stmt.strip() for stmt in PG_DDL.split(";") if stmt.strip()]
+PROTO_COLUMNS_DDL_STATEMENTS = [
+    stmt.strip() for stmt in PROTO_COLUMNS_DDL.split(";") if stmt.strip()
+]

tests/system/_helpers.py

@@ -65,6 +65,8 @@
     )
 )
 
+PROTO_COLUMNS_DDL_STATEMENTS = _fixtures.PROTO_COLUMNS_DDL_STATEMENTS
+
 retry_true = retry.RetryResult(operator.truth)
 retry_false = retry.RetryResult(operator.not_)
 

tests/system/conftest.py

@@ -73,6 +73,15 @@ def database_dialect():
         else DatabaseDialect.GOOGLE_STANDARD_SQL
     )
 
+#TODO: Check how to read through relative paths in python
+@pytest.fixture(scope="session")
+def proto_descriptor_file():
+    return open(
+        "/Users/sriharshach/GitHub/Python/python-spanner/samples/samples/testdata/descriptors.pb",
+        "rb",
+    ).read()
+    # return open("../../samples/samples/testdata/descriptors.pb", "rb").read()
+
 
 @pytest.fixture(scope="session")
 def spanner_client():
@@ -177,7 +186,9 @@ def shared_instance(
 
 
 @pytest.fixture(scope="session")
-def shared_database(shared_instance, database_operation_timeout, database_dialect):
+def shared_database(
+    shared_instance, database_operation_timeout, database_dialect, proto_descriptor_file
+):
     database_name = _helpers.unique_id("test_database")
     pool = spanner_v1.BurstyPool(labels={"testcase": "database_api"})
     if database_dialect == DatabaseDialect.POSTGRESQL:
@@ -198,6 +209,7 @@ def shared_database(shared_instance, database_operation_timeout, database_dialec
             ddl_statements=_helpers.DDL_STATEMENTS,
             pool=pool,
             database_dialect=database_dialect,
+            proto_descriptors=proto_descriptor_file,
         )
         operation = database.create()
         operation.result(database_operation_timeout)  # raises on failure / timeout.

tests/system/test_database_api.py

@@ -11,7 +11,7 @@
 # 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.
-
+import base64
 import time
 import uuid
 
@@ -22,6 +22,7 @@
 from google.cloud import spanner_v1
 from google.cloud.spanner_v1.pool import FixedSizePool, PingingPool
 from google.type import expr_pb2
+from samples.samples.testdata import singer_pb2
 from . import _helpers
 from . import _sample_data
 
@@ -562,3 +563,214 @@ def _unit_of_work(transaction, name):
         rows = list(after.read(sd.COUNTERS_TABLE, sd.COUNTERS_COLUMNS, sd.ALL))
 
     assert len(rows) == 2
+
+
+def test_create_table_with_proto_columns(
+        not_emulator,
+        not_postgres,
+        shared_instance,
+        databases_to_delete,
+        proto_descriptor_file,
+):
+    proto_cols_db_id = _helpers.unique_id("proto-columns")
+
+    proto_cols_database = shared_instance.database(
+        proto_cols_db_id,
+        ddl_statements=_helpers.PROTO_COLUMNS_DDL_STATEMENTS,
+        proto_descriptors=proto_descriptor_file,
+    )
+    operation = proto_cols_database.create()
+    operation.result(DBAPI_OPERATION_TIMEOUT)  # raises on failure / timeout.
+
+    databases_to_delete.append(proto_cols_database)
+
+    proto_cols_database.reload()
+    assert proto_cols_database.proto_descriptors is not None
+    assert any("PROTO BUNDLE" in stmt for stmt in proto_cols_database.ddl_statements)
+
+
+def test_compatibility_of_proto_columns(not_emulator, not_postgres, shared_database):
+    singer_info = singer_pb2.SingerInfo()
+    singer_genre = singer_pb2.Genre.ROCK
+
+    singer_info.singer_id = 1
+    singer_info.birth_date = "January"
+    singer_info.nationality = "Country1"
+    singer_info.genre = singer_genre
+
+    singer_info_array = [singer_info]
+    singer_genre_array = [singer_genre]
+
+    singer_info_bytes = base64.b64encode(singer_info.SerializeToString())
+    singer_info_genre_val = 3
+    singer_info_bytes_array = [singer_info_bytes]
+    singer_info_genre_val_array = [singer_info_genre_val]
+
+    # Test Compatibility between (Bytes, ProtoMessage), (Int, ProtoEnum),
+    # (BytesArray, ProtoMessageArray) and (IntArray, ProtoEnumArray)
+    with shared_database.batch() as batch:
+        batch.insert(
+            table="Types",
+            columns=(
+                "RowID",
+                "Int64a",
+                "Bytes",
+                "Int64Array",
+                "BytesArray",
+                "ProtoMessage",
+                "ProtoEnum",
+                "ProtoMessageArray",
+                "ProtoEnumArray",
+            ),
+            values=[
+                (
+                    1,
+                    singer_info_genre_val,
+                    singer_info_bytes,
+                    singer_info_genre_val_array,
+                    singer_info_bytes_array,
+                    singer_info,
+                    singer_genre,
+                    singer_info_array,
+                    singer_genre_array,
+                ),
+                (
+                    2,
+                    singer_genre,
+                    singer_info,
+                    singer_genre_array,
+                    singer_info_array,
+                    singer_info_bytes,
+                    singer_info_genre_val,
+                    singer_info_bytes_array,
+                    singer_info_genre_val_array,
+                ),
+            ],
+        )
+
+    with shared_database.snapshot() as snapshot:
+        rows = list(
+            snapshot.read(
+                table="Types",
+                columns=(
+                    "RowID",
+                    "Int64a",
+                    "Bytes",
+                    "Int64Array",
+                    "BytesArray",
+                    "ProtoMessage",
+                    "ProtoEnum",
+                    "ProtoMessageArray",
+                    "ProtoEnumArray",
+                ),
+                keyset=spanner_v1.KeySet(keys=[[1]]),
+                column_info={
+                    "ProtoMessage": singer_pb2.SingerInfo(),
+                    "ProtoEnum": singer_pb2.Genre,
+                    "ProtoMessageArray": singer_pb2.SingerInfo(),
+                    "ProtoEnumArray": singer_pb2.Genre,
+                },
+            )
+        )
+
+    assert len(rows) == 1
+
+    # Verify data
+    assert rows[0][0] == 1
+    assert rows[0][1] == singer_info_genre_val
+    assert rows[0][2] == singer_info_bytes
+    assert rows[0][3] == singer_info_genre_val_array
+    assert rows[0][4] == singer_info_bytes_array
+    assert rows[0][5] == singer_info
+    assert rows[0][6] == "ROCK"
+    assert rows[0][7] == singer_info_array
+    assert rows[0][8] == ["ROCK"]
+
+
+def test_proto_columns_dml_parameterized_queries_pk_indexes(
+        not_emulator, not_postgres, shared_database
+):
+    def _unit_of_work(transaction):
+        singer1_proto_enum = singer_pb2.Genre.ROCK
+        singer1_proto_message = singer_pb2.SingerInfo()
+        singer1_proto_message.singer_id = 1
+        singer1_proto_message.birth_date = "January"
+        singer1_proto_message.nationality = "Country1"
+        singer1_proto_message.genre = singer1_proto_enum
+
+        singer2_proto_enum = singer_pb2.Genre.FOLK
+        singer2_proto_message = singer_pb2.SingerInfo()
+        singer2_proto_message.singer_id = 2
+        singer2_proto_message.birth_date = "February"
+        singer2_proto_message.nationality = "Country2"
+        singer2_proto_message.genre = singer2_proto_enum
+
+        # Test parameterized queries with proto columns
+        transaction.execute_update(
+            "INSERT INTO Singers (SingerId, FirstName, LastName, SingerInfo, SingerGenre)"
+            " VALUES (1, 'Singer1', 'Singer1', @singerInfo, @singerGenre)",
+            params={
+                "singerInfo": singer1_proto_message,
+                "singerGenre": singer1_proto_enum,
+            },
+            param_types={
+                "singerInfo": spanner_v1.param_types.ProtoMessage(
+                    singer1_proto_message
+                ),
+                "singerGenre": spanner_v1.param_types.ProtoEnum(singer_pb2.Genre),
+            },
+        )
+        transaction.execute_update(
+            "INSERT INTO Singers (SingerId, FirstName, LastName, SingerInfo, SingerGenre)"
+            " VALUES (2, 'Singer2', 'Singer2', @singerInfo, @singerGenre)",
+            params={
+                "singerInfo": singer2_proto_message,
+                "singerGenre": singer2_proto_enum,
+            },
+            param_types={
+                "singerInfo": spanner_v1.param_types.ProtoMessage(
+                    singer2_proto_message
+                ),
+                "singerGenre": spanner_v1.param_types.ProtoEnum(singer_pb2.Genre),
+            },
+        )
+
+    shared_database.run_in_transaction(_unit_of_work)
+
+    # Test parameterized queries with proto columns
+    with shared_database.snapshot() as snapshot:
+        rows = list(
+            snapshot.execute_sql(
+                "SELECT SingerId FROM Singers WHERE SingerGenre=@genre",
+                params={"genre": singer_pb2.Genre.FOLK},
+                param_types={
+                    "genre": spanner_v1.param_types.ProtoEnum(singer_pb2.Genre)
+                },
+            )
+        )
+        assert rows[0][0] == 2
+
+    with shared_database.snapshot() as snapshot:
+        rows = list(
+            snapshot.execute_sql(
+                "SELECT SingerId FROM Singers WHERE SingerInfo.Nationality=@country",
+                params={"country": "Country1"},
+                param_types={"country": spanner_v1.param_types.STRING},
+            )
+        )
+        assert rows[0][0] == 1
+
+    # Test read using index with proto columns
+    with shared_database.snapshot() as snapshot:
+        keyset = spanner_v1.KeySet(keys=[[singer_pb2.Genre.ROCK]])
+        rows = list(
+            snapshot.read(
+                table="Singers",
+                columns=("SingerId", "FirstName", "LastName"),
+                keyset=keyset,
+                index="SingerByGenre",
+            )
+        )
+        assert len(rows) == 1
+        assert rows[0][0] == 1
+        assert rows[0][1] == "Singer1"