Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 22 additions & 7 deletions tests/system/test_asyncpg_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ async def create_sqlalchemy_engine(
user: str,
password: str,
db: str,
ip_type: str = "public",
refresh_strategy: str = "background",
resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver,
) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, Connector]:
Expand Down Expand Up @@ -63,6 +64,8 @@ async def create_sqlalchemy_engine(
The database user's password, e.g., secret-password
db (str):
The name of the database, e.g., mydb
ip_type (str):
The IP type of the Cloud SQL instance. Can be one of "public", "private", or "psc".
refresh_strategy (Optional[str]):
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
or "background". For serverless environments use "lazy" to avoid
Expand All @@ -87,7 +90,7 @@ async def create_sqlalchemy_engine(
user=user,
password=password,
db=db,
ip_type="public", # can also be "private" or "psc"
ip_type=ip_type, # can be "public", "private" or "psc"
),
execution_options={"isolation_level": "AUTOCOMMIT"},
)
Expand All @@ -99,6 +102,7 @@ async def create_asyncpg_pool(
user: str,
password: str,
db: str,
ip_type: str = "public",
refresh_strategy: str = "background",
) -> tuple[asyncpg.Pool, Connector]:
"""Creates a native asyncpg connection pool for a Cloud SQL instance and
Expand Down Expand Up @@ -128,6 +132,8 @@ async def create_asyncpg_pool(
The database user's password, e.g., secret-password
db (str):
The name of the database, e.g., mydb
ip_type (str):
The IP type of the Cloud SQL instance. Can be one of "public", "private", or "psc".
refresh_strategy (Optional[str]):
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
or "background". For serverless environments use "lazy" to avoid
Expand All @@ -145,7 +151,7 @@ async def getconn(
user=user,
password=password,
db=db,
ip_type="public", # can also be "private" or "psc",
ip_type=ip_type, # can be "public", "private" or "psc"
**kwargs,
)
return conn
Expand All @@ -161,8 +167,11 @@ async def test_sqlalchemy_connection_with_asyncpg() -> None:
user = os.environ["POSTGRES_USER"]
password = os.environ["POSTGRES_PASS"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, password, db)
pool, connector = await create_sqlalchemy_engine(
inst_conn_name, user, password, db, ip_type
)

async with pool.connect() as conn:
res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone()
Expand All @@ -177,9 +186,10 @@ async def test_lazy_sqlalchemy_connection_with_asyncpg() -> None:
user = os.environ["POSTGRES_USER"]
password = os.environ["POSTGRES_PASS"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

pool, connector = await create_sqlalchemy_engine(
inst_conn_name, user, password, db, "lazy"
inst_conn_name, user, password, db, ip_type, "lazy"
)

async with pool.connect() as conn:
Expand All @@ -195,9 +205,10 @@ async def test_custom_SAN_with_dns_sqlalchemy_connection_with_asyncpg() -> None:
user = os.environ["POSTGRES_USER"]
password = os.environ["POSTGRES_CUSTOMER_CAS_PASS"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

pool, connector = await create_sqlalchemy_engine(
inst_conn_name, user, password, db, resolver=DnsResolver
inst_conn_name, user, password, db, ip_type, resolver=DnsResolver
)

async with pool.connect() as conn:
Expand All @@ -213,8 +224,11 @@ async def test_connection_with_asyncpg() -> None:
user = os.environ["POSTGRES_USER"]
password = os.environ["POSTGRES_PASS"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

pool, connector = await create_asyncpg_pool(inst_conn_name, user, password, db)
pool, connector = await create_asyncpg_pool(
inst_conn_name, user, password, db, ip_type
)

async with pool.acquire() as conn:
res = await conn.fetch("SELECT 1")
Expand All @@ -229,9 +243,10 @@ async def test_lazy_connection_with_asyncpg() -> None:
user = os.environ["POSTGRES_USER"]
password = os.environ["POSTGRES_PASS"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

pool, connector = await create_asyncpg_pool(
inst_conn_name, user, password, db, "lazy"
inst_conn_name, user, password, db, ip_type, "lazy"
)

async with pool.acquire() as conn:
Expand Down
13 changes: 10 additions & 3 deletions tests/system/test_asyncpg_iam_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ async def create_sqlalchemy_engine(
instance_connection_name: str,
user: str,
db: str,
ip_type: str = "public",
refresh_strategy: str = "background",
) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, Connector]:
"""Creates a connection pool for a Cloud SQL instance and returns the pool
Expand Down Expand Up @@ -55,6 +56,8 @@ async def create_sqlalchemy_engine(
e.g., [email protected], [email protected]
db (str):
The name of the database, e.g., mydb
ip_type (str):
The IP type of the Cloud SQL instance. Can be one of "public", "private", or "psc".
refresh_strategy (Optional[str]):
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
or "background". For serverless environments use "lazy" to avoid
Expand All @@ -71,7 +74,7 @@ async def create_sqlalchemy_engine(
"asyncpg",
user=user,
db=db,
ip_type="public", # can also be "private" or "psc"
ip_type=ip_type, # can be "public", "private" or "psc"
enable_iam_auth=True,
),
execution_options={"isolation_level": "AUTOCOMMIT"},
Expand All @@ -84,8 +87,9 @@ async def test_iam_authn_connection_with_asyncpg() -> None:
inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
user = os.environ["POSTGRES_IAM_USER"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, db)
pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, db, ip_type)

async with pool.connect() as conn:
res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone()
Expand All @@ -99,8 +103,11 @@ async def test_lazy_iam_authn_connection_with_asyncpg() -> None:
inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
user = os.environ["POSTGRES_IAM_USER"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, db, "lazy")
pool, connector = await create_sqlalchemy_engine(
inst_conn_name, user, db, ip_type, "lazy"
)

async with pool.connect() as conn:
res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone()
Expand Down
1 change: 1 addition & 0 deletions tests/system/test_connector_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ def getconn() -> pymysql.connections.Connection:
user=os.environ["MYSQL_USER"],
password=os.environ["MYSQL_PASS"],
db=os.environ["MYSQL_DB"],
ip_type=os.environ.get("IP_TYPE", "public"),
)
return conn

Expand Down
2 changes: 1 addition & 1 deletion tests/system/test_ip_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ def getconn() -> pymysql.connections.Connection:
conn: pymysql.connections.Connection = connector.connect(
os.environ["MYSQL_CONNECTION_NAME"],
"pymysql",
ip_type=ip_type,
ip_type=ip_type, # can be "public", "private" or "psc"
user=os.environ["MYSQL_USER"],
password=os.environ["MYSQL_PASS"],
db=os.environ["MYSQL_DB"],
Expand Down
26 changes: 20 additions & 6 deletions tests/system/test_pg8000_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ def create_sqlalchemy_engine(
user: str,
password: str,
db: str,
ip_type: str = "public",
refresh_strategy: str = "background",
resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver,
) -> tuple[sqlalchemy.engine.Engine, Connector]:
Expand Down Expand Up @@ -64,6 +65,8 @@ def create_sqlalchemy_engine(
The database user's password, e.g., secret-password
db (str):
The name of the database, e.g., mydb
ip_type (str):
The IP type of the Cloud SQL instance. Can be one of "public", "private", or "psc".
refresh_strategy (Optional[str]):
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
or "background". For serverless environments use "lazy" to avoid
Expand All @@ -85,7 +88,7 @@ def create_sqlalchemy_engine(
user=user,
password=password,
db=db,
ip_type="public", # can also be "private" or "psc"
ip_type=ip_type, # can be "public", "private" or "psc"
),
)
return engine, connector
Expand All @@ -100,8 +103,11 @@ def test_pg8000_connection() -> None:
user = os.environ["POSTGRES_USER"]
password = os.environ["POSTGRES_PASS"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
engine, connector = create_sqlalchemy_engine(
inst_conn_name, user, password, db, ip_type
)
with engine.connect() as conn:
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
conn.commit()
Expand All @@ -116,9 +122,10 @@ def test_lazy_pg8000_connection() -> None:
user = os.environ["POSTGRES_USER"]
password = os.environ["POSTGRES_PASS"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

engine, connector = create_sqlalchemy_engine(
inst_conn_name, user, password, db, "lazy"
inst_conn_name, user, password, db, ip_type, "lazy"
)
with engine.connect() as conn:
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
Expand All @@ -134,8 +141,11 @@ def test_CAS_pg8000_connection() -> None:
user = os.environ["POSTGRES_USER"]
password = os.environ["POSTGRES_CAS_PASS"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
engine, connector = create_sqlalchemy_engine(
inst_conn_name, user, password, db, ip_type
)
with engine.connect() as conn:
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
conn.commit()
Expand All @@ -150,8 +160,11 @@ def test_customer_managed_CAS_pg8000_connection() -> None:
user = os.environ["POSTGRES_USER"]
password = os.environ["POSTGRES_CUSTOMER_CAS_PASS"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
engine, connector = create_sqlalchemy_engine(
inst_conn_name, user, password, db, ip_type
)
with engine.connect() as conn:
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
conn.commit()
Expand All @@ -166,9 +179,10 @@ def test_custom_SAN_with_dns_pg8000_connection() -> None:
user = os.environ["POSTGRES_USER"]
password = os.environ["POSTGRES_CUSTOMER_CAS_PASS"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

engine, connector = create_sqlalchemy_engine(
inst_conn_name, user, password, db, resolver=DnsResolver
inst_conn_name, user, password, db, ip_type, resolver=DnsResolver
)
with engine.connect() as conn:
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
Expand Down
13 changes: 10 additions & 3 deletions tests/system/test_pg8000_iam_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ def create_sqlalchemy_engine(
instance_connection_name: str,
user: str,
db: str,
ip_type: str = "public",
refresh_strategy: str = "background",
) -> tuple[sqlalchemy.engine.Engine, Connector]:
"""Creates a connection pool for a Cloud SQL instance and returns the pool
Expand Down Expand Up @@ -55,6 +56,8 @@ def create_sqlalchemy_engine(
e.g., [email protected], [email protected]
db (str):
The name of the database, e.g., mydb
ip_type (str):
The IP type of the Cloud SQL instance. Can be one of "public", "private", or "psc".
refresh_strategy (Optional[str]):
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
or "background". For serverless environments use "lazy" to avoid
Expand All @@ -70,7 +73,7 @@ def create_sqlalchemy_engine(
"pg8000",
user=user,
db=db,
ip_type="public", # can also be "private" or "psc"
ip_type=ip_type, # can be "public", "private" or "psc"
enable_iam_auth=True,
),
)
Expand All @@ -82,8 +85,9 @@ def test_pg8000_iam_authn_connection() -> None:
inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
user = os.environ["POSTGRES_IAM_USER"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db)
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db, ip_type)
with engine.connect() as conn:
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
conn.commit()
Expand All @@ -97,8 +101,11 @@ def test_lazy_pg8000_iam_authn_connection() -> None:
inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
user = os.environ["POSTGRES_IAM_USER"]
db = os.environ["POSTGRES_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db, "lazy")
engine, connector = create_sqlalchemy_engine(
inst_conn_name, user, db, ip_type, "lazy"
)
with engine.connect() as conn:
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
conn.commit()
Expand Down
13 changes: 10 additions & 3 deletions tests/system/test_pymysql_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ def create_sqlalchemy_engine(
user: str,
password: str,
db: str,
ip_type: str = "public",
refresh_strategy: str = "background",
) -> tuple[sqlalchemy.engine.Engine, Connector]:
"""Creates a connection pool for a Cloud SQL instance and returns the pool
Expand Down Expand Up @@ -59,6 +60,8 @@ def create_sqlalchemy_engine(
The database user's password, e.g., secret-password
db (str):
The name of the database, e.g., mydb
ip_type (str):
The IP type of the Cloud SQL instance. Can be one of "public", "private", or "psc".
refresh_strategy (Optional[str]):
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
or "background". For serverless environments use "lazy" to avoid
Expand All @@ -75,7 +78,7 @@ def create_sqlalchemy_engine(
user=user,
password=password,
db=db,
ip_type="public", # can also be "private" or "psc"
ip_type=ip_type, # can be "public", "private" or "psc"
),
)
return engine, connector
Expand All @@ -90,8 +93,11 @@ def test_pymysql_connection() -> None:
user = os.environ["MYSQL_USER"]
password = os.environ["MYSQL_PASS"]
db = os.environ["MYSQL_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
engine, connector = create_sqlalchemy_engine(
inst_conn_name, user, password, db, ip_type
)
with engine.connect() as conn:
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
conn.commit()
Expand All @@ -106,9 +112,10 @@ def test_lazy_pymysql_connection() -> None:
user = os.environ["MYSQL_USER"]
password = os.environ["MYSQL_PASS"]
db = os.environ["MYSQL_DB"]
ip_type = os.environ.get("IP_TYPE", "public")

engine, connector = create_sqlalchemy_engine(
inst_conn_name, user, password, db, "lazy"
inst_conn_name, user, password, db, ip_type, "lazy"
)
with engine.connect() as conn:
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
Expand Down
Loading
Loading