From cd67d3d2fe085b207268be649ef282fc6032a8cc Mon Sep 17 00:00:00 2001 From: Blake Rouse Date: Mon, 8 Sep 2014 12:05:28 -0400 Subject: [PATCH] Modify truncate to use lo_truncate64. Use HAVE_LO64 define to use new lo_*64 methods. Check size of offset and length for versions without LO64. --- psycopg/lobject_int.c | 12 ++++++++ psycopg/lobject_type.c | 20 +++++++++++-- setup.py | 7 +++++ tests/test_lobject.py | 67 ++++++++++++++++++++++++++++++------------ 4 files changed, 86 insertions(+), 20 deletions(-) diff --git a/psycopg/lobject_int.c b/psycopg/lobject_int.c index 731d3da7..19d1b9d3 100644 --- a/psycopg/lobject_int.c +++ b/psycopg/lobject_int.c @@ -389,7 +389,11 @@ lobject_seek(lobjectObject *self, long pos, int whence) Py_BEGIN_ALLOW_THREADS; pthread_mutex_lock(&(self->conn->lock)); +#if HAVE_LO64 where = lo_lseek64(self->conn->pgconn, self->fd, pos, whence); +#else + where = (long)lo_lseek(self->conn->pgconn, self->fd, (int)pos, whence); +#endif Dprintf("lobject_seek: where = %Ld", where); if (where < 0) collect_error(self->conn, &error); @@ -416,7 +420,11 @@ lobject_tell(lobjectObject *self) Py_BEGIN_ALLOW_THREADS; pthread_mutex_lock(&(self->conn->lock)); +#if HAVE_LO64 where = lo_tell64(self->conn->pgconn, self->fd); +#else + where = (long)lo_tell(self->conn->pgconn, self->fd); +#endif Dprintf("lobject_tell: where = %Ld", where); if (where < 0) collect_error(self->conn, &error); @@ -473,7 +481,11 @@ lobject_truncate(lobjectObject *self, size_t len) Py_BEGIN_ALLOW_THREADS; pthread_mutex_lock(&(self->conn->lock)); +#if HAVE_LO64 + retvalue = lo_truncate64(self->conn->pgconn, self->fd, len); +#else retvalue = lo_truncate(self->conn->pgconn, self->fd, len); +#endif Dprintf("lobject_truncate: result = %d", retvalue); if (retvalue < 0) collect_error(self->conn, &error); diff --git a/psycopg/lobject_type.c b/psycopg/lobject_type.c index ba27e4fb..82aa186d 100644 --- a/psycopg/lobject_type.c +++ b/psycopg/lobject_type.c @@ -175,6 +175,14 @@ psyco_lobj_seek(lobjectObject *self, PyObject *args) EXC_IF_LOBJ_LEVEL0(self); EXC_IF_LOBJ_UNMARKED(self); +#if !HAVE_LO64 + if (offset > INT_MAX) { + psyco_set_error(InterfaceError, NULL, + "offset out of range"); + return NULL; + } +#endif + if ((pos = lobject_seek(self, offset, whence)) < 0) return NULL; @@ -255,15 +263,23 @@ psyco_lobj_get_closed(lobjectObject *self, void *closure) static PyObject * psyco_lobj_truncate(lobjectObject *self, PyObject *args) { - int len = 0; + long len = 0; - if (!PyArg_ParseTuple(args, "|i", &len)) + if (!PyArg_ParseTuple(args, "|l", &len)) return NULL; EXC_IF_LOBJ_CLOSED(self); EXC_IF_LOBJ_LEVEL0(self); EXC_IF_LOBJ_UNMARKED(self); +#if !HAVE_LO64 + if (len > INT_MAX) { + psyco_set_error(InterfaceError, NULL, + "len out of range"); + return NULL; + } +#endif + if (lobject_truncate(self, len) < 0) return NULL; diff --git a/setup.py b/setup.py index 49069f71..9fdf0b56 100644 --- a/setup.py +++ b/setup.py @@ -415,6 +415,13 @@ class psycopg_build_ext(build_ext): define_macros.append(("PG_VERSION_HEX", "0x%02X%02X%02X" % (int(pgmajor), int(pgminor), int(pgpatch)))) + + # enable lo64 if postgres >= 9.3 + if int(pgmajor) >= 9 and int(pgminor) >= 3: + define_macros.append(("HAVE_LO64", "1")) + else: + define_macros.append(("HAVE_LO64", "0")) + except Warning: w = sys.exc_info()[1] # work around py 2/3 different syntax sys.stderr.write("Error: %s\n" % w) diff --git a/tests/test_lobject.py b/tests/test_lobject.py index 36b24273..c692fcd0 100755 --- a/tests/test_lobject.py +++ b/tests/test_lobject.py @@ -225,24 +225,6 @@ class LargeObjectTests(LargeObjectTestCase): self.assertEqual(lo.seek(-2, 2), length - 2) self.assertEqual(lo.read(), "ta") - def test_seek_tell_greater_than_2gb(self): - lo = self.conn.lobject() - - # write chunks until its 3gb - length = 0 - for _ in range(24): - # each chunk is written with 128mb - length += lo.write("data" * (1 << 25)) - self.assertEqual(lo.tell(), length) - lo.close() - lo = self.conn.lobject(lo.oid) - - # seek to 3gb - 4, last written text should be data - offset = (1 << 31) + (1 << 30) - 4 # 2gb + 1gb - 4 - self.assertEqual(lo.seek(offset, 0), offset) - self.assertEqual(lo.tell(), offset) - self.assertEqual(lo.read(), "data") - def test_unlink(self): lo = self.conn.lobject() lo.unlink() @@ -458,6 +440,55 @@ decorate_all_tests(LargeObjectTruncateTests, skip_if_no_lo, skip_lo_if_green, skip_if_no_truncate) +def skip_if_no_lo64(f): + @wraps(f) + def skip_if_no_lo64_(self): + if self.conn.server_version < 90300: + return self.skipTest("large objects 64bit only supported from PG 9.3") + else: + return f(self) + + return skip_if_no_lo64_ + +class LargeObject64Tests(LargeObjectTestCase): + def test_seek_tell_truncate_greater_than_2gb(self): + lo = self.conn.lobject() + + length = (1 << 31) + (1 << 30) # 2gb + 1gb = 3gb + lo.truncate(length) + + self.assertEqual(lo.seek(length, 0), length) + self.assertEqual(lo.tell(), length) + +decorate_all_tests(LargeObject64Tests, + skip_if_no_lo, skip_lo_if_green, skip_if_no_truncate, skip_if_no_lo64) + + +def skip_if_lo64(f): + @wraps(f) + def skip_if_lo64_(self): + if self.conn.server_version >= 90300: + return self.skipTest("large objects 64bit only supported from PG 9.3") + else: + return f(self) + + return skip_if_lo64_ + +class LargeObjectNot64Tests(LargeObjectTestCase): + def test_seek_larger_than_2gb(self): + lo = self.conn.lobject() + offset = 1 << 32 # 4gb + self.assertRaises(psycopg2.InterfaceError, lo.seek, offset, 0) + + def test_truncate_larger_than_2gb(self): + lo = self.conn.lobject() + length = 1 << 32 # 4gb + self.assertRaises(psycopg2.InterfaceError, lo.truncate, length) + +decorate_all_tests(LargeObjectNot64Tests, + skip_if_no_lo, skip_lo_if_green, skip_if_no_truncate, skip_if_lo64) + + def test_suite(): return unittest.TestLoader().loadTestsFromName(__name__)