import sqlalchemy
from google.cloud.alloydbconnector import Connector
def create_sqlalchemy_engine(
inst_uri: str,
user: str,
password: str,
db: str,
refresh_strategy: str = "background",
ip_type: str = "PRIVATE",
) -> tuple[sqlalchemy.engine.Engine, Connector]:
"""Creates a connection pool for an AlloyDB instance and returns the pool
and the connector. Callers are responsible for closing the pool and the
connector.
A sample invocation looks like:
engine, connector = create_sqlalchemy_engine(
inst_uri,
user,
password,
db,
)
with engine.connect() as conn:
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
conn.commit()
curr_time = time[0]
# do something with query result
connector.close()
Args:
instance_uri (str):
The instance URI specifies the instance relative to the project,
region, and cluster. For example:
"projects/my-project/locations/us-central1/clusters/my-cluster/instances/my-instance"
user (str):
The database user name, e.g., postgres
password (str):
The database user's password, e.g., secret-password
db (str):
The name of the database, e.g., mydb
refresh_strategy (Optional[str]):
Refresh strategy for the AlloyDB Connector. Can be one of "lazy"
or "background". For serverless environments use "lazy" to avoid
errors resulting from CPU being throttled.
ip_type (str):
IP type used to connect. One of "PUBLIC", "PRIVATE" or "PSC".
"""
connector = Connector(refresh_strategy=refresh_strategy)
# create SQLAlchemy connection pool
engine = sqlalchemy.create_engine(
"postgresql+psycopg://",
creator=lambda: connector.connect(
inst_uri,
"psycopg",
user=user,
password=password,
db=db,
ip_type=ip_type,
),
)
return engine, connector