Improve handling of sqlstate

src/errors.cpp

@@ -57,71 +57,6 @@ static PyObject* ExceptionFromSqlState(const char* sqlstate)
 }
 
 
-PyObject* RaiseErrorV(const char* sqlstate, PyObject* exc_class, const char* format, ...)
-{
-    PyObject *pAttrs = 0, *pError = 0;
-
-    if (!sqlstate || !*sqlstate)
-        sqlstate = "HY000";
-
-    if (!exc_class)
-        exc_class = ExceptionFromSqlState(sqlstate);
-
-    // Note: Don't use any native strprintf routines.  With Py_ssize_t, we need "%zd", but VC .NET doesn't support it.
-    // PyUnicode_FromFormatV already takes this into account.
-
-    va_list marker;
-    va_start(marker, format);
-    PyObject* pMsg = PyUnicode_FromFormatV(format, marker);
-    va_end(marker);
-    if (!pMsg)
-    {
-        PyErr_NoMemory();
-        return 0;
-    }
-
-    // Create an exception with a 'sqlstate' attribute (set to None if we don't have one) whose 'args' attribute is a
-    // tuple containing the message and sqlstate value.  The 'sqlstate' attribute ensures it is easy to access in
-    // Python (and more understandable to the reader than ex.args[1]), but putting it in the args ensures it shows up
-    // in logs because of the default repr/str.
-
-    pAttrs = Py_BuildValue("(Os)", pMsg, sqlstate);
-    if (pAttrs)
-    {
-        pError = PyObject_CallObject(exc_class, pAttrs);
-        if (pError)
-            RaiseErrorFromException(pError);
-    }
-
-    Py_DECREF(pMsg);
-    Py_XDECREF(pAttrs);
-    Py_XDECREF(pError);
-
-    return 0;
-}
-
-
-bool HasSqlState(PyObject* ex, const char* szSqlState)
-{
-  // Returns true if `ex` is an exception and has the given SQLSTATE.  It is safe to pass 0 for
-  // `ex`.
-
-  if (!ex)
-    return false;
-
-  Object args(PyObject_GetAttrString(ex, "args"));
-  if (!args)
-    return false;
-
-  Object sqlstate(PySequence_GetItem(args, 1));
-  if (!sqlstate || !PyBytes_Check(sqlstate))
-    return false;
-
-  const char* sz = PyBytes_AsString(sqlstate);
-  return (sz && _strcmpi(sz, szSqlState) == 0);
-}
-
-
 static PyObject* GetError(const char* sqlstate, PyObject* exc_class, PyObject* pMsg)
 {
     // pMsg
@@ -155,13 +90,48 @@ static PyObject* GetError(const char* sqlstate, PyObject* exc_class, PyObject* p
     PyTuple_SetItem(pAttrs, 0, pSqlState); // pAttrs now owns the pSqlState reference
 
     pError = PyObject_CallObject(exc_class, pAttrs); // pError will incref pAttrs
+    if (pError)
+        PyObject_SetAttrString(pError, "sqlstate", pSqlState);
 
     Py_XDECREF(pAttrs);
 
     return pError;
 }
 
 
+PyObject* RaiseErrorV(const char* sqlstate, PyObject* exc_class, const char* format, ...)
+{
+    if (!sqlstate || !*sqlstate)
+        sqlstate = "HY000";
+
+    if (!exc_class)
+        exc_class = ExceptionFromSqlState(sqlstate);
+
+    // Note: Don't use any native strprintf routines.  With Py_ssize_t, we need "%zd", but VC .NET doesn't support it.
+    // PyUnicode_FromFormatV already takes this into account.
+
+    va_list marker;
+    va_start(marker, format);
+    PyObject* pMsg = PyUnicode_FromFormatV(format, marker);
+    va_end(marker);
+    if (!pMsg)
+    {
+        PyErr_NoMemory();
+        return 0;
+    }
+
+    // GetError() takes ownership of pMsg, so no Py_DECREF on pMsg after this.
+    PyObject* pError = GetError(sqlstate, exc_class, pMsg);
+    if (pError)
+    {
+        RaiseErrorFromException(pError);
+        Py_DECREF(pError);
+    }
+
+    return 0;
+}
+
+
 static const char* DEFAULT_ERROR = "The driver did not supply an error!";
 
 PyObject* RaiseErrorFromHandle(Connection *conn, const char* szFunction, HDBC hdbc, HSTMT hstmt)
@@ -326,24 +296,3 @@ PyObject* GetErrorFromHandle(Connection *conn, const char* szFunction, HDBC hdbc
 
     return GetError(sqlstate, 0, msg.Detach());
 }
-
-
-static bool GetSqlState(HSTMT hstmt, char* szSqlState)
-{
-    SQLSMALLINT cchMsg;
-    SQLRETURN ret;
-
-    Py_BEGIN_ALLOW_THREADS
-    ret = SQLGetDiagField(SQL_HANDLE_STMT, hstmt, 1, SQL_DIAG_SQLSTATE, (SQLCHAR*)szSqlState, 5, &cchMsg);
-    Py_END_ALLOW_THREADS
-    return SQL_SUCCEEDED(ret);
-}
-
-
-bool HasSqlState(HSTMT hstmt, const char* szSqlState)
-{
-    char szActual[6];
-    if (!GetSqlState(hstmt, szActual))
-        return false;
-    return memcmp(szActual, szSqlState, 5) == 0;
-}

src/errors.h

@@ -29,9 +29,9 @@ PyObject* RaiseErrorV(const char* sqlstate, PyObject* exc_class, const char* for
 
 // Constructs an exception and returns it.
 //
-// This function is like RaiseErrorFromHandle, but gives you the ability to examine the error first (in particular,
-// used to examine the SQLSTATE using HasSqlState).  If you want to use the error, call PyErr_SetObject(ex->ob_type,
-// ex).  Otherwise, dispose of the error using Py_DECREF(ex).
+// This function is like RaiseErrorFromHandle, but gives you the ability to examine the error first.
+// If you want to use the error, call PyErr_SetObject(ex->ob_type, ex).  Otherwise, dispose of the
+// error using Py_DECREF(ex).
 //
 // conn
 //   The connection object, from which it will use the Unicode encoding. May be null if not available.
@@ -42,19 +42,6 @@ PyObject* RaiseErrorV(const char* sqlstate, PyObject* exc_class, const char* for
 //
 PyObject* GetErrorFromHandle(Connection *conn, const char* szFunction, HDBC hdbc, HSTMT hstmt);
 
-
-// Returns true if `ex` is a database exception with SQLSTATE `szSqlState`.  Returns false otherwise.
-// 
-// It is safe to call with ex set to zero.  The SQLSTATE comparison is case-insensitive.
-//
-bool HasSqlState(PyObject* ex, const char* szSqlState);
-
-
-// Returns true if the HSTMT has a diagnostic record with the given SQLSTATE.  This is used after SQLGetData call that
-// returned SQL_SUCCESS_WITH_INFO to see if it also has SQLSTATE 01004, indicating there is more data.
-//
-bool HasSqlState(HSTMT hstmt, const char* szSqlState);
-
 inline PyObject* RaiseErrorFromException(PyObject* pError)
 {
     // PyExceptionInstance_Class doesn't exist in 2.4

src/pyodbc.pyi

@@ -318,8 +318,8 @@ pooling: bool = True
 
 # exceptions
 # https://www.python.org/dev/peps/pep-0249/#exceptions
-class Warning(Exception): ...
-class Error(Exception): ...
+class Warning(Exception): sqlstate: str | None
+class Error(Exception): sqlstate: str | None
 class InterfaceError(Error): ...
 class DatabaseError(Error): ...
 class DataError(DatabaseError): ...

src/pyodbcmodule.cpp

@@ -311,7 +311,7 @@ static bool AllocateEnv()
         }
     }
     Py_DECREF(odbcversion);
-    
+
     if (!SQL_SUCCEEDED(SQLSetEnvAttr(henv, SQL_ATTR_ODBC_VERSION, defaultVersion, sizeof(int))))
     {
         PyErr_SetString(PyExc_RuntimeError, "Unable to set SQL_ATTR_ODBC_VERSION attribute.");
@@ -1121,17 +1121,17 @@ static bool CreateExceptions()
         if (!classdict)
             return false;
 
-        PyObject* doc = PyUnicode_FromString(info.szDoc);
-        if (!doc)
+        if (!strcmp(info.szFullName, "pyodbc.Error") || !strcmp(info.szFullName, "pyodbc.Warning"))
         {
-            Py_DECREF(classdict);
-            return false;
+            Py_INCREF(Py_None);
+            if (PyDict_SetItemString(classdict, "sqlstate", Py_None) < 0) {
+                Py_DECREF(Py_None);
+                Py_DECREF(classdict);
+                return false;
+            }
         }
 
-        PyDict_SetItemString(classdict, "__doc__", doc);
-        Py_DECREF(doc);
-
-        *info.ppexc = PyErr_NewException((char*)info.szFullName, *info.ppexcParent, classdict);
+        *info.ppexc = PyErr_NewExceptionWithDoc((char*)info.szFullName, info.szDoc, *info.ppexcParent, classdict);
         if (*info.ppexc == 0)
         {
             Py_DECREF(classdict);

tests/mysql_test.py

@@ -489,3 +489,19 @@ def _generate_str(length, encoding=None):
     v = v[:length]
 
     return v
+
+
+def test_sqlstate(cursor: pyodbc.Cursor):
+    """Validate new exception property and order of args tuple."""
+    cursor.execute("create table t1 (id int primary key)")
+    cursor.execute("insert into t1 values (1)")
+    try:
+        cursor.execute("insert into t1 values (1)")
+    except pyodbc.IntegrityError as e:
+        assert e.sqlstate == "23000"
+        assert e.args[0] == "23000"
+    try:
+        cursor.execute("insert into t1 values (?)", (1, "bogus"))
+    except pyodbc.ProgrammingError as e:
+        assert e.sqlstate == "HY000"
+        assert e.args[0] == "HY000"

tests/postgresql_test.py

@@ -620,3 +620,19 @@ def _test():
     count_after = sys.getrefcount(encoding)
 
     assert count_after == count_before
+
+
+def test_sqlstate(cursor: pyodbc.Cursor):
+    """Validate new exception property and order of args tuple."""
+    cursor.execute("create table t1 (id int primary key)")
+    cursor.execute("insert into t1 values (1)")
+    try:
+        cursor.execute("insert into t1 values (1)")
+    except pyodbc.IntegrityError as e:
+        assert e.sqlstate == "23505"
+        assert e.args[0] == "23505"
+    try:
+        cursor.execute("insert into t1 values (?)", (1, "bogus"))
+    except pyodbc.ProgrammingError as e:
+        assert e.sqlstate == "HY000"
+        assert e.args[0] == "HY000"

tests/sqlite_test.py

@@ -726,3 +726,19 @@ def test_pickling(cnxn: pyodbc.Connection):
     pickled_rows = pickle.dumps(original_rows)
     unpickled_rows = pickle.loads(pickled_rows)
     assert unpickled_rows == original_rows
+
+
+def test_sqlstate(cursor: pyodbc.Cursor):
+    """Validate new exception property and order of args tuple."""
+    cursor.execute("create table t1 (id int primary key)")
+    cursor.execute("insert into t1 values (1)")
+    try:
+        cursor.execute("insert into t1 values (1)")
+    except pyodbc.Error as e:
+        assert e.sqlstate == "HY000"
+        assert e.args[0] == "HY000"
+    try:
+        cursor.execute("insert into t1 values (?)", (1, "bogus"))
+    except pyodbc.ProgrammingError as e:
+        assert e.sqlstate == "HY000"
+        assert e.args[0] == "HY000"

tests/sqlserver_test.py

@@ -323,6 +323,22 @@ def test_nextset_with_raiserror(cursor: pyodbc.Cursor):
         cursor.nextset()
 
 
+def test_sqlstate(cursor: pyodbc.Cursor):
+    """Validate new exception property and order of args tuple."""
+    cursor.execute("create table t1 (id int primary key)")
+    cursor.execute("insert into t1 values (1)")
+    try:
+        cursor.execute("insert into t1 values (1)")
+    except pyodbc.IntegrityError as e:
+        assert e.sqlstate == "23000"
+        assert e.args[0] == "23000"
+    try:
+        cursor.execute("insert into t1 values (?)", (1, "bogus"))
+    except pyodbc.ProgrammingError as e:
+        assert e.sqlstate == "HY000"
+        assert e.args[0] == "HY000"
+
+
 def test_fixed_unicode(cursor: pyodbc.Cursor):
     value = "t\xebsting"
     cursor.execute("create table t1(s nchar(7))")