fix: support storing columns for indices

.github/workflows/test_suite.yml

@@ -126,6 +126,30 @@ jobs:
           SPANNER_EMULATOR_HOST: localhost:9010
           GOOGLE_CLOUD_PROJECT: appdev-soda-spanner-staging
 
+  system:
+    runs-on: ubuntu-latest
+
+    services:
+      emulator-0:
+        image: gcr.io/cloud-spanner-emulator/emulator:latest
+        ports:
+          - 9010:9010
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v2
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.12
+      - name: Install nox
+        run: python -m pip install nox
+      - name: Run System Tests
+        run: nox -s system
+        env:
+          SPANNER_EMULATOR_HOST: localhost:9010
+          GOOGLE_CLOUD_PROJECT: appdev-soda-spanner-staging
+
   migration_tests:
     runs-on: ubuntu-latest
 

.kokoro/build.sh

@@ -45,4 +45,5 @@ if [[ -n "${NOX_SESSION:-}" ]]; then
     python3 -m nox -s ${NOX_SESSION:-}
 else
     python3 -m nox -s unit
+    python3 -m nox -s system
 fi

create_test_config.py

@@ -18,18 +18,18 @@
 import sys
 
 
-def set_test_config(project, instance, user=None, password=None, host=None, port=None):
+def set_test_config(project, instance, database, user=None, password=None, host=None, port=None):
     config = configparser.ConfigParser()
     if user is not None and password is not None and host is not None and port is not None:
         url = (
             f"spanner+spanner://{user}:{password}@{host}:{port}"
             f"/projects/{project}/instances/{instance}/"
-            "databases/compliance-test"
+            f"databases/{database}"
         )
     else:
         url = (
             f"spanner+spanner:///projects/{project}/instances/{instance}/"
-            "databases/compliance-test"
+            f"databases/{database}"
         )
     config.add_section("db")
     config["db"]["default"] = url
@@ -41,17 +41,18 @@ def set_test_config(project, instance, user=None, password=None, host=None, port
 def main(argv):
     project = argv[0]
     instance = argv[1]
-    if len(argv) == 6:
-        user = argv[2]
-        password = argv[3]
-        host = argv[4]
-        port = argv[5]
+    database = argv[2]
+    if len(argv) == 7:
+        user = argv[3]
+        password = argv[4]
+        host = argv[5]
+        port = argv[6]
     else:
         user = None
         password = None
         host = None
         port = None
-    set_test_config(project, instance, user, password, host, port)
+    set_test_config(project, instance, database, user, password, host, port)
 
 
 if __name__ == "__main__":

create_test_database.py

@@ -14,14 +14,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import configparser
 import os
 import time
 
 from create_test_config import set_test_config
+from google.api_core import datetime_helpers
 from google.api_core.exceptions import AlreadyExists, ResourceExhausted
 from google.cloud.spanner_v1 import Client
 from google.cloud.spanner_v1.instance import Instance
+from google.cloud.spanner_v1.database import Database
 
 
 USE_EMULATOR = os.getenv("SPANNER_EMULATOR_HOST") is not None
@@ -66,43 +67,61 @@ def delete_stale_test_instances():
             )
 
 
-def create_test_instance():
-    configs = list(CLIENT.list_instance_configs())
-    if not USE_EMULATOR:
-        # Filter out non "us" locations
-        configs = [config for config in configs if "asia-southeast1" in config.name]
+def delete_stale_test_databases():
+    """Delete test databases that are older than four hours."""
+    cutoff = (int(time.time()) - 4 * 60 * 60) * 1000
+    instance = CLIENT.instance("sqlalchemy-dialect-test")
+    if not instance.exists():
+        return
+    database_pbs = instance.list_databases()
+    for database_pb in database_pbs:
+        database = Database.from_pb(database_pb, instance)
+        if database.create_time is None:
+            continue
+        create_time = datetime_helpers.to_milliseconds(database_pb.create_time)
+        if create_time > cutoff:
+            continue
+        try:
+            database.drop()
+        except ResourceExhausted:
+            print(
+                "Unable to drop stale database '{}'. May need manual delete.".format(
+                    database.database_id
+                )
+            )
 
-    instance_config = configs[0].name
-    create_time = str(int(time.time()))
-    unique_resource_id = "%s%d" % ("-", 1000 * time.time())
-    instance_id = (
-        "sqlalchemy-dialect-test"
-        if USE_EMULATOR
-        else "sqlalchemy-test" + unique_resource_id
-    )
-    labels = {"python-spanner-sqlalchemy-systest": "true", "created": create_time}
 
-    instance = CLIENT.instance(instance_id, instance_config, labels=labels)
+def create_test_instance():
+    instance_id = "sqlalchemy-dialect-test"
+    instance = CLIENT.instance(instance_id)
+    if not instance.exists():
+        instance_config = f"projects/{PROJECT}/instanceConfigs/regional-us-east1"
+        if USE_EMULATOR:
+            configs = list(CLIENT.list_instance_configs())
+            instance_config = configs[0].name
+        create_time = str(int(time.time()))
+        labels = {"python-spanner-sqlalchemy-systest": "true", "created": create_time}
+
+        instance = CLIENT.instance(instance_id, instance_config, labels=labels)
 
-    try:
-        created_op = instance.create()
-        created_op.result(1800)  # block until completion
-    except AlreadyExists:
-        pass  # instance was already created
+        try:
+            created_op = instance.create()
+            created_op.result(1800)  # block until completion
+        except AlreadyExists:
+            pass  # instance was already created
 
-    if USE_EMULATOR:
-        database = instance.database("compliance-test")
-        database.drop()
+    unique_resource_id = "%s%d" % ("-", 1000 * time.time())
+    database_id = "sqlalchemy-test" + unique_resource_id
 
     try:
-        database = instance.database("compliance-test")
+        database = instance.database(database_id)
         created_op = database.create()
         created_op.result(1800)
     except AlreadyExists:
-        pass  # instance was already created
+        pass  # database was already created
 
-    set_test_config(PROJECT, instance_id)
+    set_test_config(PROJECT, instance_id, database_id)
 
 
-delete_stale_test_instances()
+delete_stale_test_databases()
 create_test_instance()

google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py

@@ -492,6 +492,26 @@ def post_create_table(self, table):
 
         return post_cmds
 
+    def visit_create_index(
+        self, create, include_schema=False, include_table_schema=True, **kw
+    ):
+        text = super().visit_create_index(
+            create, include_schema, include_table_schema, **kw
+        )
+        index = create.element
+        if "spanner" in index.dialect_options:
+            options = index.dialect_options["spanner"]
+            if "storing" in options:
+                storing = options["storing"]
+                storing_columns = [
+                    index.table.c[col] if isinstance(col, str) else col
+                    for col in storing
+                ]
+                text += " STORING (%s)" % ", ".join(
+                    [self.preparer.quote(c.name) for c in storing_columns]
+                )
+        return text
+
     def get_identity_options(self, identity_options):
         text = ["sequence_kind = 'bit_reversed_positive'"]
         if identity_options.start is not None:
@@ -997,15 +1017,35 @@ def get_multi_indexes(
                i.table_schema,
                i.table_name,
                i.index_name,
-               ARRAY_AGG(ic.column_name),
+               (
+                   SELECT ARRAY_AGG(ic.column_name)
+                   FROM information_schema.index_columns ic
+                   WHERE ic.index_name = i.index_name
+                   AND ic.table_catalog = i.table_catalog
+                   AND ic.table_schema = i.table_schema
+                   AND ic.table_name = i.table_name
+                   AND ic.column_ordering is not null
+               ) as columns,
                i.is_unique,
-               ARRAY_AGG(ic.column_ordering)
+               (
+                   SELECT ARRAY_AGG(ic.column_ordering)
+                   FROM information_schema.index_columns ic
+                   WHERE ic.index_name = i.index_name
+                   AND ic.table_catalog = i.table_catalog
+                   AND ic.table_schema = i.table_schema
+                   AND ic.table_name = i.table_name
+                   AND ic.column_ordering is not null
+               ) as column_orderings,
+               (
+                   SELECT ARRAY_AGG(storing.column_name)
+                   FROM information_schema.index_columns storing
+                   WHERE storing.index_name = i.index_name
+                   AND storing.table_catalog = i.table_catalog
+                   AND storing.table_schema = i.table_schema
+                   AND storing.table_name = i.table_name
+                   AND storing.column_ordering is null
+               ) as storing_columns,
             FROM information_schema.indexes as i
-            JOIN information_schema.index_columns AS ic
-                ON  ic.index_name = i.index_name
-                AND ic.table_catalog = i.table_catalog
-                AND ic.table_schema = i.table_schema
-                AND ic.table_name = i.table_name
             JOIN information_schema.tables AS t
                 ON  i.table_catalog = t.table_catalog
                 AND i.table_schema = t.table_schema
@@ -1016,7 +1056,8 @@ def get_multi_indexes(
                 {schema_filter_query}
                 i.index_type != 'PRIMARY_KEY'
                 AND i.spanner_is_managed = FALSE
-            GROUP BY i.table_schema, i.table_name, i.index_name, i.is_unique
+            GROUP BY i.table_catalog, i.table_schema, i.table_name,
+                     i.index_name, i.is_unique
             ORDER BY i.index_name
         """.format(
             table_filter_query=table_filter_query,
@@ -1029,13 +1070,19 @@ def get_multi_indexes(
             result_dict = {}
 
             for row in rows:
+                dialect_options = {}
+                include_columns = row[6]
+                if include_columns:
+                    dialect_options["spanner_storing"] = include_columns
                 index_info = {
                     "name": row[2],
                     "column_names": row[3],
                     "unique": row[4],
                     "column_sorting": {
-                        col: order for col, order in zip(row[3], row[5])
+                        col: order.lower() for col, order in zip(row[3], row[5])
                     },
+                    "include_columns": include_columns if include_columns else [],
+                    "dialect_options": dialect_options,
                 }
                 row[0] = row[0] or None
                 table_info = result_dict.get((row[0], row[1]), [])

migration_test_cleanup.py

@@ -25,10 +25,11 @@ def main(argv):
 
   project = re.findall(r"projects(.*?)instances", db_url)
   instance_id = re.findall(r"instances(.*?)databases", db_url)
+  database_id = re.findall(r"databases(.*?)$", db_url)
 
   client = spanner.Client(project="".join(project).replace("/", ""))
   instance = client.instance(instance_id="".join(instance_id).replace("/", ""))
-  database = instance.database("compliance-test")
+  database = instance.database("".join(database_id).replace("/", ""))
 
   database.update_ddl(["DROP TABLE account", "DROP TABLE alembic_version"]).result(120)
 

noxfile.py

@@ -252,6 +252,41 @@ def compliance_test_20(session):
     )
 
 
+@nox.session()
+def system(session):
+    """Run SQLAlchemy dialect system test suite."""
+
+    # Sanity check: Only run tests if the environment variable is set.
+    if not os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", "") and not os.environ.get(
+        "SPANNER_EMULATOR_HOST", ""
+    ):
+        session.skip(
+            "Credentials or emulator host must be set via environment variable"
+        )
+
+    if os.environ.get("RUN_COMPLIANCE_TESTS", "true") == "false" and not os.environ.get(
+        "SPANNER_EMULATOR_HOST", ""
+    ):
+        session.skip("RUN_COMPLIANCE_TESTS is set to false, skipping")
+
+    session.install(
+        "pytest",
+        "pytest-cov",
+        "pytest-asyncio",
+    )
+
+    session.install("mock")
+    session.install(".[tracing]")
+    session.install("opentelemetry-api==1.27.0")
+    session.install("opentelemetry-sdk==1.27.0")
+    session.install("opentelemetry-instrumentation==0.48b0")
+    session.run("python", "create_test_database.py")
+
+    session.install("sqlalchemy>=2.0")
+
+    session.run("py.test", "--quiet", os.path.join("test", "system"), *session.posargs)
+
+
 @nox.session(python=DEFAULT_PYTHON_VERSION)
 def unit(session):
     """Run unit tests."""
@@ -263,7 +298,9 @@ def unit(session):
     session.install("opentelemetry-api==1.27.0")
     session.install("opentelemetry-sdk==1.27.0")
     session.install("opentelemetry-instrumentation==0.48b0")
-    session.run("python", "create_test_config.py", "my-project", "my-instance")
+    session.run(
+        "python", "create_test_config.py", "my-project", "my-instance", "my-database"
+    )
     session.run("py.test", "--quiet", os.path.join("test/unit"), *session.posargs)
 
 
@@ -281,6 +318,7 @@ def mockserver(session):
         "create_test_config.py",
         "my-project",
         "my-instance",
+        "my-database",
         "none",
         "AnonymousCredentials",
         "localhost",
@@ -323,21 +361,12 @@ def _migration_test(session):
 
     session.run("python", "create_test_database.py")
 
-    project = os.getenv(
-        "GOOGLE_CLOUD_PROJECT",
-        os.getenv("PROJECT_ID", "emulator-test-project"),
-    )
-    db_url = (
-        f"spanner+spanner:///projects/{project}/instances/"
-        "sqlalchemy-dialect-test/databases/compliance-test"
-    )
-
     config = configparser.ConfigParser()
     if os.path.exists("test.cfg"):
         config.read("test.cfg")
     else:
         config.read("setup.cfg")
-    db_url = config.get("db", "default", fallback=db_url)
+    db_url = config.get("db", "default")
 
     session.run("alembic", "init", "test_migration")