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 | |||||