Skip to content

Commit 8b6306b

Browse files
chandr-andrchandr-andr
authored
Merge pull request #37 from qaspen-python/feature/more_connect_parameters
Added more parameters to configure connection
2 parents 7773bc5 + ac5366d commit 8b6306b

File tree

8 files changed

+530
-46
lines changed

8 files changed

+530
-46
lines changed

‎python/psqlpy/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,11 @@
44
ConnRecyclingMethod,
55
Cursor,
66
IsolationLevel,
7+
LoadBalanceHosts,
78
QueryResult,
89
ReadVariant,
910
SingleQueryResult,
11+
TargetSessionAttrs,
1012
Transaction,
1113
connect,
1214
)
@@ -22,4 +24,6 @@
2224
"IsolationLevel",
2325
"ReadVariant",
2426
"connect",
27+
"LoadBalanceHosts",
28+
"TargetSessionAttrs",
2529
]

‎python/psqlpy/_internal/__init__.pyi‎

Lines changed: 171 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import types
22
from enum import Enum
3-
from typing import Any, Callable, Optional, TypeVar
3+
from typing import Any, Callable, List, Optional, TypeVar
44

55
from typing_extensions import Self
66

@@ -104,6 +104,24 @@ class IsolationLevel(Enum):
104104
RepeatableRead = 3
105105
Serializable = 4
106106

107+
class LoadBalanceHosts(Enum):
108+
"""Load balancing configuration."""
109+
110+
# Make connection attempts to hosts in the order provided.
111+
Disable = 1
112+
# Make connection attempts to hosts in a random order.
113+
Random = 2
114+
115+
class TargetSessionAttrs(Enum):
116+
"""Properties required of a session."""
117+
118+
# No special properties are required.
119+
Any = 1
120+
# The session must allow writes.
121+
ReadWrite = 2
122+
# The session allow only reads.
123+
ReadOnly = 3
124+
107125
class ReadVariant(Enum):
108126
"""Class for Read Variant for transaction."""
109127

@@ -869,8 +887,24 @@ class ConnectionPool:
869887
username: Optional[str] = None,
870888
password: Optional[str] = None,
871889
host: Optional[str] = None,
890+
hosts: Optional[List[str]] = None,
872891
port: Optional[int] = None,
892+
ports: Optional[List[int]] = None,
873893
db_name: Optional[str] = None,
894+
target_session_attrs: Optional[TargetSessionAttrs] = None,
895+
options: Optional[str] = None,
896+
application_name: Optional[str] = None,
897+
connect_timeout_sec: Optional[int] = None,
898+
connect_timeout_nanosec: Optional[int] = None,
899+
tcp_user_timeout_sec: Optional[int] = None,
900+
tcp_user_timeout_nanosec: Optional[int] = None,
901+
keepalives: Optional[bool] = None,
902+
keepalives_idle_sec: Optional[int] = None,
903+
keepalives_idle_nanosec: Optional[int] = None,
904+
keepalives_interval_sec: Optional[int] = None,
905+
keepalives_interval_nanosec: Optional[int] = None,
906+
keepalives_retries: Optional[int] = None,
907+
load_balance_hosts: Optional[LoadBalanceHosts] = None,
874908
max_db_pool_size: int = 2,
875909
conn_recycling_method: Optional[ConnRecyclingMethod] = None,
876910
) -> None:
@@ -879,22 +913,67 @@ class ConnectionPool:
879913
It connects to the database and create pool.
880914
881915
You cannot set the minimum size for the connection
882-
pool, by default it is 1.
916+
pool, by it is 0.
917+
`ConnectionPool` doesn't create connections on startup.
918+
It makes new connection on demand.
883919
884-
This connection pool can:
885-
- Startup itself with `startup` method
886-
- Execute queries and return `QueryResult` class as a result
887-
- Create new instance of `Transaction`
920+
If you specify `dsn` parameter then `username`, `password`,
921+
`host`, `hosts`, `port`, `ports`, `db_name` and `target_session_attrs`
922+
parameters will be ignored.
888923
889924
### Parameters:
890-
- `dsn`: full dsn connection string.
925+
- `dsn`: Full dsn connection string.
891926
`postgres://postgres:postgres@localhost:5432/postgres?target_session_attrs=read-write`
892-
- `username`: username of the user in postgres
893-
- `password`: password of the user in postgres
894-
- `host`: host of postgres
895-
- `port`: port of postgres
896-
- `db_name`: name of the database in postgres
897-
- `max_db_pool_size`: maximum size of the connection pool
927+
- `username`: Username of the user in the PostgreSQL
928+
- `password`: Password of the user in the PostgreSQL
929+
- `host`: Host of the PostgreSQL
930+
- `hosts`: Hosts of the PostgreSQL
931+
- `port`: Port of the PostgreSQL
932+
- `ports`: Ports of the PostgreSQL
933+
- `db_name`: Name of the database in PostgreSQL
934+
- `target_session_attrs`: Specifies requirements of the session.
935+
- `options`: Command line options used to configure the server
936+
- `application_name`: Sets the application_name parameter on the server.
937+
- `connect_timeout_sec`: The time limit in seconds applied to each socket-level
938+
connection attempt.
939+
Note that hostnames can resolve to multiple IP addresses,
940+
and this limit is applied to each address. Defaults to no timeout.
941+
- `connect_timeout_nanosec`: nanosec for connection timeout,
942+
can be used only with connect_timeout_sec.
943+
- `tcp_user_timeout_sec`: The time limit that
944+
transmitted data may remain unacknowledged
945+
before a connection is forcibly closed.
946+
This is ignored for Unix domain socket connections.
947+
It is only supported on systems where TCP_USER_TIMEOUT
948+
is available and will default to the system default if omitted
949+
or set to 0; on other systems, it has no effect.
950+
- `tcp_user_timeout_nanosec`: nanosec for cp_user_timeout,
951+
can be used only with tcp_user_timeout_sec.
952+
- `keepalives`: Controls the use of TCP keepalive.
953+
This option is ignored when connecting with Unix sockets.
954+
Defaults to on.
955+
- `keepalives_idle_sec`: The number of seconds of inactivity after
956+
which a keepalive message is sent to the server.
957+
This option is ignored when connecting with Unix sockets.
958+
Defaults to 2 hours.
959+
- `keepalives_idle_nanosec`: Nanosec for keepalives_idle_sec.
960+
- `keepalives_interval_sec`: The time interval between TCP keepalive probes.
961+
This option is ignored when connecting with Unix sockets.
962+
- `keepalives_interval_nanosec`: Nanosec for keepalives_interval_sec.
963+
- `keepalives_retries`: The maximum number of TCP keepalive probes
964+
that will be sent before dropping a connection.
965+
This option is ignored when connecting with Unix sockets.
966+
- `load_balance_hosts`: Controls the order in which the client tries to connect
967+
to the available hosts and addresses.
968+
Once a connection attempt is successful no other
969+
hosts and addresses will be tried.
970+
This parameter is typically used in combination with multiple host names
971+
or a DNS record that returns multiple IPs.
972+
If set to disable, hosts and addresses will be tried in the order provided.
973+
If set to random, hosts will be tried in a random order, and the IP addresses
974+
resolved from a hostname will also be tried in a random order.
975+
Defaults to disable.
976+
- `max_db_pool_size`: maximum size of the connection pool.
898977
- `conn_recycling_method`: how a connection is recycled.
899978
"""
900979
async def execute(
@@ -945,21 +1024,92 @@ def connect(
9451024
username: Optional[str] = None,
9461025
password: Optional[str] = None,
9471026
host: Optional[str] = None,
1027+
hosts: Optional[List[str]] = None,
9481028
port: Optional[int] = None,
1029+
ports: Optional[List[int]] = None,
9491030
db_name: Optional[str] = None,
1031+
target_session_attrs: Optional[TargetSessionAttrs] = None,
1032+
options: Optional[str] = None,
1033+
application_name: Optional[str] = None,
1034+
connect_timeout_sec: Optional[int] = None,
1035+
connect_timeout_nanosec: Optional[int] = None,
1036+
tcp_user_timeout_sec: Optional[int] = None,
1037+
tcp_user_timeout_nanosec: Optional[int] = None,
1038+
keepalives: Optional[bool] = None,
1039+
keepalives_idle_sec: Optional[int] = None,
1040+
keepalives_idle_nanosec: Optional[int] = None,
1041+
keepalives_interval_sec: Optional[int] = None,
1042+
keepalives_interval_nanosec: Optional[int] = None,
1043+
keepalives_retries: Optional[int] = None,
1044+
load_balance_hosts: Optional[LoadBalanceHosts] = None,
9501045
max_db_pool_size: int = 2,
9511046
conn_recycling_method: Optional[ConnRecyclingMethod] = None,
9521047
) -> ConnectionPool:
953-
"""Create new connection pool.
1048+
"""Create new PostgreSQL connection pool.
1049+
1050+
It connects to the database and create pool.
1051+
1052+
You cannot set the minimum size for the connection
1053+
pool, by it is 0.
1054+
`ConnectionPool` doesn't create connections on startup.
1055+
It makes new connection on demand.
1056+
1057+
If you specify `dsn` parameter then `username`, `password`,
1058+
`host`, `hosts`, `port`, `ports`, `db_name` and `target_session_attrs`
1059+
parameters will be ignored.
9541060
9551061
### Parameters:
956-
- `dsn`: full dsn connection string.
1062+
- `dsn`: Full dsn connection string.
9571063
`postgres://postgres:postgres@localhost:5432/postgres?target_session_attrs=read-write`
958-
- `username`: username of the user in postgres
959-
- `password`: password of the user in postgres
960-
- `host`: host of postgres
961-
- `port`: port of postgres
962-
- `db_name`: name of the database in postgres
963-
- `max_db_pool_size`: maximum size of the connection pool
1064+
- `username`: Username of the user in the PostgreSQL
1065+
- `password`: Password of the user in the PostgreSQL
1066+
- `host`: Host of the PostgreSQL
1067+
- `hosts`: Hosts of the PostgreSQL
1068+
- `port`: Port of the PostgreSQL
1069+
- `ports`: Ports of the PostgreSQL
1070+
- `db_name`: Name of the database in PostgreSQL
1071+
- `target_session_attrs`: Specifies requirements of the session.
1072+
- `options`: Command line options used to configure the server
1073+
- `application_name`: Sets the application_name parameter on the server.
1074+
- `connect_timeout_sec`: The time limit in seconds applied to each socket-level
1075+
connection attempt.
1076+
Note that hostnames can resolve to multiple IP addresses,
1077+
and this limit is applied to each address. Defaults to no timeout.
1078+
- `connect_timeout_nanosec`: nanosec for connection timeout,
1079+
can be used only with connect_timeout_sec.
1080+
- `tcp_user_timeout_sec`: The time limit that
1081+
transmitted data may remain unacknowledged
1082+
before a connection is forcibly closed.
1083+
This is ignored for Unix domain socket connections.
1084+
It is only supported on systems where TCP_USER_TIMEOUT
1085+
is available and will default to the system default if omitted
1086+
or set to 0; on other systems, it has no effect.
1087+
- `tcp_user_timeout_nanosec`: nanosec for cp_user_timeout,
1088+
can be used only with tcp_user_timeout_sec.
1089+
- `keepalives`: Controls the use of TCP keepalive.
1090+
This option is ignored when connecting with Unix sockets.
1091+
Defaults to on.
1092+
- `keepalives_idle_sec`: The number of seconds of inactivity after
1093+
which a keepalive message is sent to the server.
1094+
This option is ignored when connecting with Unix sockets.
1095+
Defaults to 2 hours.
1096+
- `keepalives_idle_nanosec`: Nanosec for keepalives_idle_sec.
1097+
- `keepalives_interval_sec`: The time interval between TCP keepalive probes.
1098+
This option is ignored when connecting with Unix sockets.
1099+
- `keepalives_interval_nanosec`: Nanosec for keepalives_interval_sec.
1100+
- `keepalives_retries`: The maximum number of TCP keepalive probes
1101+
that will be sent before dropping a connection.
1102+
This option is ignored when connecting with Unix sockets.
1103+
- `load_balance_hosts`: Controls the order in which the client tries to connect
1104+
to the available hosts and addresses.
1105+
Once a connection attempt is successful no other
1106+
hosts and addresses will be tried.
1107+
This parameter is typically used in combination with multiple host names
1108+
or a DNS record that returns multiple IPs.
1109+
If set to disable, hosts and addresses will be tried in the order provided.
1110+
If set to random, hosts will be tried in a random order, and the IP addresses
1111+
resolved from a hostname will also be tried in a random order.
1112+
Defaults to disable.
1113+
- `max_db_pool_size`: maximum size of the connection pool.
9641114
- `conn_recycling_method`: how a connection is recycled.
9651115
"""

‎python/tests/test_connection_pool.py‎

Lines changed: 77 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,15 @@
11
import pytest
22

3-
from psqlpy import Connection, ConnectionPool, ConnRecyclingMethod, QueryResult, connect
4-
from psqlpy.exceptions import RustPSQLDriverPyBaseError
3+
from psqlpy import (
4+
Connection,
5+
ConnectionPool,
6+
ConnRecyclingMethod,
7+
LoadBalanceHosts,
8+
QueryResult,
9+
TargetSessionAttrs,
10+
connect,
11+
)
12+
from psqlpy.exceptions import DBPoolConfigurationError, RustPSQLDriverPyBaseError
513

614
pytestmark = pytest.mark.anyio
715

@@ -68,6 +76,73 @@ async def test_pool_conn_recycling_method(
6876
await pg_pool.execute("SELECT 1")
6977

7078

79+
async def test_build_pool_failure() -> None:
80+
with pytest.raises(expected_exception=DBPoolConfigurationError):
81+
ConnectionPool(
82+
dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test",
83+
connect_timeout_nanosec=12,
84+
)
85+
with pytest.raises(expected_exception=DBPoolConfigurationError):
86+
ConnectionPool(
87+
dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test",
88+
connect_timeout_nanosec=12,
89+
)
90+
with pytest.raises(expected_exception=DBPoolConfigurationError):
91+
ConnectionPool(
92+
dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test",
93+
keepalives_idle_nanosec=12,
94+
)
95+
with pytest.raises(expected_exception=DBPoolConfigurationError):
96+
ConnectionPool(
97+
dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test",
98+
keepalives_interval_nanosec=12,
99+
)
100+
101+
102+
@pytest.mark.parametrize(
103+
"target_session_attrs",
104+
[
105+
TargetSessionAttrs.Any,
106+
TargetSessionAttrs.ReadWrite,
107+
TargetSessionAttrs.ReadOnly,
108+
],
109+
)
110+
async def test_pool_target_session_attrs(
111+
target_session_attrs: TargetSessionAttrs,
112+
) -> None:
113+
pg_pool = ConnectionPool(
114+
db_name="psqlpy_test",
115+
host="localhost",
116+
username="postgres",
117+
password="postgres", # noqa: S106
118+
target_session_attrs=target_session_attrs,
119+
)
120+
121+
if target_session_attrs == TargetSessionAttrs.ReadOnly:
122+
with pytest.raises(expected_exception=RustPSQLDriverPyBaseError):
123+
await pg_pool.execute("SELECT 1")
124+
else:
125+
await pg_pool.execute("SELECT 1")
126+
127+
128+
@pytest.mark.parametrize(
129+
"load_balance_hosts",
130+
[
131+
LoadBalanceHosts.Disable,
132+
LoadBalanceHosts.Random,
133+
],
134+
)
135+
async def test_pool_load_balance_hosts(
136+
load_balance_hosts: LoadBalanceHosts,
137+
) -> None:
138+
pg_pool = ConnectionPool(
139+
dsn="postgres://postgres:postgres@localhost:5432/psqlpy_test",
140+
load_balance_hosts=load_balance_hosts,
141+
)
142+
143+
await pg_pool.execute("SELECT 1")
144+
145+
71146
async def test_close_connection_pool() -> None:
72147
"""Test that `close` method closes connection pool."""
73148
pg_pool = ConnectionPool(

‎src/driver/common_options.rs‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,3 +19,44 @@ impl ConnRecyclingMethod {
1919
}
2020
}
2121
}
22+
23+
#[pyclass]
24+
#[derive(Clone, Copy)]
25+
pub enum LoadBalanceHosts {
26+
/// Make connection attempts to hosts in the order provided.
27+
Disable,
28+
/// Make connection attempts to hosts in a random order.
29+
Random,
30+
}
31+
32+
impl LoadBalanceHosts {
33+
#[must_use]
34+
pub fn to_internal(&self) -> tokio_postgres::config::LoadBalanceHosts {
35+
match self {
36+
LoadBalanceHosts::Disable => tokio_postgres::config::LoadBalanceHosts::Disable,
37+
LoadBalanceHosts::Random => tokio_postgres::config::LoadBalanceHosts::Random,
38+
}
39+
}
40+
}
41+
42+
#[pyclass]
43+
#[derive(Clone, Copy)]
44+
pub enum TargetSessionAttrs {
45+
/// No special properties are required.
46+
Any,
47+
/// The session must allow writes.
48+
ReadWrite,
49+
/// The session allow only reads.
50+
ReadOnly,
51+
}
52+
53+
impl TargetSessionAttrs {
54+
#[must_use]
55+
pub fn to_internal(&self) -> tokio_postgres::config::TargetSessionAttrs {
56+
match self {
57+
TargetSessionAttrs::Any => tokio_postgres::config::TargetSessionAttrs::Any,
58+
TargetSessionAttrs::ReadWrite => tokio_postgres::config::TargetSessionAttrs::ReadWrite,
59+
TargetSessionAttrs::ReadOnly => tokio_postgres::config::TargetSessionAttrs::ReadOnly,
60+
}
61+
}
62+
}

0 commit comments

Comments
 (0)