Skip to content

Commit ccd414c

Browse files
kgala2kgala2
authored
test: update ip type for connectors (GoogleCloudPlatform#1256)
1 parent e465eea commit ccd414c

8 files changed

+101
-28
lines changed

‎tests/system/test_asyncpg_connection.py‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ async def create_sqlalchemy_engine(
3232
user: str,
3333
password: str,
3434
db: str,
35+
ip_type: str = "public",
3536
refresh_strategy: str = "background",
3637
resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver,
3738
) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, Connector]:
@@ -63,6 +64,9 @@ async def create_sqlalchemy_engine(
6364
The database user's password, e.g., secret-password
6465
db (str):
6566
The name of the database, e.g., mydb
67+
ip_type (str):
68+
The IP type of the Cloud SQL instance to connect to. Can be one
69+
of "public", "private", or "psc".
6670
refresh_strategy (Optional[str]):
6771
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
6872
or "background". For serverless environments use "lazy" to avoid
@@ -87,7 +91,7 @@ async def create_sqlalchemy_engine(
8791
user=user,
8892
password=password,
8993
db=db,
90-
ip_type="public", # can also be "private" or "psc"
94+
ip_type=ip_type, # can be "public", "private" or "psc"
9195
),
9296
execution_options={"isolation_level": "AUTOCOMMIT"},
9397
)
@@ -99,6 +103,7 @@ async def create_asyncpg_pool(
99103
user: str,
100104
password: str,
101105
db: str,
106+
ip_type: str = "public",
102107
refresh_strategy: str = "background",
103108
) -> tuple[asyncpg.Pool, Connector]:
104109
"""Creates a native asyncpg connection pool for a Cloud SQL instance and
@@ -128,6 +133,9 @@ async def create_asyncpg_pool(
128133
The database user's password, e.g., secret-password
129134
db (str):
130135
The name of the database, e.g., mydb
136+
ip_type (str):
137+
The IP type of the Cloud SQL instance to connect to. Can be one
138+
of "public", "private", or "psc".
131139
refresh_strategy (Optional[str]):
132140
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
133141
or "background". For serverless environments use "lazy" to avoid
@@ -145,7 +153,7 @@ async def getconn(
145153
user=user,
146154
password=password,
147155
db=db,
148-
ip_type="public", # can also be "private" or "psc",
156+
ip_type=ip_type, # can be "public", "private" or "psc"
149157
**kwargs,
150158
)
151159
return conn
@@ -161,8 +169,11 @@ async def test_sqlalchemy_connection_with_asyncpg() -> None:
161169
user = os.environ["POSTGRES_USER"]
162170
password = os.environ["POSTGRES_PASS"]
163171
db = os.environ["POSTGRES_DB"]
172+
ip_type = os.environ.get("IP_TYPE", "public")
164173

165-
pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, password, db)
174+
pool, connector = await create_sqlalchemy_engine(
175+
inst_conn_name, user, password, db, ip_type
176+
)
166177

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

181193
pool, connector = await create_sqlalchemy_engine(
182-
inst_conn_name, user, password, db, "lazy"
194+
inst_conn_name, user, password, db, ip_type, "lazy"
183195
)
184196

185197
async with pool.connect() as conn:
@@ -195,9 +207,10 @@ async def test_custom_SAN_with_dns_sqlalchemy_connection_with_asyncpg() -> None:
195207
user = os.environ["POSTGRES_USER"]
196208
password = os.environ["POSTGRES_CUSTOMER_CAS_PASS"]
197209
db = os.environ["POSTGRES_DB"]
210+
ip_type = os.environ.get("IP_TYPE", "public")
198211

199212
pool, connector = await create_sqlalchemy_engine(
200-
inst_conn_name, user, password, db, resolver=DnsResolver
213+
inst_conn_name, user, password, db, ip_type, resolver=DnsResolver
201214
)
202215

203216
async with pool.connect() as conn:
@@ -213,8 +226,11 @@ async def test_connection_with_asyncpg() -> None:
213226
user = os.environ["POSTGRES_USER"]
214227
password = os.environ["POSTGRES_PASS"]
215228
db = os.environ["POSTGRES_DB"]
229+
ip_type = os.environ.get("IP_TYPE", "public")
216230

217-
pool, connector = await create_asyncpg_pool(inst_conn_name, user, password, db)
231+
pool, connector = await create_asyncpg_pool(
232+
inst_conn_name, user, password, db, ip_type
233+
)
218234

219235
async with pool.acquire() as conn:
220236
res = await conn.fetch("SELECT 1")
@@ -229,9 +245,10 @@ async def test_lazy_connection_with_asyncpg() -> None:
229245
user = os.environ["POSTGRES_USER"]
230246
password = os.environ["POSTGRES_PASS"]
231247
db = os.environ["POSTGRES_DB"]
248+
ip_type = os.environ.get("IP_TYPE", "public")
232249

233250
pool, connector = await create_asyncpg_pool(
234-
inst_conn_name, user, password, db, "lazy"
251+
inst_conn_name, user, password, db, ip_type, "lazy"
235252
)
236253

237254
async with pool.acquire() as conn:

‎tests/system/test_asyncpg_iam_auth.py‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ async def create_sqlalchemy_engine(
2727
instance_connection_name: str,
2828
user: str,
2929
db: str,
30+
ip_type: str = "public",
3031
refresh_strategy: str = "background",
3132
) -> tuple[sqlalchemy.ext.asyncio.engine.AsyncEngine, Connector]:
3233
"""Creates a connection pool for a Cloud SQL instance and returns the pool
@@ -55,6 +56,9 @@ async def create_sqlalchemy_engine(
5556
e.g., [email protected], [email protected]
5657
db (str):
5758
The name of the database, e.g., mydb
59+
ip_type (str):
60+
The IP type of the Cloud SQL instance to connect to. Can be one
61+
of "public", "private", or "psc".
5862
refresh_strategy (Optional[str]):
5963
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
6064
or "background". For serverless environments use "lazy" to avoid
@@ -71,7 +75,7 @@ async def create_sqlalchemy_engine(
7175
"asyncpg",
7276
user=user,
7377
db=db,
74-
ip_type="public", # can also be "private" or "psc"
78+
ip_type=ip_type, # can be "public", "private" or "psc"
7579
enable_iam_auth=True,
7680
),
7781
execution_options={"isolation_level": "AUTOCOMMIT"},
@@ -84,8 +88,9 @@ async def test_iam_authn_connection_with_asyncpg() -> None:
8488
inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
8589
user = os.environ["POSTGRES_IAM_USER"]
8690
db = os.environ["POSTGRES_DB"]
91+
ip_type = os.environ.get("IP_TYPE", "public")
8792

88-
pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, db)
93+
pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, db, ip_type)
8994

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

103-
pool, connector = await create_sqlalchemy_engine(inst_conn_name, user, db, "lazy")
109+
pool, connector = await create_sqlalchemy_engine(
110+
inst_conn_name, user, db, ip_type, "lazy"
111+
)
104112

105113
async with pool.connect() as conn:
106114
res = (await conn.execute(sqlalchemy.text("SELECT 1"))).fetchone()

‎tests/system/test_connector_object.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ def getconn() -> pymysql.connections.Connection:
3838
user=os.environ["MYSQL_USER"],
3939
password=os.environ["MYSQL_PASS"],
4040
db=os.environ["MYSQL_DB"],
41+
ip_type=os.environ.get("IP_TYPE", "public"),
4142
)
4243
return conn
4344

‎tests/system/test_pg8000_connection.py‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ def create_sqlalchemy_engine(
3232
user: str,
3333
password: str,
3434
db: str,
35+
ip_type: str = "public",
3536
refresh_strategy: str = "background",
3637
resolver: Union[type[DefaultResolver], type[DnsResolver]] = DefaultResolver,
3738
) -> tuple[sqlalchemy.engine.Engine, Connector]:
@@ -64,6 +65,9 @@ def create_sqlalchemy_engine(
6465
The database user's password, e.g., secret-password
6566
db (str):
6667
The name of the database, e.g., mydb
68+
ip_type (str):
69+
The IP type of the Cloud SQL instance to connect to. Can be one
70+
of "public", "private", or "psc".
6771
refresh_strategy (Optional[str]):
6872
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
6973
or "background". For serverless environments use "lazy" to avoid
@@ -85,7 +89,7 @@ def create_sqlalchemy_engine(
8589
user=user,
8690
password=password,
8791
db=db,
88-
ip_type="public", # can also be "private" or "psc"
92+
ip_type=ip_type, # can be "public", "private" or "psc"
8993
),
9094
)
9195
return engine, connector
@@ -100,8 +104,11 @@ def test_pg8000_connection() -> None:
100104
user = os.environ["POSTGRES_USER"]
101105
password = os.environ["POSTGRES_PASS"]
102106
db = os.environ["POSTGRES_DB"]
107+
ip_type = os.environ.get("IP_TYPE", "public")
103108

104-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
109+
engine, connector = create_sqlalchemy_engine(
110+
inst_conn_name, user, password, db, ip_type
111+
)
105112
with engine.connect() as conn:
106113
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
107114
conn.commit()
@@ -116,9 +123,10 @@ def test_lazy_pg8000_connection() -> None:
116123
user = os.environ["POSTGRES_USER"]
117124
password = os.environ["POSTGRES_PASS"]
118125
db = os.environ["POSTGRES_DB"]
126+
ip_type = os.environ.get("IP_TYPE", "public")
119127

120128
engine, connector = create_sqlalchemy_engine(
121-
inst_conn_name, user, password, db, "lazy"
129+
inst_conn_name, user, password, db, ip_type, "lazy"
122130
)
123131
with engine.connect() as conn:
124132
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
@@ -134,8 +142,11 @@ def test_CAS_pg8000_connection() -> None:
134142
user = os.environ["POSTGRES_USER"]
135143
password = os.environ["POSTGRES_CAS_PASS"]
136144
db = os.environ["POSTGRES_DB"]
145+
ip_type = os.environ.get("IP_TYPE", "public")
137146

138-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
147+
engine, connector = create_sqlalchemy_engine(
148+
inst_conn_name, user, password, db, ip_type
149+
)
139150
with engine.connect() as conn:
140151
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
141152
conn.commit()
@@ -150,8 +161,11 @@ def test_customer_managed_CAS_pg8000_connection() -> None:
150161
user = os.environ["POSTGRES_USER"]
151162
password = os.environ["POSTGRES_CUSTOMER_CAS_PASS"]
152163
db = os.environ["POSTGRES_DB"]
164+
ip_type = os.environ.get("IP_TYPE", "public")
153165

154-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
166+
engine, connector = create_sqlalchemy_engine(
167+
inst_conn_name, user, password, db, ip_type
168+
)
155169
with engine.connect() as conn:
156170
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
157171
conn.commit()
@@ -166,9 +180,10 @@ def test_custom_SAN_with_dns_pg8000_connection() -> None:
166180
user = os.environ["POSTGRES_USER"]
167181
password = os.environ["POSTGRES_CUSTOMER_CAS_PASS"]
168182
db = os.environ["POSTGRES_DB"]
183+
ip_type = os.environ.get("IP_TYPE", "public")
169184

170185
engine, connector = create_sqlalchemy_engine(
171-
inst_conn_name, user, password, db, resolver=DnsResolver
186+
inst_conn_name, user, password, db, ip_type, resolver=DnsResolver
172187
)
173188
with engine.connect() as conn:
174189
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()

‎tests/system/test_pg8000_iam_auth.py‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ def create_sqlalchemy_engine(
2626
instance_connection_name: str,
2727
user: str,
2828
db: str,
29+
ip_type: str = "public",
2930
refresh_strategy: str = "background",
3031
) -> tuple[sqlalchemy.engine.Engine, Connector]:
3132
"""Creates a connection pool for a Cloud SQL instance and returns the pool
@@ -55,6 +56,9 @@ def create_sqlalchemy_engine(
5556
e.g., [email protected], [email protected]
5657
db (str):
5758
The name of the database, e.g., mydb
59+
ip_type (str):
60+
The IP type of the Cloud SQL instance to connect to. Can be one
61+
of "public", "private", or "psc".
5862
refresh_strategy (Optional[str]):
5963
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
6064
or "background". For serverless environments use "lazy" to avoid
@@ -70,7 +74,7 @@ def create_sqlalchemy_engine(
7074
"pg8000",
7175
user=user,
7276
db=db,
73-
ip_type="public", # can also be "private" or "psc"
77+
ip_type=ip_type, # can be "public", "private" or "psc"
7478
enable_iam_auth=True,
7579
),
7680
)
@@ -82,8 +86,9 @@ def test_pg8000_iam_authn_connection() -> None:
8286
inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
8387
user = os.environ["POSTGRES_IAM_USER"]
8488
db = os.environ["POSTGRES_DB"]
89+
ip_type = os.environ.get("IP_TYPE", "public")
8590

86-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db)
91+
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db, ip_type)
8792
with engine.connect() as conn:
8893
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
8994
conn.commit()
@@ -97,8 +102,11 @@ def test_lazy_pg8000_iam_authn_connection() -> None:
97102
inst_conn_name = os.environ["POSTGRES_CONNECTION_NAME"]
98103
user = os.environ["POSTGRES_IAM_USER"]
99104
db = os.environ["POSTGRES_DB"]
105+
ip_type = os.environ.get("IP_TYPE", "public")
100106

101-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, db, "lazy")
107+
engine, connector = create_sqlalchemy_engine(
108+
inst_conn_name, user, db, ip_type, "lazy"
109+
)
102110
with engine.connect() as conn:
103111
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
104112
conn.commit()

‎tests/system/test_pymysql_connection.py‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ def create_sqlalchemy_engine(
2828
user: str,
2929
password: str,
3030
db: str,
31+
ip_type: str = "public",
3132
refresh_strategy: str = "background",
3233
) -> tuple[sqlalchemy.engine.Engine, Connector]:
3334
"""Creates a connection pool for a Cloud SQL instance and returns the pool
@@ -59,6 +60,9 @@ def create_sqlalchemy_engine(
5960
The database user's password, e.g., secret-password
6061
db (str):
6162
The name of the database, e.g., mydb
63+
ip_type (str):
64+
The IP type of the Cloud SQL instance to connect to. Can be one
65+
of "public", "private", or "psc".
6266
refresh_strategy (Optional[str]):
6367
Refresh strategy for the Cloud SQL Connector. Can be one of "lazy"
6468
or "background". For serverless environments use "lazy" to avoid
@@ -75,7 +79,7 @@ def create_sqlalchemy_engine(
7579
user=user,
7680
password=password,
7781
db=db,
78-
ip_type="public", # can also be "private" or "psc"
82+
ip_type=ip_type, # can be "public", "private" or "psc"
7983
),
8084
)
8185
return engine, connector
@@ -90,8 +94,11 @@ def test_pymysql_connection() -> None:
9094
user = os.environ["MYSQL_USER"]
9195
password = os.environ["MYSQL_PASS"]
9296
db = os.environ["MYSQL_DB"]
97+
ip_type = os.environ.get("IP_TYPE", "public")
9398

94-
engine, connector = create_sqlalchemy_engine(inst_conn_name, user, password, db)
99+
engine, connector = create_sqlalchemy_engine(
100+
inst_conn_name, user, password, db, ip_type
101+
)
95102
with engine.connect() as conn:
96103
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()
97104
conn.commit()
@@ -106,9 +113,10 @@ def test_lazy_pymysql_connection() -> None:
106113
user = os.environ["MYSQL_USER"]
107114
password = os.environ["MYSQL_PASS"]
108115
db = os.environ["MYSQL_DB"]
116+
ip_type = os.environ.get("IP_TYPE", "public")
109117

110118
engine, connector = create_sqlalchemy_engine(
111-
inst_conn_name, user, password, db, "lazy"
119+
inst_conn_name, user, password, db, ip_type, "lazy"
112120
)
113121
with engine.connect() as conn:
114122
time = conn.execute(sqlalchemy.text("SELECT NOW()")).fetchone()

0 commit comments

Comments
 (0)