def register_inet(oid=None, conn_or_curs=None): """Create the INET type and an Inet adapter. :param oid: oid for the PostgreSQL :sql:`inet` type, or 2-items sequence with oids of the type and the array. If not specified, use PostgreSQL standard oids. :param conn_or_curs: where to register the typecaster. If not specified, register it globally. """ if not oid: oid1 = 869 oid2 = 1041 elif isinstance(oid, (list, tuple)): oid1, oid2 = oid else: oid1 = oid oid2 = 1041 _ext.INET = _ext.new_type((oid1, ), "INET", lambda data, cursor: data and Inet(data) or None) _ext.INETARRAY = _ext.new_array_type((oid2, ), "INETARRAY", _ext.INET) _ext.register_type(_ext.INET, conn_or_curs) _ext.register_type(_ext.INETARRAY, conn_or_curs) return _ext.INET
def register_uuid(oids=None, conn_or_curs=None): """Create the UUID type and an uuid.UUID adapter. :param oids: oid for the PostgreSQL :sql:`uuid` type, or 2-items sequence with oids of the type and the array. If not specified, use PostgreSQL standard oids. :param conn_or_curs: where to register the typecaster. If not specified, register it globally. """ import uuid if not oids: oid1 = 2950 oid2 = 2951 elif isinstance(oids, (list, tuple)): oid1, oid2 = oids else: oid1 = oids oid2 = 2951 _ext.UUID = _ext.new_type( (oid1, ), "UUID", lambda data, cursor: data and uuid.UUID(data) or None) _ext.UUIDARRAY = _ext.new_array_type((oid2, ), "UUID[]", _ext.UUID) _ext.register_type(_ext.UUID, conn_or_curs) _ext.register_type(_ext.UUIDARRAY, conn_or_curs) _ext.register_adapter(uuid.UUID, UUID_adapter) return _ext.UUID
def test_encoding_name(self): self.conn.set_client_encoding("EUC_JP") # conn.encoding is 'EUCJP' now. cur = self.conn.cursor() extensions.register_type(extensions.UNICODE, cur) cur.execute("select 'foo'::text;") self.assertEqual(cur.fetchone()[0], _u(b'foo'))
def test_unicode(self): cur = self.conn.cursor() ext.register_type(ext.UNICODE, cur) snowman = u"\u2603" # unicode in statement psycopg2.extras.execute_batch( cur, "insert into testfast (id, data) values (%%s, %%s) -- %s" % snowman, [(1, 'x')]) cur.execute("select id, data from testfast where id = 1") self.assertEqual(cur.fetchone(), (1, 'x')) # unicode in data psycopg2.extras.execute_batch( cur, "insert into testfast (id, data) values (%s, %s)", [(2, snowman)]) cur.execute("select id, data from testfast where id = 2") self.assertEqual(cur.fetchone(), (2, snowman)) # unicode in both psycopg2.extras.execute_batch( cur, "insert into testfast (id, data) values (%%s, %%s) -- %s" % snowman, [(3, snowman)]) cur.execute("select id, data from testfast where id = 3") self.assertEqual(cur.fetchone(), (3, snowman))
def register_composite(name, conn_or_curs, globally=False, factory=None): """Register a typecaster to convert a composite type into a tuple. :param name: the name of a PostgreSQL composite type, e.g. created using the |CREATE TYPE|_ command :param conn_or_curs: a connection or cursor used to find the type oid and components; the typecaster is registered in a scope limited to this object, unless *globally* is set to `!True` :param globally: if `!False` (default) register the typecaster only on *conn_or_curs*, otherwise register it globally :param factory: if specified it should be a `CompositeCaster` subclass: use it to :ref:`customize how to cast composite types <custom-composite>` :return: the registered `CompositeCaster` or *factory* instance responsible for the conversion """ if factory is None: factory = CompositeCaster caster = factory._from_db(name, conn_or_curs) _ext.register_type(caster.typecaster, not globally and conn_or_curs or None) if caster.array_typecaster is not None: _ext.register_type(caster.array_typecaster, not globally and conn_or_curs or None) return caster
def test_unicode(self): curs = self.conn.cursor() curs.execute("SHOW server_encoding") server_encoding = curs.fetchone()[0] if server_encoding != "UTF8": return self.skipTest( "Unicode test skipped since server encoding is %s" % server_encoding) data = _u(b"""some data with \t chars to escape into, 'quotes', \xe2\x82\xac euro sign and \\ a backslash too. """) _unichr = chr if six.PY3 else unichr data += _u(b"").join( map(_unichr, [u for u in range(1, 65536) if not 0xD800 <= u <= 0xDFFF ])) # surrogate area self.conn.set_client_encoding('UNICODE') extensions.register_type(extensions.UNICODE, self.conn) curs.execute("SELECT %s::text;", (data, )) res = curs.fetchone()[0] self.assertEqual(res, data) self.assert_(not self.conn.notices)
def register_uuid(oids=None, conn_or_curs=None): """Create the UUID type and an uuid.UUID adapter. :param oids: oid for the PostgreSQL :sql:`uuid` type, or 2-items sequence with oids of the type and the array. If not specified, use PostgreSQL standard oids. :param conn_or_curs: where to register the typecaster. If not specified, register it globally. """ import uuid if not oids: oid1 = 2950 oid2 = 2951 elif isinstance(oids, (list, tuple)): oid1, oid2 = oids else: oid1 = oids oid2 = 2951 _ext.UUID = _ext.new_type((oid1, ), "UUID", lambda data, cursor: data and uuid.UUID(data) or None) _ext.UUIDARRAY = _ext.new_array_type((oid2,), "UUID[]", _ext.UUID) _ext.register_type(_ext.UUID, conn_or_curs) _ext.register_type(_ext.UUIDARRAY, conn_or_curs) _ext.register_adapter(uuid.UUID, UUID_adapter) return _ext.UUID
def test_stolen_reference_bug(self): def fish(val, cur): gc.collect() return 42 UUID = extensions.new_type((2950,), "UUID", fish) extensions.register_type(UUID, self.conn) curs = self.conn.cursor() curs.execute("select 'b5219e01-19ab-4994-b71e-149225dc51e4'::uuid") curs.fetchone()
def testGenericArrayNull(self): def caster(s, cur): if s is None: return "nada" return int(s) * 2 base = extensions.new_type((23,), "INT4", caster) array = extensions.new_array_type((1007,), "INT4ARRAY", base) extensions.register_type(array, self.conn) a = self.execute("select '{1,2,3}'::int4[]") self.assertEqual(a, [2,4,6]) a = self.execute("select '{1,2,NULL}'::int4[]") self.assertEqual(a, [2,4,'nada'])
def testGenericArrayNull(self): def caster(s, cur): if s is None: return "nada" return int(s) * 2 base = extensions.new_type((23, ), "INT4", caster) array = extensions.new_array_type((1007, ), "INT4ARRAY", base) extensions.register_type(array, self.conn) a = self.execute("select '{1,2,3}'::int4[]") self.assertEqual(a, [2, 4, 6]) a = self.execute("select '{1,2,NULL}'::int4[]") self.assertEqual(a, [2, 4, 'nada'])
def test_cast_specificity(self): curs = self.conn.cursor() self.assertEqual("foo", curs.cast(705, 'foo')) D = extensions.new_type((705,), "DOUBLING", lambda v, c: v * 2) extensions.register_type(D, self.conn) self.assertEqual("foofoo", curs.cast(705, 'foo')) T = extensions.new_type((705,), "TREBLING", lambda v, c: v * 3) extensions.register_type(T, curs) self.assertEqual("foofoofoo", curs.cast(705, 'foo')) curs2 = self.conn.cursor() self.assertEqual("foofoo", curs2.cast(705, 'foo'))
def test_copy_expert_textiobase(self): self.conn.set_client_encoding('latin1') self._create_temp_table() # the above call closed the xn if sys.version_info[0] < 3: abin = b''.join(map(chr, range(32, 127) + range(160, 256))) abin = abin.decode('latin1') about = abin.replace('\\', '\\\\') else: abin = bytes(list(range(32, 127)) + list(range(160, 256)))\ .decode('latin1') about = abin.replace('\\', '\\\\') import io f = io.StringIO() f.write(about) f.seek(0) curs = self.conn.cursor() extensions.register_type(extensions.UNICODE, curs) curs.copy_expert('COPY tcopy (data) FROM STDIN', f) curs.execute("select data from tcopy;") self.assertEqual(curs.fetchone()[0], abin) f = io.StringIO() curs.copy_expert('COPY tcopy (data) TO STDOUT', f) f.seek(0) self.assertEqual(f.readline().rstrip(), about) # same tests with setting size f = io.StringIO() f.write(about) f.seek(0) exp_size = 123 # hack here to leave file as is, only check size when reading real_read = f.read def read(_size, f=f, exp_size=exp_size): assert _size == exp_size return real_read(_size) f.read = read curs.copy_expert('COPY tcopy (data) FROM STDIN', f, size=exp_size) curs.execute("select data from tcopy;") self.assertEqual(curs.fetchone()[0], abin)
def test_unicode(self): curs = self.conn.cursor() curs.execute("SHOW server_encoding") server_encoding = curs.fetchone()[0] if server_encoding != "UTF8": return self.skipTest( "Unicode test skipped since server encoding is %s" % server_encoding) data = _u(b"""some data with \t chars to escape into, 'quotes', \xe2\x82\xac euro sign and \\ a backslash too. """) _unichr = chr if six.PY3 else unichr data += _u(b"").join(map(_unichr, [ u for u in range(1,65536) if not 0xD800 <= u <= 0xDFFF ])) # surrogate area self.conn.set_client_encoding('UNICODE') extensions.register_type(extensions.UNICODE, self.conn) curs.execute("SELECT %s::text;", (data,)) res = curs.fetchone()[0] self.assertEqual(res, data) self.assert_(not self.conn.notices)
def test_koi8(self): self.conn.set_client_encoding('KOI8') curs = self.conn.cursor() if sys.version_info[0] < 3: data = b''.join(map(chr, range(32, 127) + range(128, 256))) else: data = bytes(list(range(32, 127)) + list(range(128, 256)))\ .decode('koi8_r') # as string curs.execute("SELECT %s::text;", (data,)) res = curs.fetchone()[0] self.assertEqual(res, data) self.assert_(not self.conn.notices) # as unicode if sys.version_info[0] < 3: extensions.register_type(extensions.UNICODE, self.conn) data = data.decode('koi8_r') curs.execute("SELECT %s::text;", (data,)) res = curs.fetchone()[0] self.assertEqual(res, data) self.assert_(not self.conn.notices)
def test_koi8(self): self.conn.set_client_encoding('KOI8') curs = self.conn.cursor() if sys.version_info[0] < 3: data = b''.join(map(chr, range(32, 127) + range(128, 256))) else: data = bytes(list(range(32, 127)) + list(range(128, 256)))\ .decode('koi8_r') # as string curs.execute("SELECT %s::text;", (data, )) res = curs.fetchone()[0] self.assertEqual(res, data) self.assert_(not self.conn.notices) # as unicode if sys.version_info[0] < 3: extensions.register_type(extensions.UNICODE, self.conn) data = data.decode('koi8_r') curs.execute("SELECT %s::text;", (data, )) res = curs.fetchone()[0] self.assertEqual(res, data) self.assert_(not self.conn.notices)
def _register(self, scope=None): register_type(self.typecaster, scope) if self.array_typecaster is not None: register_type(self.array_typecaster, scope) register_adapter(self.range, self.adapter)
import re from collections import defaultdict, namedtuple # Import the Psycopg2 connector for PostgreSQL try: # Prefer Psycopg2cffi, which should work both for PyPy and CPython import psycopg2cffi.extensions as psycopg2ext import psycopg2cffi as psycopg2 except ImportError: # Try plain Psycopg2 for CPython import psycopg2.extensions as psycopg2ext import psycopg2 # Make Psycopg2 and PostgreSQL happy with UTF-8 psycopg2ext.register_type(psycopg2ext.UNICODE) psycopg2ext.register_type(psycopg2ext.UNICODEARRAY) # Hack to make this Python program executable from the utils subdirectory basepath, _ = os.path.split(os.path.realpath(__file__)) _UTILS = os.sep + "utils" if basepath.endswith(_UTILS): basepath = basepath[0 : -len(_UTILS)] sys.path.append(basepath) # Note: We can't use settings from ReynirPackage because it # reads package resource streams, not plain text files from settings import Settings, LineReader, ConfigError
def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None, array_oid=None): """Register adapter and typecaster for `!dict`\-\ |hstore| conversions. :param conn_or_curs: a connection or cursor: the typecaster will be registered only on this object unless *globally* is set to `!True` :param globally: register the adapter globally, not only on *conn_or_curs* :param unicode: if `!True`, keys and values returned from the database will be `!unicode` instead of `!str`. The option is not available on Python 3 :param oid: the OID of the |hstore| type if known. If not, it will be queried on *conn_or_curs*. :param array_oid: the OID of the |hstore| array type if known. If not, it will be queried on *conn_or_curs*. The connection or cursor passed to the function will be used to query the database and look for the OID of the |hstore| type (which may be different across databases). If querying is not desirable (e.g. with :ref:`asynchronous connections <async-support>`) you may specify it in the *oid* parameter, which can be found using a query such as :sql:`SELECT 'hstore'::regtype::oid`. Analogously you can obtain a value for *array_oid* using a query such as :sql:`SELECT 'hstore[]'::regtype::oid`. Note that, when passing a dictionary from Python to the database, both strings and unicode keys and values are supported. Dictionaries returned from the database have keys/values according to the *unicode* parameter. The |hstore| contrib module must be already installed in the database (executing the ``hstore.sql`` script in your ``contrib`` directory). Raise `~psycopg2.ProgrammingError` if the type is not found. """ if oid is None: oid = HstoreAdapter.get_oids(conn_or_curs) if oid is None or not oid[0]: raise psycopg2.ProgrammingError( "hstore type not found in the database. " "please install it from your 'contrib/hstore.sql' file") else: array_oid = oid[1] oid = oid[0] if isinstance(oid, int): oid = (oid,) if array_oid is not None: if isinstance(array_oid, int): array_oid = (array_oid,) else: array_oid = tuple([x for x in array_oid if x]) # create and register the typecaster if _sys.version_info[0] < 3 and unicode: cast = HstoreAdapter.parse_unicode else: cast = HstoreAdapter.parse HSTORE = _ext.new_type(oid, "HSTORE", cast, unicode) _ext.register_type(HSTORE, not globally and conn_or_curs or None) _ext.register_adapter(dict, HstoreAdapter) if array_oid: HSTOREARRAY = _ext.new_array_type(array_oid, "HSTOREARRAY", HSTORE) _ext.register_type(HSTOREARRAY, not globally and conn_or_curs or None)
def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None, array_oid=None): """Register adapter and typecaster for `!dict`\-\ |hstore| conversions. :param conn_or_curs: a connection or cursor: the typecaster will be registered only on this object unless *globally* is set to `!True` :param globally: register the adapter globally, not only on *conn_or_curs* :param unicode: if `!True`, keys and values returned from the database will be `!unicode` instead of `!str`. The option is not available on Python 3 :param oid: the OID of the |hstore| type if known. If not, it will be queried on *conn_or_curs*. :param array_oid: the OID of the |hstore| array type if known. If not, it will be queried on *conn_or_curs*. The connection or cursor passed to the function will be used to query the database and look for the OID of the |hstore| type (which may be different across databases). If querying is not desirable (e.g. with :ref:`asynchronous connections <async-support>`) you may specify it in the *oid* parameter, which can be found using a query such as :sql:`SELECT 'hstore'::regtype::oid`. Analogously you can obtain a value for *array_oid* using a query such as :sql:`SELECT 'hstore[]'::regtype::oid`. Note that, when passing a dictionary from Python to the database, both strings and unicode keys and values are supported. Dictionaries returned from the database have keys/values according to the *unicode* parameter. The |hstore| contrib module must be already installed in the database (executing the ``hstore.sql`` script in your ``contrib`` directory). Raise `~psycopg2.ProgrammingError` if the type is not found. """ if oid is None: oid = HstoreAdapter.get_oids(conn_or_curs) if oid is None or not oid[0]: raise psycopg2.ProgrammingError( "hstore type not found in the database. " "please install it from your 'contrib/hstore.sql' file") else: array_oid = oid[1] oid = oid[0] if isinstance(oid, int): oid = (oid, ) if array_oid is not None: if isinstance(array_oid, int): array_oid = (array_oid, ) else: array_oid = tuple([x for x in array_oid if x]) # create and register the typecaster if _sys.version_info[0] < 3 and unicode: cast = HstoreAdapter.parse_unicode else: cast = HstoreAdapter.parse HSTORE = _ext.new_type(oid, "HSTORE", cast) _ext.register_type(HSTORE, not globally and conn_or_curs or None) _ext.register_adapter(dict, HstoreAdapter) if array_oid: HSTOREARRAY = _ext.new_array_type(array_oid, "HSTOREARRAY", HSTORE) _ext.register_type(HSTOREARRAY, not globally and conn_or_curs or None)