psycopg/psycopg2
1
1
Fork 0
mirror of https://github.com/psycopg/psycopg2.git synced 2024-11-25 10:23:43 +03:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'fix-1487'

Browse Source
This commit is contained in:
Daniele Varrazzo 2022-10-06 02:09:06 +01:00
commit 68d786b610
4 changed files with 185 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
NEWS
View File

@ -1,3 +1,11 @@
What's new in psycopg 2.9.4 (unreleased)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
- Fix `register_composite()`, `register_range()` with customized search_path
(:ticket:`#1487`).
- Handle correctly composite types with names or in schemas requiring escape.
Current release
---------------

47
lib/_range.py
View File

@ -363,33 +363,54 @@ class RangeCaster:
schema = 'public'
# get the type oid and attributes
try:
curs.execute("""\
select rngtypid, rngsubtype,
(select typarray from pg_type where oid = rngtypid)
select rngtypid, rngsubtype, typarray
from pg_range r
join pg_type t on t.oid = rngtypid
join pg_namespace ns on ns.oid = typnamespace
where typname = %s and ns.nspname = %s;
""", (tname, schema))
except ProgrammingError:
if not conn.autocommit:
conn.rollback()
raise
else:
rec = curs.fetchone()
if not rec:
# The above algorithm doesn't work for customized seach_path
# (#1487) The implementation below works better, but, to guarantee
# backwards compatibility, use it only if the original one failed.
try:
savepoint = False
# Because we executed statements earlier, we are either INTRANS
# or we are IDLE only if the transaction is autocommit, in
# which case we don't need the savepoint anyway.
if conn.status == STATUS_IN_TRANSACTION:
curs.execute("SAVEPOINT register_type")
savepoint = True
curs.execute("""\
SELECT rngtypid, rngsubtype, typarray, typname, nspname
from pg_range r
join pg_type t on t.oid = rngtypid
join pg_namespace ns on ns.oid = typnamespace
WHERE t.oid = %s::regtype
""", (name, ))
except ProgrammingError:
pass
else:
rec = curs.fetchone()
if rec:
tname, schema = rec[3:]
finally:
if savepoint:
curs.execute("ROLLBACK TO SAVEPOINT register_type")
# revert the status of the connection as before the command
if (conn_status != STATUS_IN_TRANSACTION
and not conn.autocommit):
if conn_status != STATUS_IN_TRANSACTION and not conn.autocommit:
conn.rollback()
if not rec:
raise ProgrammingError(
f"PostgreSQL type '{name}' not found")
f"PostgreSQL range '{name}' not found")
type, subtype, array = rec
type, subtype, array = rec[:3]
return RangeCaster(name, pyrange,
oid=type, subtype_oid=subtype, array_oid=array)

48
lib/extras.py
View File

@ -357,10 +357,6 @@ class NamedTupleCursor(_cursor):
except StopIteration:
return
# ascii except alnum and underscore
_re_clean = _re.compile(
'[' + _re.escape(' !"#$%&\'()*+,-./:;<=>?@[\\]^`{|}~') + ']')
def _make_nt(self):
key = tuple(d[0] for d in self.description) if self.description else ()
return self._cached_make_nt(key)
@ -369,7 +365,7 @@ class NamedTupleCursor(_cursor):
def _do_make_nt(cls, key):
fields = []
for s in key:
s = cls._re_clean.sub('_', s)
s = _re_clean.sub('_', s)
# Python identifier cannot start with numbers, namedtuple fields
# cannot start with underscore. So...
if s[0] == '_' or '0' <= s[0] <= '9':
@ -1061,6 +1057,7 @@ class CompositeCaster:
return rv
def _create_type(self, name, attnames):
name = _re_clean.sub('_', name)
self.type = namedtuple(name, attnames)
self._ctor = self.type._make
@ -1098,9 +1095,41 @@ ORDER BY attnum;
recs = curs.fetchall()
if not recs:
# The above algorithm doesn't work for customized seach_path
# (#1487) The implementation below works better, but, to guarantee
# backwards compatibility, use it only if the original one failed.
try:
savepoint = False
# Because we executed statements earlier, we are either INTRANS
# or we are IDLE only if the transaction is autocommit, in
# which case we don't need the savepoint anyway.
if conn.status == _ext.STATUS_IN_TRANSACTION:
curs.execute("SAVEPOINT register_type")
savepoint = True
curs.execute("""\
SELECT t.oid, %s, attname, atttypid, typname, nspname
FROM pg_type t
JOIN pg_namespace ns ON typnamespace = ns.oid
JOIN pg_attribute a ON attrelid = typrelid
WHERE t.oid = %%s::regtype
AND attnum > 0 AND NOT attisdropped
ORDER BY attnum;
""" % typarray, (name, ))
except psycopg2.ProgrammingError:
pass
else:
recs = curs.fetchall()
if recs:
tname = recs[0][4]
schema = recs[0][5]
finally:
if savepoint:
curs.execute("ROLLBACK TO SAVEPOINT register_type")
# revert the status of the connection as before the command
if (conn_status != _ext.STATUS_IN_TRANSACTION
and not conn.autocommit):
if conn_status != _ext.STATUS_IN_TRANSACTION and not conn.autocommit:
conn.rollback()
if not recs:
@ -1304,3 +1333,8 @@ def _split_sql(sql):
raise ValueError("the query doesn't contain any '%s' placeholder")
return pre, post
# ascii except alnum and underscore
_re_clean = _re.compile(
'[' + _re.escape(' !"#$%&\'()*+,-./:;<=>?@[\\]^`{|}~') + ']')

110
tests/test_types_extras.py
View File

@ -584,6 +584,68 @@ class AdaptTypeTestCase(ConnectingTestCase):
curs.execute("select (4,8)::typens.typens_ii")
self.assertEqual(curs.fetchone()[0], (4, 8))
@skip_if_no_composite
def test_composite_namespace_path(self):
curs = self.conn.cursor()
curs.execute("""
select nspname from pg_namespace
where nspname = 'typens';
""")
if not curs.fetchone():
curs.execute("create schema typens;")
self.conn.commit()
self._create_type("typens.typensp_ii",
[("a", "integer"), ("b", "integer")])
curs.execute("set search_path=typens,public")
t = psycopg2.extras.register_composite(
"typensp_ii", self.conn)
self.assertEqual(t.schema, 'typens')
curs.execute("select (4,8)::typensp_ii")
self.assertEqual(curs.fetchone()[0], (4, 8))
@skip_if_no_composite
def test_composite_weird_name(self):
curs = self.conn.cursor()
curs.execute("""
select nspname from pg_namespace
where nspname = 'qux.quux';
""")
if not curs.fetchone():
curs.execute('create schema "qux.quux";')
self._create_type('"qux.quux"."foo.bar"',
[("a", "integer"), ("b", "integer")])
t = psycopg2.extras.register_composite(
'"qux.quux"."foo.bar"', self.conn)
self.assertEqual(t.name, 'foo.bar')
self.assertEqual(t.schema, 'qux.quux')
curs.execute('select (4,8)::"qux.quux"."foo.bar"')
self.assertEqual(curs.fetchone()[0], (4, 8))
@skip_if_no_composite
def test_composite_not_found(self):
self.assertRaises(
psycopg2.ProgrammingError, psycopg2.extras.register_composite,
"nosuchtype", self.conn)
self.assertEqual(self.conn.status, ext.STATUS_READY)
cur = self.conn.cursor()
cur.execute("select 1")
self.assertRaises(
psycopg2.ProgrammingError, psycopg2.extras.register_composite,
"nosuchtype", self.conn)
self.assertEqual(self.conn.status, ext.STATUS_IN_TRANSACTION)
self.conn.rollback()
self.conn.autocommit = True
self.assertRaises(
psycopg2.ProgrammingError, psycopg2.extras.register_composite,
"nosuchtype", self.conn)
self.assertEqual(self.conn.status, ext.STATUS_READY)
@skip_if_no_composite
@skip_before_postgres(8, 4)
def test_composite_array(self):
@ -710,22 +772,15 @@ class AdaptTypeTestCase(ConnectingTestCase):
def _create_type(self, name, fields):
curs = self.conn.cursor()
try:
curs.execute("savepoint x")
curs.execute(f"drop type {name} cascade;")
except psycopg2.ProgrammingError:
self.conn.rollback()
curs.execute("rollback to savepoint x")
curs.execute("create type {} as ({});".format(name,
", ".join(["%s %s" % p for p in fields])))
if '.' in name:
schema, name = name.split('.')
else:
schema = 'public'
curs.execute("""\
SELECT t.oid
FROM pg_type t JOIN pg_namespace ns ON typnamespace = ns.oid
WHERE typname = %s and nspname = %s;
""", (name, schema))
curs.execute("SELECT %s::regtype::oid", (name, ))
oid = curs.fetchone()[0]
self.conn.commit()
return oid
@ -1560,6 +1615,18 @@ class RangeCasterTestCase(ConnectingTestCase):
cur = self.conn.cursor()
self.assertRaises(psycopg2.ProgrammingError,
register_range, 'nosuchrange', 'FailRange', cur)
self.assertEqual(self.conn.status, ext.STATUS_READY)
cur.execute("select 1")
self.assertRaises(psycopg2.ProgrammingError,
register_range, 'nosuchrange', 'FailRange', cur)
self.assertEqual(self.conn.status, ext.STATUS_IN_TRANSACTION)
self.conn.rollback()
self.conn.autocommit = True
self.assertRaises(psycopg2.ProgrammingError,
register_range, 'nosuchrange', 'FailRange', cur)
@restore_types
def test_schema_range(self):
@ -1574,7 +1641,7 @@ class RangeCasterTestCase(ConnectingTestCase):
register_range('r1', 'r1', cur)
ra2 = register_range('r2', 'r2', cur)
rars2 = register_range('rs.r2', 'r2', cur)
register_range('rs.r3', 'r3', cur)
rars3 = register_range('rs.r3', 'r3', cur)
self.assertNotEqual(
ra2.typecaster.values[0],
@ -1588,6 +1655,27 @@ class RangeCasterTestCase(ConnectingTestCase):
register_range, 'rs.r1', 'FailRange', cur)
cur.execute("rollback to savepoint x;")
cur2 = self.conn.cursor()
cur2.execute("set local search_path to rs,public")
ra3 = register_range('r3', 'r3', cur2)
self.assertEqual(ra3.typecaster.values[0], rars3.typecaster.values[0])
@skip_if_no_composite
def test_rang_weird_name(self):
cur = self.conn.cursor()
cur.execute("""
select nspname from pg_namespace
where nspname = 'qux.quux';
""")
if not cur.fetchone():
cur.execute('create schema "qux.quux";')
cur.execute('create type "qux.quux"."foo.range" as range (subtype=text)')
r = psycopg2.extras.register_range(
'"qux.quux"."foo.range"', "foorange", cur)
cur.execute('''select '[a,z]'::"qux.quux"."foo.range"''')
self.assertEqual(cur.fetchone()[0], r.range('a', 'z', '[]'))
def test_suite():
return unittest.TestLoader().loadTestsFromName(__name__)