mirror of
https://github.com/psycopg/psycopg2.git
synced 2024-11-29 20:23:45 +03:00
409 lines
13 KiB
Python
409 lines
13 KiB
Python
# Copyright (C) 2006 Fabio Tranchitella <fabio@tranchitella.it>
|
|
#
|
|
# psycopg2da is free software: you can redistribute it and/or modify it
|
|
# under the terms of the GNU Lesser General Public License as published
|
|
# by the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
#
|
|
# In addition, as a special exception, the copyright holders give
|
|
# permission to link this program with the OpenSSL library (or with
|
|
# modified versions of OpenSSL that use the same license as OpenSSL),
|
|
# and distribute linked combinations including the two.
|
|
#
|
|
# You must obey the GNU Lesser General Public License in all respects for
|
|
# all of the code used other than OpenSSL.
|
|
#
|
|
# psycopg2da is distributed in the hope that it will be useful, but WITHOUT
|
|
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
|
|
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public
|
|
# License for more details.
|
|
#
|
|
# Based on ZPsycopgDA.
|
|
#
|
|
# If you prefer you can use this package using the ZPL license as
|
|
# published on the Zope web site, http://www.zope.org/Resources/ZPL.
|
|
|
|
"""PostgreSQL Database Adapter for Zope 3"""
|
|
|
|
from zope.interface import implements
|
|
from zope.rdb import ZopeDatabaseAdapter, parseDSN, ZopeConnection, ZopeCursor
|
|
from zope.rdb.interfaces import DatabaseException, IZopeConnection
|
|
from zope.publisher.interfaces import Retry
|
|
|
|
from datetime import date, time, datetime, timedelta
|
|
|
|
import psycopg2
|
|
import psycopg2.extensions
|
|
import re
|
|
import sys
|
|
|
|
# OIDs from psycopg/pgtypes.h
|
|
DATE_OID = 1082
|
|
TIME_OID = 1083
|
|
TIMETZ_OID = 1266
|
|
TIMESTAMP_OID = 1114
|
|
TIMESTAMPTZ_OID = 1184
|
|
INTERVAL_OID = 1186
|
|
CHAR_OID = 18
|
|
TEXT_OID = 25
|
|
BPCHAR_OID = 1042
|
|
VARCHAR_OID = 1043
|
|
|
|
# date/time parsing functions
|
|
_dateFmt = re.compile(r"^(\d\d\d\d)-?([01]\d)-?([0-3]\d)$")
|
|
|
|
def parse_date(s):
|
|
"""Parses ISO-8601 compliant dates and returns a tuple (year, month,
|
|
day).
|
|
|
|
The following formats are accepted:
|
|
YYYY-MM-DD (extended format)
|
|
YYYYMMDD (basic format)
|
|
"""
|
|
m = _dateFmt.match(s)
|
|
if m is None:
|
|
raise ValueError, 'invalid date string: %s' % s
|
|
year, month, day = m.groups()
|
|
return int(year), int(month), int(day)
|
|
|
|
|
|
_timeFmt = re.compile(
|
|
r"^([0-2]\d)(?::?([0-5]\d)(?::?([0-5]\d)(?:[.,](\d+))?)?)?$")
|
|
|
|
def parse_time(s):
|
|
"""Parses ISO-8601 compliant times and returns a tuple (hour, minute,
|
|
second).
|
|
|
|
The following formats are accepted:
|
|
HH:MM:SS.ssss or HHMMSS.ssss
|
|
HH:MM:SS,ssss or HHMMSS,ssss
|
|
HH:MM:SS or HHMMSS
|
|
HH:MM or HHMM
|
|
HH
|
|
"""
|
|
m = _timeFmt.match(s)
|
|
if m is None:
|
|
raise ValueError, 'invalid time string: %s' % s
|
|
hr, mn, sc, msc = m.groups(0)
|
|
if msc != 0:
|
|
sc = float("%s.%s" % (sc, msc))
|
|
else:
|
|
sc = int(sc)
|
|
return int(hr), int(mn), sc
|
|
|
|
|
|
_tzFmt = re.compile(r"^([+-])([0-2]\d)(?::?([0-5]\d))?$")
|
|
|
|
def parse_tz(s):
|
|
"""Parses ISO-8601 timezones and returns the offset east of UTC in
|
|
minutes.
|
|
|
|
The following formats are accepted:
|
|
+/-HH:MM
|
|
+/-HHMM
|
|
+/-HH
|
|
Z (equivalent to +0000)
|
|
"""
|
|
if s == 'Z':
|
|
return 0
|
|
m = _tzFmt.match(s)
|
|
if m is None:
|
|
raise ValueError, 'invalid time zone: %s' % s
|
|
d, hoff, moff = m.groups(0)
|
|
if d == "-":
|
|
return - int(hoff) * 60 - int(moff)
|
|
return int(hoff) * 60 + int(moff)
|
|
|
|
|
|
_tzPos = re.compile(r"[Z+-]")
|
|
|
|
def parse_timetz(s):
|
|
"""Parses ISO-8601 compliant times that may include timezone information
|
|
and returns a tuple (hour, minute, second, tzoffset).
|
|
|
|
tzoffset is the offset east of UTC in minutes. It will be None if s does
|
|
not include time zone information.
|
|
|
|
Formats accepted are those listed in the descriptions of parse_time() and
|
|
parse_tz(). Time zone should immediatelly follow time without intervening
|
|
spaces.
|
|
"""
|
|
m = _tzPos.search(s)
|
|
if m is None:
|
|
return parse_time(s) + (None,)
|
|
pos = m.start()
|
|
return parse_time(s[:pos]) + (parse_tz(s[pos:]),)
|
|
|
|
|
|
_datetimeFmt = re.compile(r"[T ]")
|
|
|
|
def _split_datetime(s):
|
|
"""Split date and time parts of ISO-8601 compliant timestamp and
|
|
return a tuple (date, time).
|
|
|
|
' ' or 'T' used to separate date and time parts.
|
|
"""
|
|
m = _datetimeFmt.search(s)
|
|
if m is None:
|
|
raise ValueError, 'time part of datetime missing: %s' % s
|
|
pos = m.start()
|
|
return s[:pos], s[pos + 1:]
|
|
|
|
|
|
def parse_datetime(s):
|
|
"""Parses ISO-8601 compliant timestamp and returns a tuple (year, month,
|
|
day, hour, minute, second).
|
|
|
|
Formats accepted are those listed in the descriptions of parse_date() and
|
|
parse_time() with ' ' or 'T' used to separate date and time parts.
|
|
"""
|
|
dt, tm = _split_datetime(s)
|
|
return parse_date(dt) + parse_time(tm)
|
|
|
|
|
|
def parse_datetimetz(s):
|
|
"""Parses ISO-8601 compliant timestamp that may include timezone
|
|
information and returns a tuple (year, month, day, hour, minute, second,
|
|
tzoffset).
|
|
|
|
tzoffset is the offset east of UTC in minutes. It will be None if s does
|
|
not include time zone information.
|
|
|
|
Formats accepted are those listed in the descriptions of parse_date() and
|
|
parse_timetz() with ' ' or 'T' used to separate date and time parts.
|
|
"""
|
|
dt, tm = _split_datetime(s)
|
|
return parse_date(dt) + parse_timetz(tm)
|
|
|
|
|
|
def parse_interval(s):
|
|
"""Parses PostgreSQL interval notation and returns a tuple (years, months,
|
|
days, hours, minutes, seconds).
|
|
|
|
Values accepted:
|
|
interval ::= date
|
|
| time
|
|
| date time
|
|
date ::= date_comp
|
|
| date date_comp
|
|
date_comp ::= 1 'day'
|
|
| number 'days'
|
|
| 1 'month'
|
|
| 1 'mon'
|
|
| number 'months'
|
|
| number 'mons'
|
|
| 1 'year'
|
|
| number 'years'
|
|
time ::= number ':' number
|
|
| number ':' number ':' number
|
|
| number ':' number ':' number '.' fraction
|
|
"""
|
|
years = months = days = 0
|
|
hours = minutes = seconds = 0
|
|
elements = s.split()
|
|
# Tests with 7.4.6 on Ubuntu 5.4 interval output returns 'mon' and 'mons'
|
|
# and not 'month' or 'months' as expected. I've fixed this and left
|
|
# the original matches there too in case this is dependant on
|
|
# OS or PostgreSQL release.
|
|
for i in range(0, len(elements) - 1, 2):
|
|
count, unit = elements[i:i+2]
|
|
if unit == 'day' and count == '1':
|
|
days += 1
|
|
elif unit == 'days':
|
|
days += int(count)
|
|
elif unit == 'month' and count == '1':
|
|
months += 1
|
|
elif unit == 'mon' and count == '1':
|
|
months += 1
|
|
elif unit == 'months':
|
|
months += int(count)
|
|
elif unit == 'mons':
|
|
months += int(count)
|
|
elif unit == 'year' and count == '1':
|
|
years += 1
|
|
elif unit == 'years':
|
|
years += int(count)
|
|
else:
|
|
raise ValueError, 'unknown time interval %s %s' % (count, unit)
|
|
if len(elements) % 2 == 1:
|
|
hours, minutes, seconds = parse_time(elements[-1])
|
|
return (years, months, days, hours, minutes, seconds)
|
|
|
|
|
|
# Type conversions
|
|
def _conv_date(s, cursor):
|
|
if s:
|
|
return date(*parse_date(s))
|
|
|
|
def _conv_time(s, cursor):
|
|
if s:
|
|
hr, mn, sc = parse_time(s)
|
|
sc, micro = divmod(sc, 1.0)
|
|
micro = round(micro * 1000000)
|
|
return time(hr, mn, int(sc), int(micro))
|
|
|
|
def _conv_timetz(s, cursor):
|
|
if s:
|
|
from zope.datetime import tzinfo
|
|
hr, mn, sc, tz = parse_timetz(s)
|
|
sc, micro = divmod(sc, 1.0)
|
|
micro = round(micro * 1000000)
|
|
if tz: tz = tzinfo(tz)
|
|
return time(hr, mn, int(sc), int(micro), tz)
|
|
|
|
def _conv_timestamp(s, cursor):
|
|
if s:
|
|
y, m, d, hr, mn, sc = parse_datetime(s)
|
|
sc, micro = divmod(sc, 1.0)
|
|
micro = round(micro * 1000000)
|
|
return datetime(y, m, d, hr, mn, int(sc), int(micro))
|
|
|
|
def _conv_timestamptz(s, cursor):
|
|
if s:
|
|
from zope.datetime import tzinfo
|
|
y, m, d, hr, mn, sc, tz = parse_datetimetz(s)
|
|
sc, micro = divmod(sc, 1.0)
|
|
micro = round(micro * 1000000)
|
|
if tz: tz = tzinfo(tz)
|
|
return datetime(y, m, d, hr, mn, int(sc), int(micro), tz)
|
|
|
|
def _conv_interval(s, cursor):
|
|
if s:
|
|
y, m, d, hr, mn, sc = parse_interval(s)
|
|
if (y, m) != (0, 0):
|
|
# XXX: Currently there's no way to represent years and months as
|
|
# timedeltas
|
|
return s
|
|
else:
|
|
return timedelta(days=d, hours=hr, minutes=mn, seconds=sc)
|
|
|
|
def _get_string_conv(encoding):
|
|
def _conv_string(s, cursor):
|
|
if s is not None:
|
|
s = s.decode(encoding)
|
|
return s
|
|
return _conv_string
|
|
|
|
# User-defined types
|
|
DATE = psycopg2.extensions.new_type((DATE_OID,), "ZDATE", _conv_date)
|
|
TIME = psycopg2.extensions.new_type((TIME_OID,), "ZTIME", _conv_time)
|
|
TIMETZ = psycopg2.extensions.new_type((TIMETZ_OID,), "ZTIMETZ", _conv_timetz)
|
|
TIMESTAMP = psycopg2.extensions.new_type((TIMESTAMP_OID,), "ZTIMESTAMP", _conv_timestamp)
|
|
TIMESTAMPTZ = psycopg2.extensions.new_type((TIMESTAMPTZ_OID,), "ZTIMESTAMPTZ", _conv_timestamptz)
|
|
INTERVAL = psycopg2.extensions.new_type((INTERVAL_OID,), "ZINTERVAL", _conv_interval)
|
|
|
|
def registerTypes(encoding):
|
|
"""Register type conversions for psycopg"""
|
|
psycopg2.extensions.register_type(DATE)
|
|
psycopg2.extensions.register_type(TIME)
|
|
psycopg2.extensions.register_type(TIMETZ)
|
|
psycopg2.extensions.register_type(TIMESTAMP)
|
|
psycopg2.extensions.register_type(TIMESTAMPTZ)
|
|
psycopg2.extensions.register_type(INTERVAL)
|
|
STRING = psycopg2.extensions.new_type((CHAR_OID, TEXT_OID, BPCHAR_OID, VARCHAR_OID), "ZSTRING", _get_string_conv(encoding))
|
|
psycopg2.extensions.register_type(STRING)
|
|
psycopg2.extensions.register_type(psycopg2.extensions.UNICODEARRAY)
|
|
|
|
|
|
dsn2option_mapping = {'host': 'host',
|
|
'port': 'port',
|
|
'dbname': 'dbname',
|
|
'username': 'user',
|
|
'password': 'password'}
|
|
|
|
class Psycopg2Adapter(ZopeDatabaseAdapter):
|
|
"""A psycopg2 adapter for Zope3.
|
|
|
|
The following type conversions are performed:
|
|
|
|
DATE -> datetime.date
|
|
TIME -> datetime.time
|
|
TIMETZ -> datetime.time
|
|
TIMESTAMP -> datetime.datetime
|
|
TIMESTAMPTZ -> datetime.datetime
|
|
|
|
XXX: INTERVAL cannot be represented exactly as datetime.timedelta since
|
|
it might be something like '1 month', which is a variable number of days.
|
|
"""
|
|
|
|
def connect(self):
|
|
if not self.isConnected():
|
|
try:
|
|
self._v_connection = Psycopg2Connection(
|
|
self._connection_factory(), self
|
|
)
|
|
except psycopg2.Error, error:
|
|
raise DatabaseException, str(error)
|
|
|
|
def registerTypes(self):
|
|
registerTypes(self.getEncoding())
|
|
|
|
def _connection_factory(self):
|
|
"""Create a psycopg2 DBI connection based on the DSN"""
|
|
self.registerTypes()
|
|
conn_info = parseDSN(self.dsn)
|
|
conn_list = []
|
|
for dsnname, optname in dsn2option_mapping.iteritems():
|
|
if conn_info[dsnname]:
|
|
conn_list.append('%s=%s' % (optname, conn_info[dsnname]))
|
|
conn_str = ' '.join(conn_list)
|
|
connection = psycopg2.connect(conn_str)
|
|
connection.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_SERIALIZABLE)
|
|
return connection
|
|
|
|
def disconnect(self):
|
|
if self.isConnected():
|
|
try:
|
|
self._v_connection.close()
|
|
except psycopg2.InterfaceError:
|
|
pass
|
|
self._v_connection = None
|
|
|
|
|
|
def _handle_psycopg_exception(error):
|
|
"""Called from a exception handler for psycopg2.Error.
|
|
|
|
If we have a serialization exception or a deadlock, we should retry the
|
|
transaction by raising a Retry exception. Otherwise, we reraise.
|
|
"""
|
|
if isinstance(error, psycopg2.extensions.TransactionRollbackError):
|
|
raise Retry(sys.exc_info())
|
|
raise
|
|
|
|
|
|
class IPsycopg2ZopeConnection(IZopeConnection):
|
|
"""A marker interface stating that this connection uses PostgreSQL."""
|
|
|
|
|
|
class Psycopg2Connection(ZopeConnection):
|
|
|
|
implements(IPsycopg2ZopeConnection)
|
|
|
|
def cursor(self):
|
|
"""See IZopeConnection"""
|
|
return Psycopg2Cursor(self.conn.cursor(), self)
|
|
|
|
def commit(self):
|
|
try:
|
|
ZopeConnection.commit(self)
|
|
except psycopg2.Error, error:
|
|
_handle_psycopg_exception(error)
|
|
|
|
|
|
class Psycopg2Cursor(ZopeCursor):
|
|
|
|
def execute(self, operation, parameters=None):
|
|
"""See IZopeCursor"""
|
|
try:
|
|
return ZopeCursor.execute(self, operation, parameters)
|
|
except psycopg2.Error, error:
|
|
_handle_psycopg_exception(error)
|
|
|
|
def executemany(operation, seq_of_parameters=None):
|
|
"""See IZopeCursor"""
|
|
raise RuntimeError, 'Oos'
|
|
try:
|
|
return ZopeCursor.execute(self, operation, seq_of_parameters)
|
|
except psycopg2.Error, error:
|
|
_handle_psycopg_exception(error)
|