Changeset View
Changeset View
Standalone View
Standalone View
swh/core/db/__init__.py
| Show First 20 Lines • Show All 46 Lines • ▼ Show 20 Lines | elif isinstance(data, psycopg2.extras.Range): | ||||
| ) | ) | ||||
| elif isinstance(data, enum.IntEnum): | elif isinstance(data, enum.IntEnum): | ||||
| return escape(int(data)) | return escape(int(data)) | ||||
| else: | else: | ||||
| # We don't escape here to make sure we pass literals properly | # We don't escape here to make sure we pass literals properly | ||||
| return str(data) | return str(data) | ||||
| def typecast_bytea(value, cur): | |||||
| if value is not None: | |||||
| data = psycopg2.BINARY(value, cur) | |||||
| return data.tobytes() | |||||
| 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 | ||||
| def adapt_conn(cls, conn): | |||||
| """Makes psycopg2 use 'bytes' to decode bytea instead of | |||||
| 'memoryview', for this connection.""" | |||||
| cur = conn.cursor() | |||||
| cur.execute("SELECT null::bytea, null::bytea[]") | |||||
| bytea_oid = cur.description[0][1] | |||||
| bytea_array_oid = cur.description[1][1] | |||||
| t_bytes = psycopg2.extensions.new_type( | |||||
| (bytea_oid,), "bytea", typecast_bytea) | |||||
| psycopg2.extensions.register_type(t_bytes, conn) | |||||
| t_bytes_array = psycopg2.extensions.new_array_type( | |||||
| (bytea_array_oid,), "bytea[]", t_bytes) | |||||
| psycopg2.extensions.register_type(t_bytes_array, conn) | |||||
| @classmethod | |||||
| def connect(cls, *args, **kwargs): | def connect(cls, *args, **kwargs): | ||||
| """factory method to create a DB proxy | """factory method to create a DB proxy | ||||
| Accepts all arguments of psycopg2.connect; only some specific | Accepts all arguments of psycopg2.connect; only some specific | ||||
| possibilities are reported below. | possibilities are reported below. | ||||
| Args: | Args: | ||||
| connstring: libpq2 connection string | connstring: libpq2 connection string | ||||
| """ | """ | ||||
| conn = psycopg2.connect(*args, **kwargs) | conn = psycopg2.connect(*args, **kwargs) | ||||
| cls.adapt_conn(conn) | |||||
| return cls(conn) | return cls(conn) | ||||
| @classmethod | @classmethod | ||||
| def from_pool(cls, pool): | def from_pool(cls, pool): | ||||
| return cls(pool.getconn(), pool=pool) | conn = pool.getconn() | ||||
| cls.adapt_conn(conn) | |||||
| return cls(conn, pool=pool) | |||||
| def __init__(self, conn, pool=None): | def __init__(self, conn, pool=None): | ||||
| """create a DB proxy | """create a DB proxy | ||||
| Args: | Args: | ||||
| conn: psycopg2 connection to the SWH DB | conn: psycopg2 connection to the SWH DB | ||||
| pool: psycopg2 pool of connections | pool: psycopg2 pool of connections | ||||
| ▲ Show 20 Lines • Show All 81 Lines • Show Last 20 Lines | |||||