Changeset View
Changeset View
Standalone View
Standalone View
swh/core/db/__init__.py
Show All 13 Lines | |||||
import psycopg2 | import psycopg2 | ||||
import psycopg2.extras | import psycopg2.extras | ||||
psycopg2.extras.register_uuid() | psycopg2.extras.register_uuid() | ||||
def escape(data): | |||||
if data is None: | |||||
return '' | |||||
if isinstance(data, bytes): | |||||
return '\\x%s' % binascii.hexlify(data).decode('ascii') | |||||
elif isinstance(data, str): | |||||
return '"%s"' % data.replace('"', '""') | |||||
elif isinstance(data, datetime.datetime): | |||||
# We escape twice to make sure the string generated by | |||||
# isoformat gets escaped | |||||
return escape(data.isoformat()) | |||||
elif isinstance(data, dict): | |||||
return escape(json.dumps(data)) | |||||
elif isinstance(data, list): | |||||
return escape("{%s}" % ','.join(escape(d) for d in data)) | |||||
elif isinstance(data, psycopg2.extras.Range): | |||||
# We escape twice here too, so that we make sure | |||||
# everything gets passed to copy properly | |||||
return escape( | |||||
'%s%s,%s%s' % ( | |||||
'[' if data.lower_inc else '(', | |||||
'-infinity' if data.lower_inf else escape(data.lower), | |||||
'infinity' if data.upper_inf else escape(data.upper), | |||||
']' if data.upper_inc else ')', | |||||
) | |||||
) | |||||
elif isinstance(data, enum.IntEnum): | |||||
return escape(int(data)) | |||||
else: | |||||
# We don't escape here to make sure we pass literals properly | |||||
return str(data) | |||||
class BaseDb: | class BaseDb: | ||||
"""Base class for swh.*.*Db. | """Base class for swh.*.*Db. | ||||
cf. swh.storage.db.Db, swh.archiver.db.ArchiverDb | cf. swh.storage.db.Db, swh.archiver.db.ArchiverDb | ||||
""" | """ | ||||
@classmethod | @classmethod | ||||
▲ Show 20 Lines • Show All 64 Lines • ▼ Show 20 Lines | def copy_to(self, items, tblname, columns, cur=None, item_cb=None): | ||||
Args: | Args: | ||||
items (dict): dictionary of data to copy over tblname | items (dict): dictionary of data to copy over tblname | ||||
tblname (str): Destination table's name | tblname (str): Destination table's name | ||||
columns ([str]): keys to access data in items and also the | columns ([str]): keys to access data in items and also the | ||||
column names in the destination table. | column names in the destination table. | ||||
item_cb (fn): optional function to apply to items's entry | item_cb (fn): optional function to apply to items's entry | ||||
""" | """ | ||||
def escape(data): | |||||
if data is None: | |||||
return '' | |||||
if isinstance(data, bytes): | |||||
return '\\x%s' % binascii.hexlify(data).decode('ascii') | |||||
elif isinstance(data, str): | |||||
return '"%s"' % data.replace('"', '""') | |||||
elif isinstance(data, datetime.datetime): | |||||
# We escape twice to make sure the string generated by | |||||
# isoformat gets escaped | |||||
return escape(data.isoformat()) | |||||
elif isinstance(data, dict): | |||||
return escape(json.dumps(data)) | |||||
elif isinstance(data, list): | |||||
return escape("{%s}" % ','.join(escape(d) for d in data)) | |||||
elif isinstance(data, psycopg2.extras.Range): | |||||
# We escape twice here too, so that we make sure | |||||
# everything gets passed to copy properly | |||||
return escape( | |||||
'%s%s,%s%s' % ( | |||||
'[' if data.lower_inc else '(', | |||||
'-infinity' if data.lower_inf else escape(data.lower), | |||||
'infinity' if data.upper_inf else escape(data.upper), | |||||
']' if data.upper_inc else ')', | |||||
) | |||||
) | |||||
elif isinstance(data, enum.IntEnum): | |||||
return escape(int(data)) | |||||
else: | |||||
# We don't escape here to make sure we pass literals properly | |||||
return str(data) | |||||
read_file, write_file = os.pipe() | read_file, write_file = os.pipe() | ||||
def writer(): | def writer(): | ||||
cursor = self._cursor(cur) | cursor = self._cursor(cur) | ||||
with open(read_file, 'r') as f: | with open(read_file, 'r') as f: | ||||
cursor.copy_expert('COPY %s (%s) FROM STDIN CSV' % ( | cursor.copy_expert('COPY %s (%s) FROM STDIN CSV' % ( | ||||
tblname, ', '.join(columns)), f) | tblname, ', '.join(columns)), f) | ||||
Show All 20 Lines |