from sqlalchemy.dialects import registry
from sqlalchemy.dialects.postgresql.psycopg import PGDialect_psycopg
from sqlalchemy import NullPool
from sqlalchemy import create_engine


class MyDialect(PGDialect_psycopg):
    @classmethod
    def get_pool_class(cls, url):
        return NullPool

    def connect(self, *args, **params):
        # this is your pool
        return pool.getconn()

    def do_close(self, dbapi_connection):
        # this is your pool
        pool.putconn(dbapi_connection)

# second argument is the module path, e.g. if we do `__import__(modulepath)`
registry.register("postgresql.mydialect", "__main__", "MyDialect")


e = create_engine("postgresql+mydialect://")

# this will be MyDialect
print(e.dialect)

# this will be NullPool
print(e.pool)

class MyConnection(psycopg.Connection):
    def close(self):
        if pool := getattr(self, "_pool", None):
            # Connection was in use; return it to the pool
            pool.putconn(self)
        else:
            # Connection was in pool; close it for real
            super().close()

import logging
import time

logging.basicConfig(
    level=logging.DEBUG, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
)
logging.getLogger("psycopg.pool").setLevel(logging.DEBUG)
logging.getLogger("sqlalchemy.pool").setLevel(logging.DEBUG)

import psycopg_pool
import sqlalchemy

# For database dialect in URL, use Psycopg 3 instead of SQLAlchemy
db_url = "psycopg://postgres:postgres@hostname:port/database"  # CHANGEME


# Create a connection pool that uses the custom connection class
mypool = psycopg_pool.ConnectionPool(
    conninfo=db_url,
    close_returns=True,  # Return "closed" active connections to the pool
    min_size=0,          # Minimum pool size
    max_size=1,          # Maximum pool size
    max_idle=5,          # Maximum idle time (seconds) for unused pool connections
)

# Create an engine that uses the connection pool to get a connection
engine = sqlalchemy.create_engine(
    url="postgresql+psycopg://",         # Only need the dialect now
    poolclass=sqlalchemy.pool.NullPool,  # Disable SQLAlchemy's default connection pool
    creator=mypool.getconn,              # Use Psycopg 3 connection pool to obtain connections
)


# Define a function that asks the server for the current time
def do_stuff():
    with engine.connect() as connection:
        result = connection.execute(sqlalchemy.text("SELECT now()"))
        print(f"Current time: {result.fetchone()[0]}")


# Perform two transactions, using the same pool-sourced server connection
print("Doing do_stuff() #1...")
do_stuff()
print("Doing do_stuff() #2...")
do_stuff()
print("Done doing stuff! Waiting 10 seconds...")

# Allow idle pool connection to disconnect from server
time.sleep(10)

# Perform another transaction, requiring a new pool connection to the server
print("Doing do_stuff() #3...")
do_stuff()

import asyncio
import logging
import time

logging.basicConfig(
    level=logging.DEBUG, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
)
logging.getLogger("psycopg.pool").setLevel(logging.DEBUG)
logging.getLogger("sqlalchemy.pool").setLevel(logging.DEBUG)

import psycopg_pool
import sqlalchemy
import sqlalchemy.ext.asyncio

# For database dialect in URL, use Psycopg 3 instead of SQLAlchemy
db_url = "psycopg://postgres:postgres@hostname:port/database"  # CHANGEME


async def main():
    # Create a connection pool that uses the custom connection class
    # (must be created inside an async event loop)
    mypool = psycopg_pool.AsyncConnectionPool(
        conninfo=db_url,
        close_returns=True,  # Return "closed" active connections to the pool
        min_size=0,          # Minimum pool size
        max_size=1,          # Maximum pool size
        max_idle=5,          # Maximum idle time (seconds) for unused pool connections
        open=False,          # See comment below
    )

    # Must explicitly open AsyncConnectionPool outside constructor
    # (https://www.psycopg.org/psycopg3/docs/advanced/pool.html#other-ways-to-create-a-pool)
    await mypool.open()

    # Create an engine that uses the connection pool to get a connection
    engine = sqlalchemy.ext.asyncio.create_async_engine(
        url="postgresql+psycopg://",         # Only need the dialect now
        poolclass=sqlalchemy.pool.NullPool,  # Disable SQLAlchemy's default connection pool
        async_creator=mypool.getconn,        # Use Psycopg 3 connection pool to obtain connections
    )

    # Define a function that asks the server for the current time
    async def do_stuff():
        async with engine.connect() as connection:
            result = await connection.execute(sqlalchemy.text("SELECT now()"))
            print(f"Current time: {result.fetchone()[0]}")

    # Perform two transactions, using the same pool-sourced server connection
    print("Doing do_stuff() #1...")
    await do_stuff()
    print("Doing do_stuff() #2...")
    await do_stuff()
    print("Done doing stuff! Waiting 10 seconds...")

    # Allow idle pool connection to disconnect from server
    await asyncio.sleep(10)

    # Perform another transaction, requiring a new pool connection to the server
    print("Doing do_stuff() #3...")
    await do_stuff()


if __name__ == "__main__":
    # Run the main function in an asyncio event loop
    asyncio.run(main())

import logging
import time

logging.basicConfig(
    level=logging.DEBUG, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
)
logging.getLogger("psycopg.pool").setLevel(logging.DEBUG)
logging.getLogger("sqlalchemy.pool").setLevel(logging.DEBUG)

import psycopg
import psycopg_pool
import sqlalchemy

# For database dialect in URL, use Psycopg 3 instead of SQLAlchemy
db_url = "psycopg://postgres:postgres@hostname:port/database"  # CHANGEME


# Define a custom connection class that overrides the close() method
class MyConnection(psycopg.Connection):
    def close(self):
        if pool := getattr(self, "_pool", None):
            # Connection currently checked out from its pool;
            # instead of closing it, return it to the pool.
            pool.putconn(self)
        else:
            # Connection being removed from its pool, or not part of any pool;
            # close the connection for real.
            super().close()


# Create a connection pool that uses the custom connection class
mypool = psycopg_pool.ConnectionPool(
    conninfo=db_url,
    connection_class=MyConnection,  # Use custom connection class to return closed connections to the pool
    min_size=0,  # Minimum pool size
    max_size=1,  # Maximum pool size
    max_idle=5,  # Maximum idle time (seconds) for unused pool connections
)

# Create an engine that uses the connection pool to get a connection
engine = sqlalchemy.create_engine(
    url="postgresql+psycopg://",         # Only need the dialect now
    poolclass=sqlalchemy.pool.NullPool,  # Disable SQLAlchemy's default connection pool
    creator=mypool.getconn,              # Use Psycopg 3 connection pool to obtain connections
)


# Define a function that asks the server for the current time
def do_stuff():
    with engine.connect() as connection:
        result = connection.execute(sqlalchemy.text("SELECT now()"))
        print(f"Current time: {result.fetchone()[0]}")


# Perform two transactions, using the same pool-sourced server connection
print("Doing do_stuff() #1...")
do_stuff()
print("Doing do_stuff() #2...")
do_stuff()
print("Done doing stuff! Waiting 10 seconds...")

# Allow idle pool connection to disconnect from server
time.sleep(10)

# Perform another transaction, requiring a new pool connection to the server
print("Doing do_stuff() #3...")
do_stuff()

import asyncio
import logging
import time

logging.basicConfig(
    level=logging.DEBUG, format="%(asctime)s %(levelname)s %(name)s: %(message)s"
)
logging.getLogger("psycopg.pool").setLevel(logging.DEBUG)
logging.getLogger("sqlalchemy.pool").setLevel(logging.DEBUG)

import psycopg
import psycopg_pool
import sqlalchemy
import sqlalchemy.ext.asyncio

# For database dialect in URL, use Psycopg 3 instead of SQLAlchemy
db_url = "psycopg://postgres:postgres@hostname:port/database"  # CHANGEME


# Define a custom connection class that overrides the close() method
class MyAsyncConnection(psycopg.AsyncConnection):
    async def close(self):
        if pool := getattr(self, "_pool", None):
            # Connection currently checked out from its pool;
            # instead of closing it, return it to the pool.
            await pool.putconn(self)
        else:
            # Connection being removed from its pool, or not part of any pool;
            # close the connection for real.
            await super().close()


async def main():
    # Create a connection pool that uses the custom connection class
    # (must be created inside an async event loop)
    mypool = psycopg_pool.AsyncConnectionPool(
        conninfo=db_url,
        connection_class=MyAsyncConnection,  # Use custom connection class to return closed connections to the pool
        min_size=0,  # Minimum pool size
        max_size=1,  # Maximum pool size
        max_idle=5,  # Maximum idle time (seconds) for unused pool connections
        open=False,  # See comment below
    )

    # Must explicitly open AsyncConnectionPool outside constructor
    # (https://www.psycopg.org/psycopg3/docs/advanced/pool.html#other-ways-to-create-a-pool)
    await mypool.open()

    # Create an engine that uses the connection pool to get a connection
    engine = sqlalchemy.ext.asyncio.create_async_engine(
        url="postgresql+psycopg://",         # Only need the dialect now
        poolclass=sqlalchemy.pool.NullPool,  # Disable SQLAlchemy's default connection pool
        async_creator=mypool.getconn,        # Use Psycopg 3 connection pool to obtain connections
    )

    # Define a function that asks the server for the current time
    async def do_stuff():
        async with engine.connect() as connection:
            result = await connection.execute(sqlalchemy.text("SELECT now()"))
            print(f"Current time: {result.fetchone()[0]}")

    # Perform two transactions, using the same pool-sourced server connection
    print("Doing do_stuff() #1...")
    await do_stuff()
    print("Doing do_stuff() #2...")
    await do_stuff()
    print("Done doing stuff! Waiting 10 seconds...")

    # Allow idle pool connection to disconnect from server
    await asyncio.sleep(10)

    # Perform another transaction, requiring a new pool connection to the server
    print("Doing do_stuff() #3...")
    await do_stuff()


if __name__ == "__main__":
    # Run the main function in an asyncio event loop
    asyncio.run(main())