Python asyncpg.connect() Examples

def test_ssl_connection_custom_context(self):
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
        ssl_context.load_verify_locations(SSL_CA_CERT_FILE)

        con = await self.connect(
            host='localhost',
            user='ssl_user',
            ssl=ssl_context)

        try:
            self.assertEqual(await con.fetchval('SELECT 42'), 42)

            with self.assertRaises(asyncio.TimeoutError):
                await con.execute('SELECT pg_sleep(5)', timeout=0.5)

            self.assertEqual(await con.fetchval('SELECT 43'), 43)
        finally:
            await con.close() 

def test_auth_password_cleartext_callable_coroutine(self):
        async def get_correctpassword():
            return 'correctpassword'

        async def get_wrongpassword():
            return 'wrongpassword'

        conn = await self.connect(
            user='password_user',
            password=get_correctpassword)
        await conn.close()

        with self.assertRaisesRegex(
                asyncpg.InvalidPasswordError,
                'password authentication failed for user "password_user"'):
            await self._try_connect(
                user='password_user',
                password=get_wrongpassword) 

def test_auth_password_cleartext_callable(self):
        def get_correctpassword():
            return 'correctpassword'

        def get_wrongpassword():
            return 'wrongpassword'

        conn = await self.connect(
            user='password_user',
            password=get_correctpassword)
        await conn.close()

        with self.assertRaisesRegex(
                asyncpg.InvalidPasswordError,
                'password authentication failed for user "password_user"'):
            await self._try_connect(
                user='password_user',
                password=get_wrongpassword) 

def test_find_tables(self):
        tables = importer.find_tables(test_dir / '../testlayers/testmaptiles.yaml')
        self.assertEqual(tables, ['osm_housenumber_point'])

    # async def test_pg_func(self):
    #     conn = None
    #     try:
    #         pghost, pgport, dbname, user, password = parse_pg_args(
    #             dict(args=dict(dict=lambda v: None))
    #         )
    #         conn = await asyncpg.connect(
    #             database=dbname, host=pghost, port=pgport, user=user, password=password,
    #         )
    #         PgWarnings(conn)
    #         await conn.set_builtin_type_codec('hstore', codec_name='pg_contrib.hstore')
    #
    #     finally:
    #         if conn:
    #             await conn.close() 

def __aenter__(self) -> PgRenderLocker:
        # pg_connection: asyncpg, not Django database, because we use its
        # transaction asynchronously. (Async is so much easier than threading.)
        pg_config = settings.DATABASES["default"]
        pg_connection = await asyncpg.connect(
            host=pg_config["HOST"],
            user=pg_config["USER"],
            password=pg_config["PASSWORD"],
            database=pg_config["NAME"],
            port=pg_config["PORT"],
            timeout=pg_config["CONN_MAX_AGE"],
            command_timeout=pg_config["CONN_MAX_AGE"],
        )

        self.pg_connection = pg_connection

        loop = asyncio.get_event_loop()
        interval = pg_config["CONN_MAX_AGE"]
        self.heartbeat_task = loop.create_task(
            self.send_pg_heartbeats_forever(interval)
        )

        return self 

def connect(self, loop=None, **kwargs):
        conn_info = self.get_connection_spec()
        conn_info.update(kwargs)
        return await asyncpg.connect(loop=loop, **conn_info) 

def test_ssl_connection_default_context(self):
        # XXX: uvloop artifact
        old_handler = self.loop.get_exception_handler()
        try:
            self.loop.set_exception_handler(lambda *args: None)
            with self.assertRaisesRegex(ssl.SSLError, 'verify failed'):
                await self.connect(
                    host='localhost',
                    user='ssl_user',
                    ssl=True)
        finally:
            self.loop.set_exception_handler(old_handler) 

def test_connection_implicit_host(self):
        conn_spec = self.get_connection_spec()
        con = await asyncpg.connect(
            port=conn_spec.get('port'),
            database=conn_spec.get('database'),
            user=conn_spec.get('user'))
        await con.close() 

def test_connection_ssl_unix(self):
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
        ssl_context.load_verify_locations(SSL_CA_CERT_FILE)

        with self.assertRaisesRegex(asyncpg.InterfaceError,
                                    'can only be enabled for TCP addresses'):
            await self.connect(
                host='/tmp',
                ssl=ssl_context) 

def test_connection_sslmode_no_ssl_server(self):
        async def verify_works(sslmode):
            con = None
            try:
                con = await self.connect(
                    dsn='postgresql://foo/?sslmode=' + sslmode,
                    host='localhost')
                self.assertEqual(await con.fetchval('SELECT 42'), 42)
            finally:
                if con:
                    await con.close()

        async def verify_fails(sslmode):
            con = None
            try:
                with self.assertRaises(ConnectionError):
                    await self.connect(
                        dsn='postgresql://foo/?sslmode=' + sslmode,
                        host='localhost')
                    await con.fetchval('SELECT 42')
            finally:
                if con:
                    await con.close()

        await verify_works('disable')
        await verify_works('allow')
        await verify_works('prefer')
        await verify_fails('require')
        await verify_fails('verify-ca')
        await verify_fails('verify-full') 

def test_connection_ssl_to_no_ssl_server(self):
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
        ssl_context.load_verify_locations(SSL_CA_CERT_FILE)

        with self.assertRaisesRegex(ConnectionError, 'rejected SSL'):
            await self.connect(
                host='localhost',
                user='ssl_user',
                ssl=ssl_context) 

def _run_no_explicit_close_test(self):
        con = await self.connect()
        proto = con._protocol
        conref = weakref.ref(con)
        del con

        gc.collect()
        gc.collect()
        gc.collect()

        self.assertIsNone(conref())
        self.assertTrue(proto.is_closed()) 

def test_auth_password_md5(self):
        conn = await self.connect(
            user='md5_user', password='correctpassword')
        await conn.close()

        with self.assertRaisesRegex(
                asyncpg.InvalidPasswordError,
                'password authentication failed for user "md5_user"'):
            await self._try_connect(
                user='md5_user', password='wrongpassword') 

def test_auth_password_cleartext(self):
        conn = await self.connect(
            user='password_user',
            password='correctpassword')
        await conn.close()

        with self.assertRaisesRegex(
                asyncpg.InvalidPasswordError,
                'password authentication failed for user "password_user"'):
            await self._try_connect(
                user='password_user',
                password='wrongpassword') 

def _try_connect(self, **kwargs):
        # On Windows the server sometimes just closes
        # the connection sooner than we receive the
        # actual error.
        if _system == 'Windows':
            for tried in range(3):
                try:
                    return await self.connect(**kwargs)
                except asyncpg.ConnectionDoesNotExistError:
                    pass

        return await self.connect(**kwargs) 

def DBClient():
    return await asyncpg.connect(
        user=DBUSER,
        password=DBPASSWORD,
        database=DBNAME,
        host=DBHOST,
        port=DBPORT,
    ) 

def _test_connection(self, timeout=60):
        self._connection_addr = None

        loop = asyncio.new_event_loop()

        try:
            for i in range(timeout):
                if self._connection_addr is None:
                    conn_spec = self._get_connection_spec()
                    if conn_spec is None:
                        time.sleep(1)
                        continue

                try:
                    con = loop.run_until_complete(
                        asyncpg.connect(database='postgres',
                                        user='postgres',
                                        timeout=5, loop=loop,
                                        **self._connection_addr))
                except (OSError, asyncio.TimeoutError,
                        asyncpg.CannotConnectNowError,
                        asyncpg.PostgresConnectionError):
                    time.sleep(1)
                    continue
                except asyncpg.PostgresError:
                    # Any other error other than ServerNotReadyError or
                    # ConnectionError is interpreted to indicate the server is
                    # up.
                    break
                else:
                    loop.run_until_complete(con.close())
                    break
        finally:
            loop.close()

        return 'running' 

def asyncpg_connect(args):
    conn = await asyncpg.connect(user=args.pguser, host=args.pghost,
                                 port=args.pgport)
    return conn 

def aiopg_connect(args):
    conn = await aiopg.connect(user=args.pguser, host=args.pghost,
                               port=args.pgport)
    return conn 

def psycopg_connect(args):
    conn = psycopg2.connect(user=args.pguser, host=args.pghost,
                            port=args.pgport)
    return conn 

def setUpClass(cls):
        super().setUpClass()
        cls._connection = None
        cls._cursor = None
        cls._tracer = cls.tracer_provider.get_tracer(__name__)
        AsyncPGInstrumentor().instrument(tracer_provider=cls.tracer_provider)
        cls._connection = _await(
            asyncpg.connect(
                database=POSTGRES_DB_NAME,
                user=POSTGRES_USER,
                password=POSTGRES_PASSWORD,
                host=POSTGRES_HOST,
                port=POSTGRES_PORT,
            )
        ) 

def asyncpg_init(db_uri):
    import asyncpg

    async def create_conn():
        global asyncpg_conn
        asyncpg_conn = await asyncpg.connect(db_uri)

    async_run(create_conn)


# asyncpg_init(config.DATABASE_URI) 

def _execute_pg_query(url, query):
    conn = await asyncpg.connect(url)
    if callable(query):
        await query(conn)
    else:
        await conn.execute(query)
    await conn.close() 

def collect_postgres_stats(self, ident: str, vacuum_full: bool = True):
        creds = self.postgres_creds

        conn = await asyncpg.connect(
            host=self.internal_host,
            port="5432",
            user=creds["admin_account"],
            password=creds["admin_password"],
            database=self.wallet_name,
        )

        tables = ("items", "tags_encrypted", "tags_plaintext")
        for t in tables:
            await conn.execute(f"VACUUM FULL {t}" if vacuum_full else f"VACUUM {t}")

        sizes = await conn.fetch(
            """
            SELECT relname AS "relation",
                pg_size_pretty(pg_total_relation_size(C.oid)) AS "total_size"
            FROM pg_class C
            LEFT JOIN pg_namespace N ON (N.oid = C.relnamespace)
            WHERE nspname = 'public'
            ORDER BY pg_total_relation_size(C.oid) DESC;
            """
        )
        results = {k: [0, "0B"] for k in tables}
        for row in sizes:
            if row["relation"] in results:
                results[row["relation"]][1] = row["total_size"].replace(" ", "")
        for t in tables:
            row = await conn.fetchrow(f"""SELECT COUNT(*) AS "count" FROM {t}""")
            results[t][0] = row["count"]
        self.wallet_stats.append((ident, results))

        await conn.close() 

def connect(
        self,
        dbname: str,
        dbver: bytes
    ) -> CompilerDatabaseState:
        self._dbname = dbname
        self._cached_db = None
        await self._get_database(dbver) 

def new_connection(self):
        con_args = self._connect_args.copy()
        con_args['database'] = self._dbname
        try:
            return await asyncpg.connect(**con_args)
        except asyncpg.InvalidCatalogNameError as ex:
            raise errors.AuthenticationError(str(ex)) from ex
        except Exception as ex:
            raise errors.InternalServerError(str(ex)) from ex 

def get_remote_pg_cluster(dsn: str) -> RemoteCluster:
    addrs, params = pgconnparams.parse_dsn(dsn)
    if len(addrs) > 1:
        raise ValueError('multiple hosts in Postgres DSN are not supported')
    rcluster = RemoteCluster(addrs[0], params)

    loop = asyncio.new_event_loop()

    async def _is_rds():
        conn = await rcluster.connect()

        try:
            rds_super = await conn.fetch(
                "SELECT * FROM pg_roles WHERE rolname = 'rds_superuser'"
            )
        finally:
            await conn.close()

        return bool(rds_super)

    try:
        is_rds = loop.run_until_complete(_is_rds())
    finally:
        loop.close()

    if is_rds:
        return RDSCluster(addrs[0], params)
    else:
        return rcluster 

def _test_connection(self, timeout=60):
        self._connection_addr = None

        loop = asyncio.new_event_loop()

        try:
            for _ in range(timeout):
                if self._connection_addr is None:
                    conn_addr = self._get_connection_addr()
                    if conn_addr is None:
                        time.sleep(1)
                        continue

                try:
                    con = loop.run_until_complete(
                        asyncpg.connect(database='postgres',
                                        user='postgres',
                                        timeout=5,
                                        loop=loop,
                                        host=self._connection_addr[0],
                                        port=self._connection_addr[1]))
                except (OSError, asyncio.TimeoutError,
                        asyncpg.CannotConnectNowError,
                        asyncpg.PostgresConnectionError):
                    time.sleep(1)
                    continue
                except asyncpg.PostgresError:
                    # Any other error other than ServerNotReadyError or
                    # ConnectionError is interpreted to indicate the server is
                    # up.
                    break
                else:
                    loop.run_until_complete(con.close())
                    break
        finally:
            loop.close()

        return 'running' 

def connect(self, loop=None, **kwargs):
        conn_info = self.get_connection_spec()
        conn_info.update(kwargs)
        conn = await asyncpg.connect(loop=loop, **conn_info)

        if (not kwargs.get('user')
                and self._default_session_auth
                and conn_info.get('user') != self._default_session_auth):
            # No explicit user given, and the default
            # SESSION AUTHORIZATION is different from the user
            # used to connect.
            await conn.execute(
                f'SET SESSION AUTHORIZATION {self._default_session_auth};'
            )

        return conn 

def DBClient():
    return await asyncpg.connect(
        user=DBUSER,
        password=DBPASSWORD,
        database=DBNAME,
        host=DBHOST,
        port=DBPORT,
    )