From 5f84b93a11d5ef088d34fede1bd9e34e2766e79f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= <koloite@gmail.com>
Date: Thu, 5 Dec 2024 16:16:07 +0100
Subject: [PATCH] fix: map now() to current_timestamp

Fixes #497
---
 .../sqlalchemy_spanner/sqlalchemy_spanner.py  |  3 ++
 .../mockserver_tests/mock_server_test_base.py | 17 ++++++----
 test/mockserver_tests/test_basics.py          | 33 ++++++++++++++++++-
 3 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py
index e2fb651d..6983a709 100644
--- a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py
+++ b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py
@@ -233,6 +233,9 @@ def get_from_hint_text(self, _, text):
         """
         return text
 
+    def visit_now_func(self, func, **kwargs):
+        return "current_timestamp"
+
     def visit_empty_set_expr(self, type_, **kw):
         """Return an empty set expression of the given type.
 
diff --git a/test/mockserver_tests/mock_server_test_base.py b/test/mockserver_tests/mock_server_test_base.py
index 71e1bf1f..a193dbcb 100644
--- a/test/mockserver_tests/mock_server_test_base.py
+++ b/test/mockserver_tests/mock_server_test_base.py
@@ -21,6 +21,7 @@
     Client,
     ResultSet,
     PingingPool,
+    TypeCode,
 )
 from google.cloud.spanner_v1.database import Database
 from google.cloud.spanner_v1.instance import Instance
@@ -36,6 +37,12 @@ def add_result(sql: str, result: ResultSet):
 
 
 def add_select1_result():
+    add_single_result("select 1", "c", TypeCode.INT64, [("1",)])
+
+
+def add_single_result(
+    sql: str, column_name: str, type_code: spanner_type.TypeCode, row
+):
     result = result_set.ResultSet(
         dict(
             metadata=result_set.ResultSetMetadata(
@@ -45,10 +52,8 @@ def add_select1_result():
                             fields=[
                                 spanner_type.StructType.Field(
                                     dict(
-                                        name="c",
-                                        type=spanner_type.Type(
-                                            dict(code=spanner_type.TypeCode.INT64)
-                                        ),
+                                        name=column_name,
+                                        type=spanner_type.Type(dict(code=type_code)),
                                     )
                                 )
                             ]
@@ -58,8 +63,8 @@ def add_select1_result():
             ),
         )
     )
-    result.rows.extend(["1"])
-    MockServerTestBase.spanner_service.mock_spanner.add_result("select 1", result)
+    result.rows.extend(row)
+    MockServerTestBase.spanner_service.mock_spanner.add_result(sql, result)
 
 
 class MockServerTestBase(fixtures.TestBase):
diff --git a/test/mockserver_tests/test_basics.py b/test/mockserver_tests/test_basics.py
index b6c916c4..543981fa 100644
--- a/test/mockserver_tests/test_basics.py
+++ b/test/mockserver_tests/test_basics.py
@@ -12,8 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
 from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest
-from sqlalchemy import create_engine, select, MetaData, Table, Column, Integer, String
+from sqlalchemy import (
+    create_engine,
+    select,
+    MetaData,
+    Table,
+    Column,
+    Integer,
+    String,
+    func,
+)
 from sqlalchemy.testing import eq_, is_instance_of
 from google.cloud.spanner_v1 import (
     FixedSizePool,
@@ -21,11 +31,13 @@
     ExecuteSqlRequest,
     ResultSet,
     PingingPool,
+    TypeCode,
 )
 from test.mockserver_tests.mock_server_test_base import (
     MockServerTestBase,
     add_select1_result,
     add_result,
+    add_single_result,
 )
 
 
@@ -59,6 +71,25 @@ def test_sqlalchemy_select1(self):
             results = connection.execute(select(1)).fetchall()
             self.verify_select1(results)
 
+    def test_sqlalchemy_select_now(self):
+        now = datetime.datetime.now(datetime.UTC)
+        iso_now = now.isoformat().replace("+00:00", "Z")
+        add_single_result(
+            "SELECT current_timestamp AS now_1",
+            "now_1",
+            TypeCode.TIMESTAMP,
+            [(iso_now,)],
+        )
+        engine = create_engine(
+            "spanner:///projects/p/instances/i/databases/d",
+            connect_args={"client": self.client, "pool": PingingPool(size=10)},
+        )
+        with engine.connect().execution_options(
+            isolation_level="AUTOCOMMIT"
+        ) as connection:
+            spanner_now = connection.execute(select(func.now())).fetchone()[0]
+            eq_(spanner_now.timestamp(), now.timestamp())
+
     def test_create_table(self):
         add_result(
             """SELECT true