Changeset View
Changeset View
Standalone View
Standalone View
swh/core/db/db_utils.py
Show All 20 Lines | |||||
def stored_procedure(stored_proc): | def stored_procedure(stored_proc): | ||||
"""decorator to execute remote stored procedure, specified as argument | """decorator to execute remote stored procedure, specified as argument | ||||
Generally, the body of the decorated function should be empty. If it is | Generally, the body of the decorated function should be empty. If it is | ||||
not, the stored procedure will be executed first; the function body then. | not, the stored procedure will be executed first; the function body then. | ||||
""" | """ | ||||
def wrap(meth): | def wrap(meth): | ||||
@functools.wraps(meth) | @functools.wraps(meth) | ||||
def _meth(self, *args, **kwargs): | def _meth(self, *args, **kwargs): | ||||
cur = kwargs.get('cur', None) | cur = kwargs.get('cur', None) | ||||
self._cursor(cur).execute('SELECT %s()' % stored_proc) | self._cursor(cur).execute('SELECT %s()' % stored_proc) | ||||
meth(self, *args, **kwargs) | meth(self, *args, **kwargs) | ||||
return _meth | return _meth | ||||
return wrap | return wrap | ||||
def jsonize(value): | def jsonize(value): | ||||
"""Convert a value to a psycopg2 JSON object if necessary""" | """Convert a value to a psycopg2 JSON object if necessary""" | ||||
if isinstance(value, dict): | if isinstance(value, dict): | ||||
return psycopg2.extras.Json(value) | return psycopg2.extras.Json(value) | ||||
Show All 30 Lines | for token in tokens: | ||||
if len(token) != 2 or token[:1] != b'%': | if len(token) != 2 or token[:1] != b'%': | ||||
curr.append(token) | curr.append(token) | ||||
continue | continue | ||||
if token[1:] == b's': | if token[1:] == b's': | ||||
if curr is pre: | if curr is pre: | ||||
curr = post | curr = post | ||||
else: | else: | ||||
raise ValueError( | raise ValueError("the query contains more than one '%s' placeholder") | ||||
"the query contains more than one '%s' placeholder") | |||||
elif token[1:] == b'%': | elif token[1:] == b'%': | ||||
curr.append(b'%') | curr.append(b'%') | ||||
else: | else: | ||||
raise ValueError("unsupported format character: '%s'" | raise ValueError( | ||||
% token[1:].decode('ascii', 'replace')) | "unsupported format character: '%s'" | ||||
% token[1:].decode('ascii', 'replace') | |||||
) | |||||
if curr is pre: | if curr is pre: | ||||
raise ValueError("the query doesn't contain any '%s' placeholder") | raise ValueError("the query doesn't contain any '%s' placeholder") | ||||
return pre, post | return pre, post | ||||
def execute_values_generator(cur, sql, argslist, template=None, page_size=100): | def execute_values_generator(cur, sql, argslist, template=None, page_size=100): | ||||
Show All 30 Lines | def execute_values_generator(cur, sql, argslist, template=None, page_size=100): | ||||
After the execution of the function the `cursor.rowcount` property will | After the execution of the function the `cursor.rowcount` property will | ||||
**not** contain a total result. | **not** contain a total result. | ||||
''' | ''' | ||||
# we can't just use sql % vals because vals is bytes: if sql is bytes | # we can't just use sql % vals because vals is bytes: if sql is bytes | ||||
# there will be some decoding error because of stupid codec used, and Py3 | # there will be some decoding error because of stupid codec used, and Py3 | ||||
# doesn't implement % on bytes. | # doesn't implement % on bytes. | ||||
if not isinstance(sql, bytes): | if not isinstance(sql, bytes): | ||||
sql = sql.encode( | sql = sql.encode(psycopg2.extensions.encodings[cur.connection.encoding]) | ||||
psycopg2.extensions.encodings[cur.connection.encoding] | |||||
) | |||||
pre, post = _split_sql(sql) | pre, post = _split_sql(sql) | ||||
for page in _paginate(argslist, page_size=page_size): | for page in _paginate(argslist, page_size=page_size): | ||||
if template is None: | if template is None: | ||||
template = b'(' + b','.join([b'%s'] * len(page[0])) + b')' | template = b'(' + b','.join([b'%s'] * len(page[0])) + b')' | ||||
parts = pre[:] | parts = pre[:] | ||||
for args in page: | for args in page: | ||||
parts.append(cur.mogrify(template, args)) | parts.append(cur.mogrify(template, args)) | ||||
parts.append(b',') | parts.append(b',') | ||||
parts[-1:] = post | parts[-1:] = post | ||||
cur.execute(b''.join(parts)) | cur.execute(b''.join(parts)) | ||||
yield from cur | yield from cur |