From 6eb67283df0db55b6497c057c4445bb545ccd7ad Mon Sep 17 00:00:00 2001
From: vimiix <i@vimiix.com>
Date: Wed, 21 Aug 2024 12:23:47 +0800
Subject: [PATCH] fix: use timedelta instead of time to cast TIME column in B
 compatibility mode(#IAIE7I)

In mysql compatibility mode, we should convert the TIME column to a python datetime.timedelta type,
just like the behavior of PyMySQL and mysql-connector-python.

See also:
- PyMySQL: https://github.com/PyMySQL/PyMySQL/blob/95635f587ba9076e71a223b113efb08ac34a361d/pymysql/converters.py#L344
- mysql-connector-python: https://github.com/mysql/mysql-connector-python/blob/59817f3de42d98ac2875a14ea265d77bd0ac3805/mysql-connector-python/lib/mysql/connector/conversion.py#L593
---
 psycopg/connection.h        |  2 +-
 psycopg/connection_type.c   | 25 +++++++++++++++++++------
 psycopg/cursor_int.c        | 10 ++++++++++
 psycopg/typecast.h          |  2 ++
 psycopg/typecast_builtins.c | 15 +++++++++++++++
 5 files changed, 47 insertions(+), 7 deletions(-)

diff --git a/psycopg/connection.h b/psycopg/connection.h
index ff3894e..932ca6a 100644
--- a/psycopg/connection.h
+++ b/psycopg/connection.h
@@ -49,7 +49,7 @@ extern "C" {
 /* sql_compatibility values */
 #define SQL_COMPATIBILITY_A         1
 #define SQL_COMPATIBILITY_OTHER     5
-// #define SQL_COMPATIBILITY_B     2
+#define SQL_COMPATIBILITY_B         2
 // #define SQL_COMPATIBILITY_C     3
 // #define SQL_COMPATIBILITY_PG    4
 
diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
index 1672c99..d3eb601 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -1364,6 +1364,24 @@ end:
 }
 
 
+static int set_sql_compatibility(connectionObject *self, char *value)
+{
+    switch (value[0]) {
+        case 'A':
+        case 'a':
+            self->sql_compatibility = SQL_COMPATIBILITY_A;
+            break;
+        case 'B':
+        case 'b':
+            self->sql_compatibility = SQL_COMPATIBILITY_B;
+            break;
+        default:
+            self->sql_compatibility = SQL_COMPATIBILITY_OTHER;
+            break;
+    }
+    return 0;
+}
+
 /* initialization and finalization methods */
 
 static int
@@ -1414,16 +1432,11 @@ connection_setup(connectionObject *self, const char *dsn, long int async)
     Py_BEGIN_ALLOW_THREADS;
     pthread_mutex_lock(&self->lock);
     sql_compatibility_value = pq_get_guc_locked(self, "sql_compatibility", &_save);
+    set_sql_compatibility(self, sql_compatibility_value);
     if (register_type_uint(self, &_save)) { goto exit; }
     pthread_mutex_unlock(&self->lock);
     Py_END_ALLOW_THREADS;
 
-    if (strcmp(sql_compatibility_value, "A") == 0) {
-        self->sql_compatibility = SQL_COMPATIBILITY_A;
-    } else {
-        self->sql_compatibility = SQL_COMPATIBILITY_OTHER;
-    }
-
 exit:
     if (sql_compatibility_value){
         free(sql_compatibility_value);
diff --git a/psycopg/cursor_int.c b/psycopg/cursor_int.c
index 7009ee8..487a529 100644
--- a/psycopg/cursor_int.c
+++ b/psycopg/cursor_int.c
@@ -57,6 +57,16 @@ curs_get_cast(cursorObject *self, PyObject *oid)
     if (cast) { return cast; }
 
     /* global lookup */
+    if (self->conn->sql_compatibility == SQL_COMPATIBILITY_B) {
+        /*
+        In mysql compatibility mode, we should convert the TIME column to a python datetime.timedelta type,
+        just like the behavior of PyMySQL and mysql-connector-python.
+        */
+        if (PyLong_Check(oid)) {
+            long value = convert_time_oid_to_interval_oid(PyLong_AsLong(oid));
+            oid = PyLong_FromLong(value);
+        }
+    }
     cast = PyDict_GetItem(psyco_types, oid);
     Dprintf("curs_get_cast:        global dict: %p", cast);
     if (cast) { return cast; }
diff --git a/psycopg/typecast.h b/psycopg/typecast.h
index 84af98f..4f7a5bc 100644
--- a/psycopg/typecast.h
+++ b/psycopg/typecast.h
@@ -88,5 +88,7 @@ HIDDEN PyObject *typecast_array_from_python(
 HIDDEN PyObject *typecast_cast(
     PyObject *self, const char *str, Py_ssize_t len, PyObject *curs);
 
+HIDDEN long convert_time_oid_to_interval_oid(long value);
+
 #endif /* !defined(PSYCOPG_TYPECAST_H) */
 PyObject *typecast_new(PyObject *name, PyObject *values, PyObject *cast, PyObject *base);
\ No newline at end of file
diff --git a/psycopg/typecast_builtins.c b/psycopg/typecast_builtins.c
index e6bc799..4b53323 100644
--- a/psycopg/typecast_builtins.c
+++ b/psycopg/typecast_builtins.c
@@ -31,6 +31,21 @@ static long int typecast_MACADDRARRAY_types[] = {1040, 0};
 static long int typecast_UNKNOWN_types[] = {705, 0};
 
 
+long convert_time_oid_to_interval_oid(long value)
+{
+    int i = 0;
+    while (typecast_TIME_types[i] != 0) {
+        if (typecast_TIME_types[i] == value) {
+            Dprintf("convert time oid to interval oid");
+            value = typecast_INTERVAL_types[0];
+            break;
+        }
+        i++;
+    }
+    return value;
+}
+
+
 static typecastObject_initlist typecast_builtins[] = {
   {"NUMBER", typecast_NUMBER_types, typecast_NUMBER_cast, NULL},
   {"LONGINTEGER", typecast_LONGINTEGER_types, typecast_LONGINTEGER_cast, NULL},
-- 
Gitee