diff --git a/Lib/sqlite3/test/dbapi.py b/Lib/sqlite3/test/dbapi.py index 518b8ae..202bd38 100644 --- a/Lib/sqlite3/test/dbapi.py +++ b/Lib/sqlite3/test/dbapi.py @@ -225,6 +225,13 @@ class CursorTests(unittest.TestCase): def CheckExecuteArgString(self): self.cu.execute("insert into test(name) values (?)", ("Hugo",)) + def CheckExecuteArgStringWithZeroByte(self): + self.cu.execute("insert into test(name) values (?)", ("Hu\x00go",)) + + self.cu.execute("select name from test where id=?", (self.cu.lastrowid,)) + row = self.cu.fetchone() + self.assertEqual(row[0], "Hu\x00go") + def CheckExecuteWrongNoOfArgs1(self): # too many parameters try: diff --git a/Lib/sqlite3/test/factory.py b/Lib/sqlite3/test/factory.py index 1adab2f..7f6f347 100644 --- a/Lib/sqlite3/test/factory.py +++ b/Lib/sqlite3/test/factory.py @@ -189,13 +189,48 @@ class TextFactoryTests(unittest.TestCase): def tearDown(self): self.con.close() +class TextFactoryTestsWithEmbeddedZeroBytes(unittest.TestCase): + def setUp(self): + self.con = sqlite.connect(":memory:") + self.con.execute("create table test (value text)") + self.con.execute("insert into test (value) values (?)", ("a\x00b",)) + + def CheckString(self): + # text_factory defaults to str + row = self.con.execute("select value from test").fetchone() + self.assertIs(type(row[0]), str) + self.assertEqual(row[0], "a\x00b") + + def CheckBytes(self): + self.con.text_factory = bytes + row = self.con.execute("select value from test").fetchone() + self.assertIs(type(row[0]), bytes) + self.assertEqual(row[0], b"a\x00b") + + def CheckBytearray(self): + self.con.text_factory = bytearray + row = self.con.execute("select value from test").fetchone() + self.assertIs(type(row[0]), bytearray) + self.assertEqual(row[0], b"a\x00b") + + def CheckCustom(self): + # A custom factory should receive a bytes argument + self.con.text_factory = lambda x: x + row = self.con.execute("select value from test").fetchone() + self.assertIs(type(row[0]), bytes) + self.assertEqual(row[0], b"a\x00b") + + def tearDown(self): + self.con.close() + def suite(): connection_suite = unittest.makeSuite(ConnectionFactoryTests, "Check") cursor_suite = unittest.makeSuite(CursorFactoryTests, "Check") row_suite_compat = unittest.makeSuite(RowFactoryTestsBackwardsCompat, "Check") row_suite = unittest.makeSuite(RowFactoryTests, "Check") text_suite = unittest.makeSuite(TextFactoryTests, "Check") - return unittest.TestSuite((connection_suite, cursor_suite, row_suite_compat, row_suite, text_suite)) + text_zero_bytes_suite = unittest.makeSuite(TextFactoryTestsWithEmbeddedZeroBytes, "Check") + return unittest.TestSuite((connection_suite, cursor_suite, row_suite_compat, row_suite, text_suite, text_zero_bytes_suite)) def test(): runner = unittest.TextTestRunner() diff --git a/Modules/_sqlite/cursor.c b/Modules/_sqlite/cursor.c index 75ee73a..99a8612 100644 --- a/Modules/_sqlite/cursor.c +++ b/Modules/_sqlite/cursor.c @@ -267,9 +267,9 @@ PyObject* _pysqlite_build_column_name(const char* colname) } } -PyObject* pysqlite_unicode_from_string(const char* val_str, int optimize) +PyObject* pysqlite_unicode_from_string(const char* val_str, Py_ssize_t size, int optimize) { - return PyUnicode_FromString(val_str); + return PyUnicode_FromStringAndSize(val_str, size); } /* @@ -354,10 +354,11 @@ PyObject* _pysqlite_fetch_one_row(pysqlite_Cursor* self) converted = PyFloat_FromDouble(sqlite3_column_double(self->statement->st, i)); } else if (coltype == SQLITE_TEXT) { val_str = (const char*)sqlite3_column_text(self->statement->st, i); + nbytes = sqlite3_column_bytes(self->statement->st, i); if ((self->connection->text_factory == (PyObject*)&PyUnicode_Type) || (self->connection->text_factory == pysqlite_OptimizedUnicode)) { - converted = pysqlite_unicode_from_string(val_str, + converted = pysqlite_unicode_from_string(val_str, nbytes, self->connection->text_factory == pysqlite_OptimizedUnicode ? 1 : 0); if (!converted) { @@ -382,11 +383,11 @@ PyObject* _pysqlite_fetch_one_row(pysqlite_Cursor* self) } } } else if (self->connection->text_factory == (PyObject*)&PyBytes_Type) { - converted = PyBytes_FromString(val_str); + converted = PyBytes_FromStringAndSize(val_str, nbytes); } else if (self->connection->text_factory == (PyObject*)&PyByteArray_Type) { - converted = PyByteArray_FromStringAndSize(val_str, strlen(val_str)); + converted = PyByteArray_FromStringAndSize(val_str, nbytes); } else { - converted = PyObject_CallFunction(self->connection->text_factory, "y", val_str); + converted = PyObject_CallFunction(self->connection->text_factory, "y#", val_str, nbytes); } } else { /* coltype == SQLITE_BLOB */ diff --git a/Modules/_sqlite/statement.c b/Modules/_sqlite/statement.c index 4e039c1..cc9c310 100644 --- a/Modules/_sqlite/statement.c +++ b/Modules/_sqlite/statement.c @@ -129,9 +129,9 @@ int pysqlite_statement_bind_parameter(pysqlite_Statement* self, int pos, PyObjec rc = sqlite3_bind_double(self->st, pos, PyFloat_AsDouble(parameter)); break; case TYPE_UNICODE: - string = _PyUnicode_AsString(parameter); + string = _PyUnicode_AsStringAndSize(parameter, &buflen); if (string != NULL) - rc = sqlite3_bind_text(self->st, pos, string, -1, SQLITE_TRANSIENT); + rc = sqlite3_bind_text(self->st, pos, string, buflen, SQLITE_TRANSIENT); else rc = -1; break;