| 1 | import re |
|---|
| 2 | import os.path |
|---|
| 3 | from .._compat import PY2, with_metaclass, iterkeys, to_unicode, long |
|---|
| 4 | from .._globals import IDENTITY, THREAD_LOCAL |
|---|
| 5 | from ..drivers import psycopg2_adapt |
|---|
| 6 | from .base import SQLAdapter |
|---|
| 7 | from ..utils import split_uri_args |
|---|
| 8 | from . import AdapterMeta, adapters, with_connection, with_connection_or_raise |
|---|
| 9 | |
|---|
| 10 | |
|---|
| 11 | class PostgreMeta(AdapterMeta): |
|---|
| 12 | def __call__(cls, *args, **kwargs): |
|---|
| 13 | if cls not in [Postgre, PostgreNew, PostgreBoolean]: |
|---|
| 14 | return AdapterMeta.__call__(cls, *args, **kwargs) |
|---|
| 15 | # choose driver according uri |
|---|
| 16 | available_drivers = [ |
|---|
| 17 | driver |
|---|
| 18 | for driver in cls.drivers |
|---|
| 19 | if driver in iterkeys(kwargs["db"]._drivers_available) |
|---|
| 20 | ] |
|---|
| 21 | uri_items = kwargs["uri"].split("://", 1)[0].split(":") |
|---|
| 22 | uri_driver = uri_items[1] if len(uri_items) > 1 else None |
|---|
| 23 | if uri_driver and uri_driver in available_drivers: |
|---|
| 24 | driver = uri_driver |
|---|
| 25 | else: |
|---|
| 26 | driver = available_drivers[0] if available_drivers else cls.drivers[0] |
|---|
| 27 | cls = adapters._registry_[uri_items[0] + ":" + driver] |
|---|
| 28 | return AdapterMeta.__call__(cls, *args, **kwargs) |
|---|
| 29 | |
|---|
| 30 | |
|---|
| 31 | @adapters.register_for("postgres") |
|---|
| 32 | class Postgre(with_metaclass(PostgreMeta, SQLAdapter)): |
|---|
| 33 | dbengine = "postgres" |
|---|
| 34 | drivers = ("psycopg2",) |
|---|
| 35 | support_distributed_transaction = True |
|---|
| 36 | |
|---|
| 37 | REGEX_URI = ( |
|---|
| 38 | "^(?P<user>[^:@]+)(:(?P<password>[^@]*))?" |
|---|
| 39 | r"@(?P<host>[^:/]*|\[[^\]]+\])(:(?P<port>\d+))?" |
|---|
| 40 | "/(?P<db>[^?]+)" |
|---|
| 41 | r"(\?(?P<uriargs>.*))?$" |
|---|
| 42 | ) # sslmode, ssl (no value) and unix_socket |
|---|
| 43 | |
|---|
| 44 | def __init__( |
|---|
| 45 | self, |
|---|
| 46 | db, |
|---|
| 47 | uri, |
|---|
| 48 | pool_size=0, |
|---|
| 49 | folder=None, |
|---|
| 50 | db_codec="UTF-8", |
|---|
| 51 | credential_decoder=IDENTITY, |
|---|
| 52 | driver_args={}, |
|---|
| 53 | adapter_args={}, |
|---|
| 54 | srid=4326, |
|---|
| 55 | after_connection=None, |
|---|
| 56 | ): |
|---|
| 57 | self.srid = srid |
|---|
| 58 | super(Postgre, self).__init__( |
|---|
| 59 | db, |
|---|
| 60 | uri, |
|---|
| 61 | pool_size, |
|---|
| 62 | folder, |
|---|
| 63 | db_codec, |
|---|
| 64 | credential_decoder, |
|---|
| 65 | driver_args, |
|---|
| 66 | adapter_args, |
|---|
| 67 | after_connection, |
|---|
| 68 | ) |
|---|
| 69 | self._config_json() |
|---|
| 70 | |
|---|
| 71 | def _initialize_(self): |
|---|
| 72 | super(Postgre, self)._initialize_() |
|---|
| 73 | ruri = self.uri.split("://", 1)[1] |
|---|
| 74 | m = re.match(self.REGEX_URI, ruri) |
|---|
| 75 | if not m: |
|---|
| 76 | raise SyntaxError("Invalid URI string in DAL") |
|---|
| 77 | user = self.credential_decoder(m.group("user")) |
|---|
| 78 | password = self.credential_decoder(m.group("password")) |
|---|
| 79 | host = m.group("host") |
|---|
| 80 | uriargs = m.group("uriargs") |
|---|
| 81 | if uriargs: |
|---|
| 82 | uri_args = split_uri_args(uriargs, need_equal=False) |
|---|
| 83 | else: |
|---|
| 84 | uri_args = dict() |
|---|
| 85 | socket = uri_args.get("unix_socket") |
|---|
| 86 | if not host and not socket: |
|---|
| 87 | raise SyntaxError("Host or UNIX socket name required") |
|---|
| 88 | db = m.group("db") |
|---|
| 89 | self.driver_args.update(user=user, database=db) |
|---|
| 90 | if password is not None: |
|---|
| 91 | self.driver_args["password"] = password |
|---|
| 92 | if socket: |
|---|
| 93 | if not os.path.exists(socket): |
|---|
| 94 | raise ValueError("UNIX socket %r not found" % socket) |
|---|
| 95 | if self.driver_name == "psycopg2": |
|---|
| 96 | # the psycopg2 driver let you configure the socket directory |
|---|
| 97 | # only (not the socket file name) by passing it as the host |
|---|
| 98 | # (must be an absolute path otherwise the driver tries a TCP/IP |
|---|
| 99 | # connection to host); this behaviour is due to the underlying |
|---|
| 100 | # libpq used by the driver |
|---|
| 101 | socket_dir = os.path.abspath(os.path.dirname(socket)) |
|---|
| 102 | self.driver_args["host"] = socket_dir |
|---|
| 103 | else: |
|---|
| 104 | port = int(m.group("port") or 5432) |
|---|
| 105 | self.driver_args.update(host=host, port=port) |
|---|
| 106 | sslmode = uri_args.get("sslmode") |
|---|
| 107 | if sslmode and self.driver_name == "psycopg2": |
|---|
| 108 | self.driver_args["sslmode"] = sslmode |
|---|
| 109 | if self.driver: |
|---|
| 110 | self.__version__ = "%s %s" % (self.driver.__name__, self.driver.__version__) |
|---|
| 111 | else: |
|---|
| 112 | self.__version__ = None |
|---|
| 113 | THREAD_LOCAL._pydal_last_insert_ = None |
|---|
| 114 | self.get_connection() |
|---|
| 115 | |
|---|
| 116 | def _get_json_dialect(self): |
|---|
| 117 | from ..dialects.postgre import PostgreDialectJSON |
|---|
| 118 | |
|---|
| 119 | return PostgreDialectJSON |
|---|
| 120 | |
|---|
| 121 | def _get_json_parser(self): |
|---|
| 122 | from ..parsers.postgre import PostgreAutoJSONParser |
|---|
| 123 | |
|---|
| 124 | return PostgreAutoJSONParser |
|---|
| 125 | |
|---|
| 126 | @property |
|---|
| 127 | def _last_insert(self): |
|---|
| 128 | return THREAD_LOCAL._pydal_last_insert_ |
|---|
| 129 | |
|---|
| 130 | @_last_insert.setter |
|---|
| 131 | def _last_insert(self, value): |
|---|
| 132 | THREAD_LOCAL._pydal_last_insert_ = value |
|---|
| 133 | |
|---|
| 134 | def connector(self): |
|---|
| 135 | return self.driver.connect(**self.driver_args) |
|---|
| 136 | |
|---|
| 137 | def after_connection(self): |
|---|
| 138 | self.execute("SET CLIENT_ENCODING TO 'UTF8'") |
|---|
| 139 | self.execute("SET standard_conforming_strings=on;") |
|---|
| 140 | |
|---|
| 141 | def lastrowid(self, table): |
|---|
| 142 | if self._last_insert: |
|---|
| 143 | return long(self.cursor.fetchone()[0]) |
|---|
| 144 | sequence_name = table._sequence_name |
|---|
| 145 | self.execute("SELECT currval(%s);" % self.adapt(sequence_name)) |
|---|
| 146 | return long(self.cursor.fetchone()[0]) |
|---|
| 147 | |
|---|
| 148 | def _insert(self, table, fields): |
|---|
| 149 | self._last_insert = None |
|---|
| 150 | if fields: |
|---|
| 151 | retval = None |
|---|
| 152 | if hasattr(table, "_id"): |
|---|
| 153 | self._last_insert = (table._id, 1) |
|---|
| 154 | retval = table._id._rname |
|---|
| 155 | return self.dialect.insert( |
|---|
| 156 | table._rname, |
|---|
| 157 | ",".join(el[0]._rname for el in fields), |
|---|
| 158 | ",".join(self.expand(v, f.type) for f, v in fields), |
|---|
| 159 | retval, |
|---|
| 160 | ) |
|---|
| 161 | return self.dialect.insert_empty(table._rname) |
|---|
| 162 | |
|---|
| 163 | @with_connection |
|---|
| 164 | def prepare(self, key): |
|---|
| 165 | self.execute("PREPARE TRANSACTION '%s';" % key) |
|---|
| 166 | |
|---|
| 167 | @with_connection |
|---|
| 168 | def commit_prepared(self, key): |
|---|
| 169 | self.execute("COMMIT PREPARED '%s';" % key) |
|---|
| 170 | |
|---|
| 171 | @with_connection |
|---|
| 172 | def rollback_prepared(self, key): |
|---|
| 173 | self.execute("ROLLBACK PREPARED '%s';" % key) |
|---|
| 174 | |
|---|
| 175 | |
|---|
| 176 | @adapters.register_for("postgres:psycopg2") |
|---|
| 177 | class PostgrePsyco(Postgre): |
|---|
| 178 | drivers = ("psycopg2",) |
|---|
| 179 | |
|---|
| 180 | def _config_json(self): |
|---|
| 181 | use_json = ( |
|---|
| 182 | self.driver.__version__ >= "2.0.12" |
|---|
| 183 | and self.connection.server_version >= 90200 |
|---|
| 184 | ) |
|---|
| 185 | if use_json: |
|---|
| 186 | self.dialect = self._get_json_dialect()(self) |
|---|
| 187 | if self.driver.__version__ >= "2.5.0": |
|---|
| 188 | self.parser = self._get_json_parser()(self) |
|---|
| 189 | |
|---|
| 190 | def adapt(self, obj): |
|---|
| 191 | adapted = psycopg2_adapt(obj) |
|---|
| 192 | # deal with new relic Connection Wrapper (newrelic>=2.10.0.8) |
|---|
| 193 | cxn = getattr(self.connection, "__wrapped__", self.connection) |
|---|
| 194 | adapted.prepare(cxn) |
|---|
| 195 | rv = adapted.getquoted() |
|---|
| 196 | if not PY2: |
|---|
| 197 | if isinstance(rv, bytes): |
|---|
| 198 | return rv.decode("utf-8") |
|---|
| 199 | return rv |
|---|
| 200 | |
|---|
| 201 | |
|---|
| 202 | @adapters.register_for("postgres2") |
|---|
| 203 | class PostgreNew(Postgre): |
|---|
| 204 | def _get_json_dialect(self): |
|---|
| 205 | from ..dialects.postgre import PostgreDialectArraysJSON |
|---|
| 206 | |
|---|
| 207 | return PostgreDialectArraysJSON |
|---|
| 208 | |
|---|
| 209 | def _get_json_parser(self): |
|---|
| 210 | from ..parsers.postgre import PostgreNewAutoJSONParser |
|---|
| 211 | |
|---|
| 212 | return PostgreNewAutoJSONParser |
|---|
| 213 | |
|---|
| 214 | |
|---|
| 215 | @adapters.register_for("postgres2:psycopg2") |
|---|
| 216 | class PostgrePsycoNew(PostgrePsyco, PostgreNew): |
|---|
| 217 | pass |
|---|
| 218 | |
|---|
| 219 | |
|---|
| 220 | @adapters.register_for("postgres3") |
|---|
| 221 | class PostgreBoolean(PostgreNew): |
|---|
| 222 | def _get_json_dialect(self): |
|---|
| 223 | from ..dialects.postgre import PostgreDialectBooleanJSON |
|---|
| 224 | |
|---|
| 225 | return PostgreDialectBooleanJSON |
|---|
| 226 | |
|---|
| 227 | def _get_json_parser(self): |
|---|
| 228 | from ..parsers.postgre import PostgreBooleanAutoJSONParser |
|---|
| 229 | |
|---|
| 230 | return PostgreBooleanAutoJSONParser |
|---|
| 231 | |
|---|
| 232 | |
|---|
| 233 | @adapters.register_for("postgres3:psycopg2") |
|---|
| 234 | class PostgrePsycoBoolean(PostgrePsycoNew, PostgreBoolean): |
|---|
| 235 | pass |
|---|
| 236 | |
|---|
| 237 | |
|---|
| 238 | @adapters.register_for("jdbc:postgres") |
|---|
| 239 | class JDBCPostgre(Postgre): |
|---|
| 240 | drivers = ("zxJDBC",) |
|---|
| 241 | |
|---|
| 242 | REGEX_URI = ( |
|---|
| 243 | "^(?P<user>[^:@]+)(:(?P<password>[^@]*))?" |
|---|
| 244 | r"@(?P<host>[^:/]+|\[[^\]]+\])(:(?P<port>\d+))?" |
|---|
| 245 | "/(?P<db>[^?]+)$" |
|---|
| 246 | ) |
|---|
| 247 | |
|---|
| 248 | def _initialize_(self): |
|---|
| 249 | super(Postgre, self)._initialize_() |
|---|
| 250 | ruri = self.uri.split("://", 1)[1] |
|---|
| 251 | m = re.match(self.REGEX_URI, ruri) |
|---|
| 252 | if not m: |
|---|
| 253 | raise SyntaxError("Invalid URI string in DAL") |
|---|
| 254 | user = self.credential_decoder(m.group("user")) |
|---|
| 255 | password = self.credential_decoder(m.group("password")) |
|---|
| 256 | if password is None: |
|---|
| 257 | password = "" |
|---|
| 258 | host = m.group("host") |
|---|
| 259 | db = m.group("db") |
|---|
| 260 | port = m.group("port") or "5432" |
|---|
| 261 | self.dsn = ("jdbc:postgresql://%s:%s/%s" % (host, port, db), user, password) |
|---|
| 262 | if self.driver: |
|---|
| 263 | self.__version__ = "%s %s" % (self.driver.__name__, self.driver.__version__) |
|---|
| 264 | else: |
|---|
| 265 | self.__version__ = None |
|---|
| 266 | THREAD_LOCAL._pydal_last_insert_ = None |
|---|
| 267 | self.get_connection() |
|---|
| 268 | |
|---|
| 269 | def connector(self): |
|---|
| 270 | return self.driver.connect(*self.dsn, **self.driver_args) |
|---|
| 271 | |
|---|
| 272 | def after_connection(self): |
|---|
| 273 | self.connection.set_client_encoding("UTF8") |
|---|
| 274 | self.execute("BEGIN;") |
|---|
| 275 | self.execute("SET CLIENT_ENCODING TO 'UNICODE';") |
|---|
| 276 | |
|---|
| 277 | def _config_json(self): |
|---|
| 278 | use_json = self.connection.dbversion >= "9.2.0" |
|---|
| 279 | if use_json: |
|---|
| 280 | self.dialect = self._get_json_dialect()(self) |
|---|
