Python asyncpg.create_pool() Examples

def run_bot():
    """Launches the bot."""
    log.info("Launching bot...")
    loop = asyncio.get_event_loop()

    pool: asyncpg.pool.Pool = loop.run_until_complete(create_pool(get_uri(), command_timeout=60))
    if pool is None:
        log.error('Could not set up PostgreSQL. Exiting.')
        return

    result = loop.run_until_complete(check_database(pool))
    if not result:
        log.error('Failed to check database')
        return

    bot = NabBot()
    bot.pool = pool
    bot.run() 

def create_pool(*args,
                dialect=None,
                connection_class=_SAConnection,
                **connect_kwargs):

    class SAConnection(connection_class):
        def __init__(self, *args, dialect=dialect, **kwargs):
            super().__init__(*args, dialect=dialect, **kwargs)

    connection_class = SAConnection

    # dict is fine on the pool object as there is usually only one of them
    # asyncpg.pool.Pool.__slots__ += ('__dict__',)

    # monkey patch pool to have some extra methods
    def transaction(self, **kwargs):
        return ConnectionTransactionContextManager(self, **kwargs)
    asyncpg.pool.Pool.transaction = transaction
    asyncpg.pool.Pool.begin = transaction
    pool = asyncpg.create_pool(*args, connection_class=connection_class,
                               **connect_kwargs)
    return pool 

def on_ready(self):
        guild = self.bot.guilds[0]

        if self.staff_role is None:
            self.staff_role = guild.get_role(STAFF_ROLE_ID)

        if self.fake_staff_role is None:
            self.fake_staff_role = guild.get_role(FAKE_ROLE_ID)

        self.bot.pool = await asyncpg.create_pool(
            host=PostgreSQL.PGHOST,
            port=PostgreSQL.PGPORT,
            user=PostgreSQL.PGUSER,
            password=PostgreSQL.PGPASSWORD,
            database=PostgreSQL.PGDATABASE,
        )

        await self.migrate_quotes() 

def empty():
    """Empties out the database.

    Drops all tables and functions from the saved PostgreSQL database.
    This action is irreversible, so use with caution."""
    loop = asyncio.get_event_loop()
    pool: asyncpg.pool.Pool = loop.run_until_complete(create_pool(get_uri(), command_timeout=60))
    if pool is None:
        log.error('Could not set up PostgreSQL. Exiting.')
        return

    db_name = loop.run_until_complete(get_db_name(pool))

    confirm = click.confirm(f"You are about to drop all the tables and functions of the database '{db_name}'.\n"
                            "Are you sure you want to continue? This action is irreversible.")
    if not confirm:
        log.warning("Operation aborted.")
        return

    log.info("Clearing database...")
    loop.run_until_complete(drop_tables(pool))
    log.info("Database cleared") 

def create_pool(uri, **kwargs) -> asyncpg.pool.Pool:
    """Creates a connection pool to the specified PostgreSQL server"""
    def _encode_jsonb(value):
        return b'\x01' + json.dumps(value).encode('utf-8')

    def _decode_jsonb(value):
        return json.loads(value[1:].decode('utf-8'))

    async def init(con):
        await con.set_type_codec('jsonb', schema='pg_catalog', encoder=_encode_jsonb, decoder=_decode_jsonb,
                                 format="binary")
    try:
        log.debug("Creating connection pool")
        pool = await asyncpg.create_pool(uri, init=init, **kwargs)
    except ValueError:
        log.error("PostgreSQL error: Invalid URI, check postgresql.txt. "
                  "Format must be 'postresql://user:password@host/database'")
    except asyncpg.PostgresError as e:
        log.error(f"PostgreSQL error: {e}")
    except TimeoutError:
        log.error("PostgreSQL error: Connection timed out.")
    except Exception as e:
        log.error(f"Unexpected error: {e.__class__.__name__}: {e}")
    else:
        return pool 

def async_main(concurrency):
    print(f'concurrency={concurrency}')
    async with asyncpg.create_pool(
            user='rssant', password='rssant',
            database='rssant', host='127.0.0.1',
            command_timeout=60, min_size=5, max_size=5
    ) as pool:
        async with pool.acquire() as conn:
            await init_table(conn)
        await benchmark(
            pool,
            num_fid=1000,
            num_round=3,
            concurrency=concurrency,
        )
    await pool.close() 

def create_db():
    dsn = current_app.config["FLOWDB_DSN"]
    current_app.db_conn_pool = await asyncpg.create_pool(dsn, max_size=20) 

def connect(self):
        self._conn = await asyncpg.create_pool(user='root', password='root',
                                               database='pokerpg', host='127.0.0.1')

    # User functions
    ######################################################################## 

def db_init(db_url):
    """Initializes the database connection"""
    global Pool
    Pool = await asyncpg.create_pool(dsn=db_url, command_timeout=15) 

def startup(app: web.Application):
    settings: Settings = app['settings']
    await prepare_database(settings, False)
    app['pg'] = await asyncpg.create_pool(dsn=settings.pg_dsn, min_size=2) 

def run(self):
        print(f'Connecting to PostgreSQL at {self.pghost}:{self.pgport}, '
              f'db={self.dbname}, user={self.user}...')
        async with asyncpg.create_pool(
            database=self.dbname, host=self.pghost, port=self.pgport, user=self.user,
            password=self.password, min_size=1, max_size=1,
        ) as pool:
            async with pool.acquire() as conn:
                self.results.created = dt.utcnow().isoformat()
                self.results.tileset = str(Path(self.tileset.filename).resolve())
                await self._run(conn)
                self.results.tests = [v.result for v in self.test_cases]
                self.save_results() 

def setup_db_pg(app, loop):
    """
    initiate postgresql connection
    """
    if Config.dsn:
        try:
            Config.db_pg = await asyncpg.create_pool(Config.dsn, loop=loop)
        except socket.gaierror:
            print(f'Cannot establish connection to {Config.dsn}. \
Did you pass correct values to --pghost?')
            raise
        except asyncpg.exceptions.InvalidPasswordError:
            print(f'Cannot connect to {Config.dsn}. \
Please check values passed to --pguser and --pgpassword')
            raise 

def get_instance(cls, logger=None, connect_kwargs: dict = None,
                           pool: Pool = None, schema: str = 'nanochan'):
        """
        Get a new instance of `PostgresController`
        This method will create the appropriate tables needed.
        :param logger: the logger object.
        :param connect_kwargs:
            Keyword arguments for the
            :func:`asyncpg.connection.connect` function.
        :param pool: an existing connection pool.
        One of `pool` or `connect_kwargs` must not be None.
        :param schema: the schema name used. Defaults to `minoshiro`
        :return: a new instance of `PostgresController`
        """
        assert logger, (
            'Please provide a logger to the data_controller'
        )
        assert connect_kwargs or pool, (
            'Please either provide a connection pool or '
            'a dict of connection data for creating a new '
            'connection pool.'
        )
        if not pool:
            try:
                pool = await create_pool(**connect_kwargs)
                logger.info('Connection pool made.')
            except InterfaceError as e:
                logger.error(str(e))
                raise e
        logger.info('Creating tables...')
        await make_tables(pool, schema)
        logger.info('Tables created.')
        return cls(pool, logger, schema) 

def async_pool(user=None, password=None, host=None, port=None,
                     db=None, max_connections=150, timeout=60):
    user = user or os.environ.get('PGUSER', 'postgres')
    password = password or os.environ.get('PGPASSWORD', '')
    host = host or os.environ.get('PGHOST', 'localhost')
    port = port or os.environ.get('PGPORT', '5432')
    db = db or os.environ.get('PGDATABASE', 'postgres')
    db_pool = await asyncpg.create_pool(user=user, password=password,
                                        database=db, host=host,
                                        port=port,
                                        command_timeout=timeout,
                                        max_size=max_connections)
    return db_pool 

def create_app(loop, db):
    app = web.Application(loop=loop, middlewares=[
        session_middleware(SimpleCookieStorage()),
    ])
    app.middlewares.append(aiohttp_login.flash.middleware)
    aiohttp_jinja2.setup(
        app,
        loader=jinja_app_loader.Loader(),
        context_processors=[aiohttp_login.flash.context_processor],
    )

    if db == 'asyncpg':
        pool = await asyncpg.create_pool(
            dsn='postgres:///' + DATABASE, loop=loop)
        storage = AsyncpgStorage(pool)
    elif db == 'motor':
        app['db'] = AsyncIOMotorClient(io_loop=loop)[DATABASE]
        storage = MotorStorage(app['db'])
    else:
        assert 0, 'unknown storage'

    aiohttp_login.setup(app, storage, {
        'CSRF_SECRET': 'secret',
        'LOGIN_REDIRECT': 'auth_change_email',
        'SMTP_SENDER': 'Your Name <your@gmail.com>',
        'SMTP_HOST': 'smtp.gmail.com',
        'SMTP_PORT': 465,
        'SMTP_USERNAME': 'your@gmail.com',
        'SMTP_PASSWORD': 'password'
    })

    @restricted_api
    async def api_hello_handler(request):
        return {'hello': 'world'}

    app.router.add_get('/api/hello', api_hello_handler, name='api_hello')

    return app 

def create_app(loop):
    app = web.Application(loop=loop, debug=settings.DEBUG)
    setup_jinja(app, settings.DEBUG)
    aiohttp_session.setup(app, EncryptedCookieStorage(
        settings.SESSION_SECRET.encode('utf-8'),
        max_age=settings.SESSION_MAX_AGE))
    app.middlewares.append(aiohttp_login.flash.middleware)

    app.router.add_get('/', handlers.index)
    app.router.add_get('/users/', handlers.users, name='users')

    app['db'] = await asyncpg.create_pool(dsn=settings.DATABASE, loop=loop)
    aiohttp_login.setup(app, AsyncpgStorage(app['db']), settings.AUTH)

    return app 

def get_poll(self) -> asyncpg.pool.Pool:
        if not self.poll:
            self.poll = await asyncpg.create_pool('postgresql://duckhunt:duckhunt@localhost/dhv3')

        return self.poll 

def migrate(path):
    """Migrates a v1.x.x SQLite to a PostgreSQL database.

    This is a time consuming operation and caution must be taken.
    The original SQLite file is not affected."""
    loop = asyncio.get_event_loop()
    pool: asyncpg.pool.Pool = loop.run_until_complete(create_pool(get_uri(), command_timeout=240))
    if pool is None:
        log.error('Could not set up PostgreSQL. Exiting.')
        return

    db_name = loop.run_until_complete(get_db_name(pool))

    confirm = click.confirm("Migrating a SQL database requires an empty PostgreSQL database.\n"
                            f"Confirming will delete all data from the database '{db_name}'.\n"
                            f"The SQL database located in {path} will be imported afterwards.\n"
                            "Are you sure you want to continue? This action is irreversible.")
    if not confirm:
        log.warning("Operation aborted.")
        return

    log.info("Clearing database...")
    loop.run_until_complete(drop_tables(pool))
    log.info("Database cleared")

    log.info("Starting migration...")
    result = loop.run_until_complete(check_database(pool))
    if not result:
        log.error('Failed to check database')
        return

    loop.run_until_complete(import_legacy_db(pool, path))
    log.info("Migration complete") 

def connect_to_db(app: FastAPI) -> None:
    logger.info("Connecting to {0}", repr(DATABASE_URL))

    app.state.pool = await asyncpg.create_pool(
        str(DATABASE_URL),
        min_size=MIN_CONNECTIONS_COUNT,
        max_size=MAX_CONNECTIONS_COUNT,
    )

    logger.info("Connection established") 

def connect_redis(self):
        self.redis = await aioredis.create_pool("redis://localhost", minsize=5, maxsize=10, loop=self.loop, db=0)
        info = (await self.redis.execute("INFO")).decode()
        for line in info.split("\n"):
            if line.startswith("redis_version"):
                self.redis_version = line.split(":")[1]
                break 

def connect_postgres(self):
        self.pool = await asyncpg.create_pool(**self.config.database, max_size=20, command_timeout=60) 

def create_db_pool(current_app=current_app):
    return await asyncpg.create_pool(
        host=current_app.config["DB_HOST"],
        port=current_app.config["DB_PORT"],
        user=current_app.config["DB_USER"],
        password=current_app.config["DB_PASSWORD"],
        database=current_app.config["DB_NAME"],
        # we want to rely on pgbouncer
        statement_cache_size=0,
    ) 

def start(self) -> None:
        self.db_args["loop"] = self.loop
        self.log.debug(f"Connecting to {self.url}")
        self.pool = await asyncpg.create_pool(self.url, **self.db_args)
        try:
            await self.upgrade_table.upgrade(self.pool)
        except Exception:
            self.log.critical("Failed to upgrade database", exc_info=True)
            sys.exit(25) 

def message_loop(self):
        try:
            if self.dsn:
                self.db = await asyncpg.create_pool(dsn=self.dsn, min_size=1, max_size=1)
            self.reader = await self.get_pipe_reader(self.in_reader)
            self.writer = await self.get_pipe_writer(self.out_writer)
            while True:
                msg_type, msg = await self.pipe_get_msg(self.reader)
                if msg_type ==  b'pipe_read_error':
                    return

                if msg_type == b'get':
                    self.loop.create_task(self.load_blocks(bytes_to_int(msg), self.rpc_batch_limit))
                    continue

                if msg_type == b'rpc_batch_limit':
                    self.rpc_batch_limit = bytes_to_int(msg)
                    continue

                if msg_type == b'target_height':
                    self.target_height = bytes_to_int(msg)
                    continue


        except:
            pass 

def run_asyncpg():
    async with asyncpg.create_pool(
            user='rssant', password='rssant',
            database='rssant', host='127.0.0.1',
            command_timeout=60, min_size=5, max_size=5
    ) as pool:
        t0 = time.time()
        for i in range(1000):
            async with pool.acquire() as conn:
                values = await conn.fetch("SELECT 1")
                assert values == [(1,)]
        print('run_asyncpg', time.time() - t0)
    await pool.close() 

def run_aiopg():
    pool = await aiopg.create_pool(dsn, minsize=5, maxsize=5)
    t0 = time.time()
    for i in range(1000):
        async with pool.acquire() as conn:
            async with conn.cursor() as cur:
                await cur.execute("SELECT 1")
                ret = []
                async for row in cur:
                    ret.append(row)
                assert ret == [(1,)]
    print('run_aiopg', time.time() - t0)
    pool.close()
    await pool.wait_closed() 

def _get_pool(self):
        if not self._pool:
            self._pool = await asyncpg.create_pool(
                self.dsn,
                min_size=self.pool_size,
                max_size=self.pool_size
            )
        return self._pool 

def generate(self, tileset, reset, auto_minmax, show_ranges,
                       pghost, pgport, dbname, user, password):
        ts = Tileset.parse(tileset)
        print(
            f'Connecting to PostgreSQL at {pghost}:{pgport}, db={dbname}, user={user}...')
        try:
            async with asyncpg.create_pool(
                database=dbname, host=pghost, port=pgport, user=user,
                password=password, min_size=1, max_size=1,
            ) as pool:
                async with pool.acquire() as conn:
                    mvt = MvtGenerator(
                        ts,
                        postgis_ver=await get_postgis_version(conn),
                        zoom='$1', x='$2', y='$3',
                    )
                    json_data = dict(vector_layers=await get_vector_layers(conn, mvt))
        except ConnectionError as err:
            print(f"Unable to connect to Postgres database: {err}")
            raise err

        # Convert tileset to the metadata object according to mbtiles 1.3 spec
        # https://github.com/mapbox/mbtiles-spec/blob/master/1.3/spec.md#content
        metadata = dict(
            # MUST
            name=ts.name,
            format="pbf",
            json=json.dumps(json_data, ensure_ascii=False, separators=(',', ':')),
            # SHOULD
            bounds=",".join((str(v) for v in ts.bounds)),
            center=",".join((str(v) for v in ts.center)),
            minzoom=str(ts.minzoom),
            maxzoom=str(ts.maxzoom),
            # MAY
            attribution=ts.attribution,
            description=ts.description,
            version=ts.version,
            # EXTRAS
            id=ts.id,
        )

        self._update_metadata(metadata, auto_minmax, reset, self.mbtiles, show_ranges,
                              ts.center[2])