def get_cursor(cls, cursor_type=_CursorType.PLAIN, use_replica=False) -> Cursor: """ Yields: new client-side cursor from existing db connection pool """ if cls._use_pool: _connection_source = yield from cls.get_pool(use_replica) else: if use_replica: _connection_source = yield from aiopg.connect( echo=False, **cls_replica_connection_params) else: _connection_source = yield from aiopg.connect( echo=False, **cls._connection_params) if cursor_type == _CursorType.PLAIN: _cur = yield from _connection_source.cursor() if cursor_type == _CursorType.NAMEDTUPLE: _cur = yield from _connection_source.cursor( cursor_factory=psycopg2.extras.NamedTupleCursor) if cursor_type == _CursorType.DICT: _cur = yield from _connection_source.cursor( cursor_factory=psycopg2.extras.DictCursor) if not cls._use_pool: _cur = cursor_context_manager(_connection_source, _cur) return _cur
async def test_concurrency_dispatcher(self): model = FakeModel("foo") identifier = uuid4() broker_publisher = BrokerPublisher.from_config(config=self.config) async with broker_publisher: for x in range(60): await broker_publisher.send( model, "CommandBroker-Delete", identifier=identifier, reply_topic="TestDeleteReply" ) async with aiopg.connect(**self.broker_queue_db) as connect: async with connect.cursor() as cur: await cur.execute("SELECT COUNT(*) FROM producer_queue") records = await cur.fetchone() assert records[0] == 60 await asyncio.gather(*(self.producer.dispatch() for _ in range(6))) async with aiopg.connect(**self.broker_queue_db) as connect: async with connect.cursor() as cur: await cur.execute("SELECT COUNT(*) FROM producer_queue") records = await cur.fetchone() assert records[0] == 0
def test_instrumentor_connect(self): AiopgInstrumentor().instrument() cnx = async_call(aiopg.connect(database="test")) cursor = async_call(cnx.cursor()) query = "SELECT * FROM test" async_call(cursor.execute(query)) spans_list = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans_list), 1) span = spans_list[0] # Check version and name in span's instrumentation info self.check_span_instrumentation_info( span, opentelemetry.instrumentation.aiopg) # check that no spans are generated after uninstrument AiopgInstrumentor().uninstrument() cnx = async_call(aiopg.connect(database="test")) cursor = async_call(cnx.cursor()) query = "SELECT * FROM test" cursor.execute(query) spans_list = self.memory_exporter.get_finished_spans() self.assertEqual(len(spans_list), 1)
def go(*, no_loop=False, **kwargs): nonlocal conn pg_params.update(kwargs) useloop = None if no_loop else loop conn = yield from aiopg.connect(loop=useloop, **pg_params) conn2 = yield from aiopg.connect(loop=useloop, **pg_params) cur = yield from conn2.cursor() yield from cur.execute("DROP TABLE IF EXISTS foo") yield from conn2.close() return conn
def test_unwrap_connect(self): wrappers.wrap_connect(self.tracer, "-") aiopg_mock = AiopgMock() with mock.patch("aiopg.connect", aiopg_mock.connect): connection = async_call(aiopg.connect()) self.assertEqual(aiopg_mock.connect_call_count, 1) wrappers.unwrap_connect() connection = async_call(aiopg.connect()) self.assertEqual(aiopg_mock.connect_call_count, 2) self.assertIsInstance(connection, mock.Mock)
def connect(self, no_loop=False, **kwargs): loop = None if no_loop else self.loop conn = yield from aiopg.connect( database="aiopg", user="******", password="******", host="127.0.0.1", loop=loop, **kwargs ) conn2 = yield from aiopg.connect(database="aiopg", user="******", password="******", host="127.0.0.1", loop=loop) cur = yield from conn2.cursor() yield from cur.execute("DROP TABLE IF EXISTS foo") yield from conn2.close() self.addCleanup(conn.close) return conn
async def _is_processed(self, queue_id): async with aiopg.connect(**self.broker_queue_db) as connect: async with connect.cursor() as cur: await cur.execute( "SELECT COUNT(*) FROM consumer_queue WHERE id=%d" % (queue_id, )) return (await cur.fetchone())[0] == 0
async def put_comment(request, comment_id): query = """ UPDATE public.comments SET title = %(title)s, level = %(level)s, parent_id = %(parent_id)s WHERE id = %(comment_id)s RETURNING id, post_id, title, level, parent_id, created_at;""" if not (request.get("level") or request.get("parent_id")): comment = await get_comment_by_id(comment_id) if comment: request["title"] = comment.get("title") request["level"] = comment.get("level") else: return None params = dict( title=request.get("title"), level=request.get("level"), parent_id=request.get("parent_id"), comment_id=comment_id ) async with aiopg.connect(DB_URL) as conn: async with conn.cursor(cursor_factory=DictCursor) as cur: await cur.execute(query, params) data = await cur.fetchone() if data: return dict(data) return None
def get_cursor(self, cursor_type=_CursorType.PLAIN) -> Cursor: """ Yields: new client-side cursor from existing db connection pool """ if self._use_pool: yield from self.initialize_pool() _connection_source = self._pool else: _connection_source = yield from aiopg.connect( echo=False, **self._connection_params) if cursor_type == _CursorType.PLAIN: _cur = yield from _connection_source.cursor() if cursor_type == _CursorType.NAMEDTUPLE: _cur = yield from _connection_source.cursor( cursor_factory=psycopg2.extras.NamedTupleCursor) if cursor_type == _CursorType.DICT: _cur = yield from _connection_source.cursor( cursor_factory=psycopg2.extras.DictCursor) if not self._use_pool: _cur = cursor_context_manager(_connection_source, _cur) return _cur
def connect(self): """Connect asynchronously. """ self._conn = yield from aiopg.connect(dsn=self.dsn, timeout=self.timeout, loop=self._loop, **self.connect_kwargs)
def get_connection(dsn="", **kwargs) -> Awaitable[aiopg.Connection]: # tell aiopg not to register adapters for hstore & json by default, as # those are registered at the module level and could overwrite previously # defined adapters kwargs.setdefault("enable_json", False) kwargs.setdefault("enable_hstore", False) return aiopg.connect(dsn=dsn, **kwargs)
def test_custom_tracer_provider_instrument_connection(self): resource = resources.Resource.create( {"service.name": "db-test-service"} ) result = self.create_tracer_provider(resource=resource) tracer_provider, exporter = result cnx = async_call(aiopg.connect(database="test")) cnx = AiopgInstrumentor().instrument_connection( cnx, tracer_provider=tracer_provider ) cursor = async_call(cnx.cursor()) query = "SELECT * FROM test" async_call(cursor.execute(query)) spans_list = exporter.get_finished_spans() self.assertEqual(len(spans_list), 1) span = spans_list[0] self.assertEqual( span.resource.attributes["service.name"], "db-test-service" ) self.assertIs(span.resource, resource)
async def test_if_commands_retry_was_incremented(self): model = FakeModel("foo") identifier = uuid4() async with BrokerPublisher.from_config(config=self.config) as broker_publisher: await broker_publisher.send( model, "TestDeleteOrderReply", identifier=identifier, status=BrokerMessageStatus.SUCCESS ) await broker_publisher.send( model, "TestDeleteOrderReply", identifier=identifier, status=BrokerMessageStatus.SUCCESS ) self.producer.publish = AsyncMock(return_value=False) await self.producer.dispatch() async with aiopg.connect(**self.broker_queue_db) as connection: async with connection.cursor() as cursor: await cursor.execute("SELECT COUNT(*) FROM producer_queue WHERE topic = 'TestDeleteOrderReply'") self.assertEqual(2, (await cursor.fetchone())[0]) await cursor.execute("SELECT retry FROM producer_queue WHERE id=1;") self.assertEqual(1, (await cursor.fetchone())[0]) await cursor.execute("SELECT retry FROM producer_queue WHERE id=2;") self.assertEqual(1, (await cursor.fetchone())[0])
async def put_post(request, post_id): query = """ UPDATE public.posts SET title = %(title)s, description = %(description)s, updated_at = %(updated_at)s WHERE id = %(post_id)s RETURNING id, section_id, title, description, created_at, updated_at;""" search_query = """ UPDATE posts_search SET title = to_tsvector(%(title)s) WHERE post_id = %(post_id)s;""" params = dict( title=request.get("title"), description=request.get("description"), updated_at=datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), post_id=post_id) search_params = dict(post_id=post_id, title=request.get("title")) async with aiopg.connect(DB_URL) as conn: async with conn.cursor(cursor_factory=DictCursor) as cur: await cur.execute(query, params) data = await cur.fetchone() if data: await cur.execute(search_query, search_params) return dict(data) return None
async def post_section(request): query = """ INSERT INTO sections (title, description, created_at, updated_at) VALUES (%(title)s, %(description)s, %(created_at)s, %(updated_at)s) RETURNING id, title, description, created_at, updated_at;""" search_query = """ INSERT INTO public.sections_search (section_id, title) VALUES (%(section_id)s, to_tsvector(%(title)s)) RETURNING section_id;""" params = dict( title=request.get("title", None), description=request.get("description", None), created_at=datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), updated_at=datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') ) async with aiopg.connect(DB_URL) as conn: async with conn.cursor(cursor_factory=DictCursor) as cur: await cur.execute(query, params) data = await cur.fetchone() if data: serach_params = dict( section_id=data.get("id"), title=data.get("title") ) await cur.execute(search_query, serach_params) return dict(data) return None
def connect(self, **kwargs): conn = (yield from aiopg.connect(database='aiopg', user='******', password='******', host='127.0.0.1', loop=self.loop, **kwargs)) self.addCleanup(conn.close) cur = yield from conn.cursor() yield from cur.execute("DROP TABLE IF EXISTS tbl") yield from cur.execute("CREATE TABLE tbl (id int, name varchar(255))") for i in [(1, 'a'), (2, 'b'), (3, 'c')]: yield from cur.execute("INSERT INTO tbl VALUES(%s, %s)", i) yield from cur.execute("DROP TABLE IF EXISTS tbl2") yield from cur.execute("""CREATE TABLE tbl2 (id int, name varchar(255)) WITH OIDS""") yield from cur.execute("DROP FUNCTION IF EXISTS inc(val integer)") yield from cur.execute("""CREATE FUNCTION inc(val integer) RETURNS integer AS $$ BEGIN RETURN val + 1; END; $$ LANGUAGE PLPGSQL;""") return conn
async def get_all_users(): query = 'select user_id, name, age, phone from users limit 100 ;' async with aiopg.connect(DB_URL) as conn: async with conn.cursor(cursor_factory=DictCursor) as cur: await cur.execute(query) data = await cur.fetchall() return [dict(u) for u in data]
def connect(): conn = aiopg.connect( user=db_url.username, password=db_url.password, host=db_url.host, dbname=db_url.database, ) return conn
async def link_list(self): async with aiopg.connect(PGCONF) as conn: cur = await conn.cursor() await cur.execute( f"SELECT link FROM links WHERE user_id='{self.user_id}' ORDER BY id DESC" ) result = await cur.fetchall() return result
async def _insert_one(self, instance): async with aiopg.connect(**self.broker_queue_db) as connect: async with connect.cursor() as cur: await cur.execute( "INSERT INTO consumer_queue (topic, partition, data) VALUES (%s, %s, %s) RETURNING id;", (instance.topic, 0, instance.avro_bytes), ) return (await cur.fetchone())[0]
async def get_user_by_id(user_id): query = 'select user_id, name, age, phone from users where user_id = %(user_id)s ;' async with aiopg.connect(DB_URL) as conn: async with conn.cursor(cursor_factory=DictCursor) as cur: await cur.execute(query, {'user_id': user_id}) data = await cur.fetchone() if data: return dict(data)
async def main(): async with aiopg.connect(dsn) as listenConn: async with aiopg.create_pool(dsn) as notifyPool: async with notifyPool.acquire() as notifyConn: listener = listen(listenConn) notifier = notify(notifyConn) await asyncio.gather(listener, notifier) print("ALL DONE")
def go(): conn = yield from aiopg.connect(database='aiopg', user='******', password='******', host='127.0.0.1', loop=self.loop) yield from asyncio.gather(waiter(conn), closer(conn), loop=self.loop)
async def test_connect_context_manager(loop, pg_params): async with aiopg.connect(loop=loop, **pg_params) as conn: cursor = await conn.cursor() await cursor.execute('SELECT 42') resp = await cursor.fetchone() assert resp == (42, ) cursor.close() assert conn.closed
def go(): with self.assertRaises(psycopg2.OperationalError): yield from aiopg.connect(database='aiopg', user='******', password='******', host='127.0.0.1', port=port, loop=self.loop)
async def test_connect_context_manager(pg_params): async with aiopg.connect(**pg_params) as conn: cursor = await conn.cursor() await cursor.execute('SELECT 42') resp = await cursor.fetchone() assert resp == (42, ) cursor.close() assert conn.closed
def connect(self, no_loop=False, **kwargs): loop = None if no_loop else self.loop conn = yield from aiopg.connect(database='aiopg', user='******', password='******', host='127.0.0.1', loop=loop, **kwargs) conn2 = yield from aiopg.connect(database='aiopg', user='******', password='******', host='127.0.0.1', loop=loop) cur = yield from conn2.cursor() yield from cur.execute("DROP TABLE IF EXISTS foo") yield from conn2.close() return conn
def aiopg_connection(pg_kwargs, loop): import aiopg connection = block(aiopg.connect(loop=loop, **pg_kwargs), loop=loop, timeout=2) yield connection connection.close()
def connect(self, **kwargs): conn = yield from connect(database='aiopg', user='******', password='******', host='127.0.0.1', loop=self.loop, **kwargs) ret = sa.SAConnection(conn, sa.dialect) return ret
async def delete_user(self): async with aiopg.connect(PGCONF) as conn: cur = await conn.cursor() async with cur.begin(): await cur.execute( f"DELETE from links where user_id='{self.id}'") await cur.execute(f"DELETE from users where id='{self.id}'") return True return False
async def check_user(login): async with aiopg.connect(PGCONF) as conn: cur = await conn.cursor() await cur.execute(f"SELECT id FROM users where login='******'") ret = await cur.fetchall() if len(ret) > 0: return True else: return False
async def get_login(self): async with aiopg.connect(PGCONF) as conn: cur = await conn.cursor() await cur.execute(f"SELECT login FROM users where id='{self.id}'") result = await cur.fetchone() if len(result) > 0: return result[0][0] else: return False
async def main(config): if config.get("sentry"): sentry_sdk.init(config["sentry"]) async with connect(config["postgres_uri"]) as conn: async with conn.cursor() as cur: rabbit = AtSomeoneRabbit(config, cur) await rabbit.connect() log.info("Connected") await rabbit.consume()
async def go(): kw = dict(database='aiopg', user='******', password='******', host='127.0.0.1', loop=self.loop) async with aiopg.connect(**kw) as conn: cursor = await conn.cursor() await cursor.execute('SELECT 42') resp = await cursor.fetchone() assert resp == (42, ) cursor.close() assert conn.closed
def go(): conn = yield from aiopg.connect( database="aiopg", user="******", password="******", host="127.0.0.1", loop=self.loop ) s = socket.fromfd(conn._fileno, socket.AF_INET, socket.SOCK_STREAM) s.send(b"garbage") s.detach() cur = yield from conn.cursor() with self.assertRaises(psycopg2.DatabaseError): yield from cur.execute("SELECT 1")
def go(): conn = yield from aiopg.connect(database='aiopg', user='******', password='******', host='127.0.0.1', loop=self.loop) s = socket.fromfd(conn._fileno, socket.AF_INET, socket.SOCK_STREAM) s.send(b'garbage') s.detach() cur = yield from conn.cursor() with self.assertRaises(psycopg2.DatabaseError): yield from cur.execute('SELECT 1')
def test_patch_unpatch(self): tracer = get_dummy_tracer() writer = tracer.writer # Test patch idempotence patch() patch() service = 'fo' conn = yield from aiopg.connect(**POSTGRES_CONFIG) Pin.get_from(conn).clone(service=service, tracer=tracer).onto(conn) yield from (yield from conn.cursor()).execute('select \'blah\'') conn.close() spans = writer.pop() assert spans, spans eq_(len(spans), 1) # Test unpatch unpatch() conn = yield from aiopg.connect(**POSTGRES_CONFIG) yield from (yield from conn.cursor()).execute('select \'blah\'') conn.close() spans = writer.pop() assert not spans, spans # Test patch again patch() conn = yield from aiopg.connect(**POSTGRES_CONFIG) Pin.get_from(conn).clone(service=service, tracer=tracer).onto(conn) yield from (yield from conn.cursor()).execute('select \'blah\'') conn.close() spans = writer.pop() assert spans, spans eq_(len(spans), 1)
def test___del__(loop, pg_params, warning): exc_handler = mock.Mock() loop.set_exception_handler(exc_handler) conn = yield from aiopg.connect(loop=loop, **pg_params) with warning(ResourceWarning): del conn gc.collect() msg = {'connection': mock.ANY, # conn was deleted 'message': 'Unclosed connection'} if loop.get_debug(): msg['source_traceback'] = mock.ANY exc_handler.assert_called_with(loop, msg)
def go(): exc_handler = unittest.mock.Mock() self.loop.set_exception_handler(exc_handler) conn = yield from aiopg.connect( database="aiopg", user="******", password="******", host="127.0.0.1", loop=self.loop ) with self.assertWarns(ResourceWarning): del conn gc.collect() msg = {"connection": unittest.mock.ANY, "message": "Unclosed connection"} # conn was deleted if self.loop.get_debug(): msg["source_traceback"] = unittest.mock.ANY exc_handler.assert_called_with(self.loop, msg)
def connect(self, **kwargs): conn = yield from connect(database='aiopg', user='******', password='******', host='127.0.0.1', loop=self.loop, **kwargs) cur = yield from conn.cursor() yield from cur.execute("DROP TABLE IF EXISTS sa_tbl") yield from cur.execute("CREATE TABLE sa_tbl " "(id serial, name varchar(255))") yield from cur.execute("INSERT INTO sa_tbl (name)" "VALUES ('first')") cur.close() return sa.SAConnection(conn, sa.engine._dialect)
def connect(self, **kwargs): conn = yield from connect(database='aiopg', user='******', password='******', host='127.0.0.1', loop=self.loop, **kwargs) cur = yield from conn.cursor() yield from cur.execute("DROP TABLE IF EXISTS tbl") yield from cur.execute("""CREATE TABLE tbl ( id SERIAL, val JSON)""") cur.close() self.addCleanup(conn.close) return conn
def go(): exc_handler = unittest.mock.Mock() self.loop.set_exception_handler(exc_handler) conn = yield from aiopg.connect(database='aiopg', user='******', password='******', host='127.0.0.1', loop=self.loop) with self.assertWarns(ResourceWarning): del conn gc.collect() msg = {'connection': unittest.mock.ANY, # conn was deleted 'message': 'Unclosed connection'} if self.loop.get_debug(): msg['source_traceback'] = unittest.mock.ANY exc_handler.assert_called_with(self.loop, msg)
def connect(self, **kwargs): conn = yield from connect(database='aiopg', user='******', password='******', host='127.0.0.1', loop=self.loop, **kwargs) self.addCleanup(conn.close) engine = mock.Mock() engine.dialect = sa.engine._dialect def release(*args): return engine.release = release ret = sa.SAConnection(conn, engine) return ret
def get_cursor(cls, cursor_type=_CursorType.PLAIN) -> Cursor: """ Yields: new client-side cursor from existing db connection pool """ _cur = None if cls._use_pool: _connection_source = yield from cls.get_pool() else: _connection_source = yield from aiopg.connect(echo=False, **cls._connection_params) if cursor_type == _CursorType.PLAIN: _cur = yield from _connection_source.cursor() if cursor_type == _CursorType.NAMEDTUPLE: _cur = yield from _connection_source.cursor(cursor_factory=psycopg2.extras.NamedTupleCursor) if cursor_type == _CursorType.DICT: _cur = yield from _connection_source.cursor(cursor_factory=psycopg2.extras.DictCursor) if not cls._use_pool: _cur = cursor_context_manager(_connection_source, _cur) return _cur
async def main(loop, i): async with aiohttp.ClientSession(loop=loop) as session: try: print(i) html = await fetch(session, 'http://google.com') print(f'{i}...') clean_text = [BeautifulSoup(html, 'lxml').text] dsn = "dbname='MYDB' user='******' host='MY_IP' password='******'" async with aiopg.connect(dsn) as conn: async with conn.cursor() as cur: sql = "INSERT INTO delete_me (response) VALUES (%s)" await cur.execute(sql, clean_text) conn.commit() except aiohttp.client_exceptions.ClientConnectionError as e: print('>>>>> Connection Error') print(e) except asyncio.TimeoutError as e: print('!!!!! Timeout Error') except psycopg2.Error as e: print('##### DB Connection Error') print(type(e)) print(e)
def _get_conn_and_tracer(self): conn = self._conn = yield from aiopg.connect(**POSTGRES_CONFIG) Pin.get_from(conn).clone(tracer=self.tracer).onto(conn) return conn, self.tracer
def connect(self): self._conn = yield from aiopg.connect(dsn=self.dsn, timeout=self.timeout, loop=self._loop, enable_json=False, **self.connect_kwargs)
def connect(self): """Connect asynchronously. """ self._conn = yield from aiopg.connect( dsn=self.dsn, timeout=self.timeout, loop=self._loop, **self.connect_kwargs)
def connect(self): self.log.info('connecting to db %s', self.dbname) return aiopg.connect(dsn='dbname={dbname} user={user} password={password}'.format(dbname=self.dbname, user=self.user, password=self.password))
def go(): with self.assertRaises(psycopg2.OperationalError): yield from aiopg.connect( database="aiopg", user="******", password="******", host="127.0.0.1", port=port, loop=self.loop )
def go(): conn = yield from aiopg.connect( database="aiopg", user="******", password="******", host="127.0.0.1", loop=self.loop ) yield from asyncio.gather(waiter(conn), closer(conn), loop=self.loop)
def test_connect_to_unsupported_port(unused_port, loop, pg_params): port = unused_port() pg_params['port'] = port with pytest.raises(psycopg2.OperationalError): yield from aiopg.connect(loop=loop, **pg_params)