mirror of
https://github.com/postgresml/pgcat.git
synced 2026-03-23 09:26:30 +00:00
Compare commits
10 Commits
levkk-fix-
...
levkk-star
| Author | SHA1 | Date | |
|---|---|---|---|
|
ee23b374ae | ||
|
|
9dffebccbf | ||
|
|
4c8358b8b3 | ||
|
|
f0d1916a98 | ||
|
|
bba5f10be1 | ||
|
|
a514dbc187 | ||
|
|
d660e3e565 | ||
|
|
0d882cc204 | ||
|
|
b36746a47b | ||
|
|
9e51b8110f |
16
CONFIG.md
16
CONFIG.md
@@ -188,22 +188,6 @@ default: "admin_pass"
|
||||
|
||||
Password to access the virtual administrative database
|
||||
|
||||
### dns_cache_enabled
|
||||
```
|
||||
path: general.dns_cache_enabled
|
||||
default: false
|
||||
```
|
||||
When enabled, ip resolutions for server connections specified using hostnames will be cached
|
||||
and checked for changes every `dns_max_ttl` seconds. If a change in the host resolution is found
|
||||
old ip connections are closed (gracefully) and new connections will start using new ip.
|
||||
|
||||
### dns_max_ttl
|
||||
```
|
||||
path: general.dns_max_ttl
|
||||
default: 30
|
||||
```
|
||||
Specifies how often (in seconds) cached ip addresses for servers are rechecked (see `dns_cache_enabled`).
|
||||
|
||||
## `pools.<pool_name>` Section
|
||||
|
||||
### pool_mode
|
||||
|
||||
390
Cargo.lock
generated
390
Cargo.lock
generated
@@ -26,27 +26,6 @@ version = "1.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bddcadddf5e9015d310179a59bb28c4d4b9920ad0f11e8e14dbadf654890c9a6"
|
||||
|
||||
[[package]]
|
||||
name = "async-stream"
|
||||
version = "0.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dad5c83079eae9969be7fadefe640a1c566901f05ff91ab221de4b6f68d9507e"
|
||||
dependencies = [
|
||||
"async-stream-impl",
|
||||
"futures-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream-impl"
|
||||
version = "0.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "10f203db73a71dfa2fb6dd22763990fa26f3d2625a6da2da900d23b87d26be27"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-trait"
|
||||
version = "0.1.68"
|
||||
@@ -83,9 +62,9 @@ checksum = "a4a4ddaa51a5bc52a6948f74c06d20aaaddb71924eab79b8c97a8c556e942d6a"
|
||||
|
||||
[[package]]
|
||||
name = "bb8"
|
||||
version = "0.8.1"
|
||||
version = "0.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "98b4b0f25f18bcdc3ac72bdb486ed0acf7e185221fd4dc985bc15db5800b0ba2"
|
||||
checksum = "1627eccf3aa91405435ba240be23513eeca466b5dc33866422672264de061582"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"futures-channel",
|
||||
@@ -233,12 +212,6 @@ dependencies = [
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "data-encoding"
|
||||
version = "2.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3ee2393c4a91429dffb4bedf19f4d6abf27d8a732c8ce4980305d782e5426d57"
|
||||
|
||||
[[package]]
|
||||
name = "digest"
|
||||
version = "0.10.6"
|
||||
@@ -250,24 +223,6 @@ dependencies = [
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "either"
|
||||
version = "1.8.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7fcaabb2fef8c910e7f4c7ce9f67a1283a1715879a7c230ca9d6d1ae31f16d91"
|
||||
|
||||
[[package]]
|
||||
name = "enum-as-inner"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c9720bba047d567ffc8a3cba48bf19126600e249ab7f128e9233e6376976a116"
|
||||
dependencies = [
|
||||
"heck",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "env_logger"
|
||||
version = "0.10.0"
|
||||
@@ -320,15 +275,6 @@ version = "1.0.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1"
|
||||
|
||||
[[package]]
|
||||
name = "form_urlencoded"
|
||||
version = "1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a9c384f161156f5260c24a097c56119f9be8c798586aecc13afbcbe7b7e26bf8"
|
||||
dependencies = [
|
||||
"percent-encoding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures"
|
||||
version = "0.3.28"
|
||||
@@ -464,12 +410,6 @@ version = "0.12.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
|
||||
|
||||
[[package]]
|
||||
name = "heck"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2540771e65fc8cb83cd6e8a237f70c319bd5c29f78ed1084ba5d50eeac86f7f9"
|
||||
|
||||
[[package]]
|
||||
name = "hermit-abi"
|
||||
version = "0.2.6"
|
||||
@@ -494,17 +434,6 @@ dependencies = [
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hostname"
|
||||
version = "0.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3c731c3e10504cc8ed35cfe2f1db4c9274c3d35fa486e3b31df46f068ef3e867"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"match_cfg",
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "http"
|
||||
version = "0.2.9"
|
||||
@@ -593,27 +522,6 @@ dependencies = [
|
||||
"cxx-build",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "idna"
|
||||
version = "0.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "418a0a6fab821475f634efe3ccc45c013f742efe03d853e8d3355d5cb850ecf8"
|
||||
dependencies = [
|
||||
"matches",
|
||||
"unicode-bidi",
|
||||
"unicode-normalization",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "idna"
|
||||
version = "0.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e14ddfc70884202db2244c223200c204c2bda1bc6e0998d11b5e024d657209e6"
|
||||
dependencies = [
|
||||
"unicode-bidi",
|
||||
"unicode-normalization",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "indexmap"
|
||||
version = "1.9.2"
|
||||
@@ -634,24 +542,6 @@ dependencies = [
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ipconfig"
|
||||
version = "0.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bd302af1b90f2463a98fa5ad469fc212c8e3175a41c3068601bfa2727591c5be"
|
||||
dependencies = [
|
||||
"socket2",
|
||||
"widestring",
|
||||
"winapi",
|
||||
"winreg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ipnet"
|
||||
version = "2.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f88c5561171189e69df9d98bcf18fd5f9558300f7ea7b801eb8a0fd748bd8745"
|
||||
|
||||
[[package]]
|
||||
name = "is-terminal"
|
||||
version = "0.4.4"
|
||||
@@ -664,15 +554,6 @@ dependencies = [
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itertools"
|
||||
version = "0.10.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473"
|
||||
dependencies = [
|
||||
"either",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itoa"
|
||||
version = "1.0.5"
|
||||
@@ -708,12 +589,6 @@ dependencies = [
|
||||
"wasm-bindgen",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lazy_static"
|
||||
version = "1.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646"
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.139"
|
||||
@@ -729,12 +604,6 @@ dependencies = [
|
||||
"cc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "linked-hash-map"
|
||||
version = "0.5.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0717cef1bc8b636c6e1c1bbdefc09e6322da8a9321966e8928ef80d20f7f770f"
|
||||
|
||||
[[package]]
|
||||
name = "linux-raw-sys"
|
||||
version = "0.1.4"
|
||||
@@ -753,31 +622,13 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "log"
|
||||
version = "0.4.19"
|
||||
version = "0.4.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b06a4cde4c0f271a446782e3eff8de789548ce57dbc8eca9292c27f4a42004b4"
|
||||
|
||||
[[package]]
|
||||
name = "lru-cache"
|
||||
version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "31e24f1ad8321ca0e8a1e0ac13f23cb668e6f5466c2c57319f6a5cf1cc8e3b1c"
|
||||
checksum = "abb12e687cfb44aa40f41fc3978ef76448f9b6038cad6aef4259d3c095a2382e"
|
||||
dependencies = [
|
||||
"linked-hash-map",
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "match_cfg"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ffbee8634e0d45d258acb448e7eaab3fce7a0a467395d4d9f228e3c1f01fb2e4"
|
||||
|
||||
[[package]]
|
||||
name = "matches"
|
||||
version = "0.1.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a3e378b66a060d48947b590737b30a1be76706c8dd7b8ba0f2fe3989c68a853f"
|
||||
|
||||
[[package]]
|
||||
name = "md-5"
|
||||
version = "0.10.5"
|
||||
@@ -886,15 +737,9 @@ dependencies = [
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "percent-encoding"
|
||||
version = "2.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "478c572c3d73181ff3c2539045f6eb99e5491218eae919370993b890cdbdd98e"
|
||||
|
||||
[[package]]
|
||||
name = "pgcat"
|
||||
version = "1.0.2-alpha3"
|
||||
version = "1.0.1"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"async-trait",
|
||||
@@ -909,7 +754,6 @@ dependencies = [
|
||||
"futures",
|
||||
"hmac",
|
||||
"hyper",
|
||||
"itertools",
|
||||
"jemallocator",
|
||||
"log",
|
||||
"md-5",
|
||||
@@ -926,7 +770,6 @@ dependencies = [
|
||||
"rustls-pemfile",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
"serde_json",
|
||||
"sha-1",
|
||||
"sha2",
|
||||
"socket2",
|
||||
@@ -934,9 +777,7 @@ dependencies = [
|
||||
"stringprep",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-test",
|
||||
"toml",
|
||||
"trust-dns-resolver",
|
||||
"webpki-roots",
|
||||
]
|
||||
|
||||
@@ -984,22 +825,22 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "pin-project"
|
||||
version = "1.1.0"
|
||||
version = "1.0.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c95a7476719eab1e366eaf73d0260af3021184f18177925b07f54b30089ceead"
|
||||
checksum = "ad29a609b6bcd67fee905812e544992d216af9d755757c05ed2d0e15a74c6ecc"
|
||||
dependencies = [
|
||||
"pin-project-internal",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pin-project-internal"
|
||||
version = "1.1.0"
|
||||
version = "1.0.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "39407670928234ebc5e6e580247dd567ad73a3578460c5990f9503df207e8f07"
|
||||
checksum = "069bdb1e05adc7a8990dce9cc75370895fbe4e3d58b9b73bf1aee56359344a55"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.9",
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1047,12 +888,6 @@ dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "quick-error"
|
||||
version = "1.2.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0"
|
||||
|
||||
[[package]]
|
||||
name = "quote"
|
||||
version = "1.0.26"
|
||||
@@ -1103,9 +938,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "regex"
|
||||
version = "1.8.4"
|
||||
version = "1.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d0ab3ca65655bb1e41f2a8c8cd662eb4fb035e67c3f78da1d61dffe89d07300f"
|
||||
checksum = "ac6cf59af1067a3fb53fbe5c88c053764e930f932be1d71d3ffe032cbe147f59"
|
||||
dependencies = [
|
||||
"aho-corasick",
|
||||
"memchr",
|
||||
@@ -1114,19 +949,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "regex-syntax"
|
||||
version = "0.7.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "436b050e76ed2903236f032a59761c1eb99e1b0aead2c257922771dab1fc8c78"
|
||||
|
||||
[[package]]
|
||||
name = "resolv-conf"
|
||||
version = "0.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "52e44394d2086d010551b14b53b1f24e31647570cd1deb0379e2c21b329aba00"
|
||||
dependencies = [
|
||||
"hostname",
|
||||
"quick-error",
|
||||
]
|
||||
checksum = "b6868896879ba532248f33598de5181522d8b3d9d724dfd230911e1a7d4822f5"
|
||||
|
||||
[[package]]
|
||||
name = "ring"
|
||||
@@ -1159,9 +984,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "rustls"
|
||||
version = "0.21.1"
|
||||
version = "0.21.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c911ba11bc8433e811ce56fde130ccf32f5127cab0e0194e9c68c5a5b671791e"
|
||||
checksum = "07180898a28ed6a7f7ba2311594308f595e3dd2e3c3812fa0a80a47b45f17e5d"
|
||||
dependencies = [
|
||||
"log",
|
||||
"ring",
|
||||
@@ -1188,12 +1013,6 @@ dependencies = [
|
||||
"untrusted",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ryu"
|
||||
version = "1.0.13"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f91339c0467de62360649f8d3e185ca8de4224ff281f66000de5eb2a77a79041"
|
||||
|
||||
[[package]]
|
||||
name = "scopeguard"
|
||||
version = "1.1.0"
|
||||
@@ -1221,9 +1040,6 @@ name = "serde"
|
||||
version = "1.0.160"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bb2f3770c8bce3bcda7e149193a069a0f4365bda1fa5cd88e03bca26afc1216c"
|
||||
dependencies = [
|
||||
"serde_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_derive"
|
||||
@@ -1236,22 +1052,11 @@ dependencies = [
|
||||
"syn 2.0.9",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_json"
|
||||
version = "1.0.96"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "057d394a50403bcac12672b2b18fb387ab6d289d957dab67dd201875391e52f1"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"ryu",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_spanned"
|
||||
version = "0.6.2"
|
||||
version = "0.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "93107647184f6027e3b7dcb2e11034cf95ffa1e3a682c67951963ac69c1c007d"
|
||||
checksum = "0efd8caf556a6cebd3b285caf480045fcc1ac04f6bd786b09a6f11af30c4fcf4"
|
||||
dependencies = [
|
||||
"serde",
|
||||
]
|
||||
@@ -1326,23 +1131,11 @@ checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d"
|
||||
|
||||
[[package]]
|
||||
name = "sqlparser"
|
||||
version = "0.34.0"
|
||||
version = "0.33.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "37d3706eefb17039056234df6b566b0014f303f867f2656108334a55b8096f59"
|
||||
checksum = "355dc4d4b6207ca8a3434fc587db0a8016130a574dbcdbfb93d7f7b5bc5b211a"
|
||||
dependencies = [
|
||||
"log",
|
||||
"sqlparser_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sqlparser_derive"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "55fe75cb4a364c7f7ae06c7dbbc8d84bddd85d6cdf9975963c3935bc1991761e"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1398,26 +1191,6 @@ dependencies = [
|
||||
"winapi-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "1.0.37"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "10deb33631e3c9018b9baf9dcbbc4f737320d2b576bac10f6aefa048fa407e3e"
|
||||
dependencies = [
|
||||
"thiserror-impl",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror-impl"
|
||||
version = "1.0.37"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "982d17546b47146b28f7c22e3d08465f6b8903d0ea13c1660d9d84a6e7adcdbb"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "time"
|
||||
version = "0.1.45"
|
||||
@@ -1485,30 +1258,6 @@ dependencies = [
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-stream"
|
||||
version = "0.1.11"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d660770404473ccd7bc9f8b28494a811bc18542b915c0855c51e8f419d5223ce"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-test"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "53474327ae5e166530d17f2d956afcb4f8a004de581b3cae10f12006bc8163e3"
|
||||
dependencies = [
|
||||
"async-stream",
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-util"
|
||||
version = "0.7.7"
|
||||
@@ -1525,9 +1274,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "toml"
|
||||
version = "0.7.4"
|
||||
version = "0.7.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d6135d499e69981f9ff0ef2167955a5333c35e36f6937d382974566b3d5b94ec"
|
||||
checksum = "b403acf6f2bb0859c93c7f0d967cb4a75a7ac552100f9322faf64dc047669b21"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_spanned",
|
||||
@@ -1537,18 +1286,18 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "toml_datetime"
|
||||
version = "0.6.2"
|
||||
version = "0.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5a76a9312f5ba4c2dec6b9161fdf25d87ad8a09256ccea5a556fef03c706a10f"
|
||||
checksum = "3ab8ed2edee10b50132aed5f331333428b011c99402b5a534154ed15746f9622"
|
||||
dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_edit"
|
||||
version = "0.19.9"
|
||||
version = "0.19.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "92d964908cec0d030b812013af25a0e57fddfadb1e066ecc6681d86253129d4f"
|
||||
checksum = "08de71aa0d6e348f070457f85af8bd566e2bc452156a423ddf22861b3a953fae"
|
||||
dependencies = [
|
||||
"indexmap",
|
||||
"serde",
|
||||
@@ -1571,21 +1320,9 @@ checksum = "8ce8c33a8d48bd45d624a6e523445fd21ec13d3653cd51f681abf67418f54eb8"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"pin-project-lite",
|
||||
"tracing-attributes",
|
||||
"tracing-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tracing-attributes"
|
||||
version = "0.1.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4017f8f45139870ca7e672686113917c71c7a6e02d4924eda67186083c03081a"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 1.0.109",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tracing-core"
|
||||
version = "0.1.30"
|
||||
@@ -1595,51 +1332,6 @@ dependencies = [
|
||||
"once_cell",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "trust-dns-proto"
|
||||
version = "0.22.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4f7f83d1e4a0e4358ac54c5c3681e5d7da5efc5a7a632c90bb6d6669ddd9bc26"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"cfg-if",
|
||||
"data-encoding",
|
||||
"enum-as-inner",
|
||||
"futures-channel",
|
||||
"futures-io",
|
||||
"futures-util",
|
||||
"idna 0.2.3",
|
||||
"ipnet",
|
||||
"lazy_static",
|
||||
"rand",
|
||||
"smallvec",
|
||||
"thiserror",
|
||||
"tinyvec",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "trust-dns-resolver"
|
||||
version = "0.22.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "aff21aa4dcefb0a1afbfac26deb0adc93888c7d295fb63ab273ef276ba2b7cfe"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"futures-util",
|
||||
"ipconfig",
|
||||
"lazy_static",
|
||||
"lru-cache",
|
||||
"parking_lot",
|
||||
"resolv-conf",
|
||||
"smallvec",
|
||||
"thiserror",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"trust-dns-proto",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "try-lock"
|
||||
version = "0.2.4"
|
||||
@@ -1685,17 +1377,6 @@ version = "0.7.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a"
|
||||
|
||||
[[package]]
|
||||
name = "url"
|
||||
version = "2.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0d68c799ae75762b8c3fe375feb6600ef5602c883c5d21eb51c09f22b83c4643"
|
||||
dependencies = [
|
||||
"form_urlencoded",
|
||||
"idna 0.3.0",
|
||||
"percent-encoding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "version_check"
|
||||
version = "0.9.4"
|
||||
@@ -1797,12 +1478,6 @@ dependencies = [
|
||||
"rustls-webpki",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "widestring"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "17882f045410753661207383517a6f62ec3dbeb6a4ed2acce01f0728238d1983"
|
||||
|
||||
[[package]]
|
||||
name = "winapi"
|
||||
version = "0.3.9"
|
||||
@@ -1902,18 +1577,9 @@ checksum = "447660ad36a13288b1db4d4248e857b510e8c3a225c822ba4fb748c0aafecffd"
|
||||
|
||||
[[package]]
|
||||
name = "winnow"
|
||||
version = "0.4.6"
|
||||
version = "0.3.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "61de7bac303dc551fe038e2b3cef0f571087a47571ea6e79a87692ac99b99699"
|
||||
checksum = "faf09497b8f8b5ac5d3bb4d05c0a99be20f26fd3d5f2db7b0716e946d5103658"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "winreg"
|
||||
version = "0.10.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "80d0f4e272c85def139476380b12f9ac60926689dd2e01d4923222f40580869d"
|
||||
dependencies = [
|
||||
"winapi",
|
||||
]
|
||||
|
||||
13
Cargo.toml
13
Cargo.toml
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "pgcat"
|
||||
version = "1.0.2-alpha3"
|
||||
version = "1.0.1"
|
||||
edition = "2021"
|
||||
|
||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||
@@ -8,18 +8,18 @@ edition = "2021"
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
bytes = "1"
|
||||
md-5 = "0.10"
|
||||
bb8 = "0.8.1"
|
||||
bb8 = "0.8.0"
|
||||
async-trait = "0.1"
|
||||
rand = "0.8"
|
||||
chrono = "0.4"
|
||||
sha-1 = "0.10"
|
||||
toml = "0.7"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde = "1"
|
||||
serde_derive = "1"
|
||||
regex = "1"
|
||||
num_cpus = "1"
|
||||
once_cell = "1"
|
||||
sqlparser = {version = "0.34", features = ["visitor"] }
|
||||
sqlparser = "0.33.0"
|
||||
log = "0.4"
|
||||
arc-swap = "1"
|
||||
env_logger = "0.10"
|
||||
@@ -42,11 +42,6 @@ fallible-iterator = "0.2"
|
||||
pin-project = "1"
|
||||
webpki-roots = "0.23"
|
||||
rustls = { version = "0.21", features = ["dangerous_configuration"] }
|
||||
trust-dns-resolver = "0.22.0"
|
||||
tokio-test = "0.4.2"
|
||||
serde_json = "1"
|
||||
itertools = "0.10"
|
||||
|
||||
[target.'cfg(not(target_env = "msvc"))'.dependencies]
|
||||
jemallocator = "0.5.0"
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ PostgreSQL pooler and proxy (like PgBouncer) with support for sharding, load bal
|
||||
| Failover | **Stable** | Queries are automatically rerouted around broken replicas, validated by regular health checks. |
|
||||
| Admin database statistics | **Stable** | Pooler statistics and administration via the `pgbouncer` and `pgcat` databases. |
|
||||
| Prometheus statistics | **Stable** | Statistics are reported via a HTTP endpoint for Prometheus. |
|
||||
| SSL/TLS | **Stable** | Clients can connect to the pooler using TLS. Pooler can connect to Postgres servers using TLS. |
|
||||
| Client TLS | **Stable** | Clients can connect to the pooler using TLS/SSL. |
|
||||
| Client/Server authentication | **Stable** | Clients can connect using MD5 authentication, supported by `libpq` and all Postgres client drivers. PgCat can connect to Postgres using MD5 and SCRAM-SHA-256. |
|
||||
| Live configuration reloading | **Stable** | Identical to PgBouncer; all settings can be reloaded dynamically (except `host` and `port`). |
|
||||
| Auth passthrough | **Stable** | MD5 password authentication can be configured to use an `auth_query` so no cleartext passwords are needed in the config file.|
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
FROM rust:1.70-bullseye
|
||||
FROM rust:bullseye
|
||||
|
||||
# Dependencies
|
||||
RUN apt-get update -y \
|
||||
|
||||
@@ -25,7 +25,7 @@ x-common-env-pg:
|
||||
|
||||
services:
|
||||
main:
|
||||
image: gcr.io/google_containers/pause:3.2
|
||||
image: kubernetes/pause
|
||||
ports:
|
||||
- 6432
|
||||
|
||||
@@ -64,7 +64,7 @@ services:
|
||||
<<: *common-env-pg
|
||||
POSTGRES_INITDB_ARGS: --auth-local=md5 --auth-host=md5 --auth=md5
|
||||
PGPORT: 10432
|
||||
command: ["postgres", "-p", "10432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
|
||||
command: ["postgres", "-p", "5432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
|
||||
|
||||
toxiproxy:
|
||||
build: .
|
||||
|
||||
@@ -1,22 +0,0 @@
|
||||
# This is an example of the most basic config
|
||||
# that will mimic what PgBouncer does in transaction mode with one server.
|
||||
|
||||
[general]
|
||||
|
||||
host = "0.0.0.0"
|
||||
port = 6433
|
||||
admin_username = "pgcat"
|
||||
admin_password = "pgcat"
|
||||
|
||||
[pools.pgml.users.0]
|
||||
username = "postgres"
|
||||
password = "postgres"
|
||||
pool_size = 10
|
||||
min_pool_size = 1
|
||||
pool_mode = "transaction"
|
||||
|
||||
[pools.pgml.shards.0]
|
||||
servers = [
|
||||
["127.0.0.1", 28815, "primary"]
|
||||
]
|
||||
database = "postgres"
|
||||
110
pgcat.toml
110
pgcat.toml
@@ -60,9 +60,6 @@ tcp_keepalives_count = 5
|
||||
# Number of seconds between keepalive packets.
|
||||
tcp_keepalives_interval = 5
|
||||
|
||||
# Handle prepared statements.
|
||||
prepared_statements = true
|
||||
|
||||
# Path to TLS Certificate file to use for TLS connections
|
||||
# tls_certificate = ".circleci/server.cert"
|
||||
# Path to TLS private key file to use for TLS connections
|
||||
@@ -80,58 +77,6 @@ admin_username = "admin_user"
|
||||
# Password to access the virtual administrative database
|
||||
admin_password = "admin_pass"
|
||||
|
||||
# Default plugins that are configured on all pools.
|
||||
[plugins]
|
||||
|
||||
# Prewarmer plugin that runs queries on server startup, before giving the connection
|
||||
# to the client.
|
||||
[plugins.prewarmer]
|
||||
enabled = false
|
||||
queries = [
|
||||
"SELECT pg_prewarm('pgbench_accounts')",
|
||||
]
|
||||
|
||||
# Log all queries to stdout.
|
||||
[plugins.query_logger]
|
||||
enabled = false
|
||||
|
||||
# Block access to tables that Postgres does not allow us to control.
|
||||
[plugins.table_access]
|
||||
enabled = false
|
||||
tables = [
|
||||
"pg_user",
|
||||
"pg_roles",
|
||||
"pg_database",
|
||||
]
|
||||
|
||||
# Intercept user queries and give a fake reply.
|
||||
[plugins.intercept]
|
||||
enabled = true
|
||||
|
||||
[plugins.intercept.queries.0]
|
||||
|
||||
query = "select current_database() as a, current_schemas(false) as b"
|
||||
schema = [
|
||||
["a", "text"],
|
||||
["b", "text"],
|
||||
]
|
||||
result = [
|
||||
["${DATABASE}", "{public}"],
|
||||
]
|
||||
|
||||
[plugins.intercept.queries.1]
|
||||
|
||||
query = "select current_database(), current_schema(), current_user"
|
||||
schema = [
|
||||
["current_database", "text"],
|
||||
["current_schema", "text"],
|
||||
["current_user", "text"],
|
||||
]
|
||||
result = [
|
||||
["${DATABASE}", "public", "${USER}"],
|
||||
]
|
||||
|
||||
|
||||
# pool configs are structured as pool.<pool_name>
|
||||
# the pool_name is what clients use as database name when connecting.
|
||||
# For a pool named `sharded_db`, clients access that pool using connection string like
|
||||
@@ -201,61 +146,6 @@ idle_timeout = 40000
|
||||
# Connect timeout can be overwritten in the pool
|
||||
connect_timeout = 3000
|
||||
|
||||
# When enabled, ip resolutions for server connections specified using hostnames will be cached
|
||||
# and checked for changes every `dns_max_ttl` seconds. If a change in the host resolution is found
|
||||
# old ip connections are closed (gracefully) and new connections will start using new ip.
|
||||
# dns_cache_enabled = false
|
||||
|
||||
# Specifies how often (in seconds) cached ip addresses for servers are rechecked (see `dns_cache_enabled`).
|
||||
# dns_max_ttl = 30
|
||||
|
||||
# Plugins can be configured on a pool-per-pool basis. This overrides the global plugins setting,
|
||||
# so all plugins have to be configured here again.
|
||||
[pool.sharded_db.plugins]
|
||||
|
||||
[pools.sharded_db.plugins.prewarmer]
|
||||
enabled = true
|
||||
queries = [
|
||||
"SELECT pg_prewarm('pgbench_accounts')",
|
||||
]
|
||||
|
||||
[pools.sharded_db.plugins.query_logger]
|
||||
enabled = false
|
||||
|
||||
[pools.sharded_db.plugins.table_access]
|
||||
enabled = false
|
||||
tables = [
|
||||
"pg_user",
|
||||
"pg_roles",
|
||||
"pg_database",
|
||||
]
|
||||
|
||||
[pools.sharded_db.plugins.intercept]
|
||||
enabled = true
|
||||
|
||||
[pools.sharded_db.plugins.intercept.queries.0]
|
||||
|
||||
query = "select current_database() as a, current_schemas(false) as b"
|
||||
schema = [
|
||||
["a", "text"],
|
||||
["b", "text"],
|
||||
]
|
||||
result = [
|
||||
["${DATABASE}", "{public}"],
|
||||
]
|
||||
|
||||
[pools.sharded_db.plugins.intercept.queries.1]
|
||||
|
||||
query = "select current_database(), current_schema(), current_user"
|
||||
schema = [
|
||||
["current_database", "text"],
|
||||
["current_schema", "text"],
|
||||
["current_user", "text"],
|
||||
]
|
||||
result = [
|
||||
["${DATABASE}", "public", "${USER}"],
|
||||
]
|
||||
|
||||
# User configs are structured as pool.<pool_name>.users.<user_index>
|
||||
# This section holds the credentials for users that may connect to this cluster
|
||||
[pools.sharded_db.users.0]
|
||||
|
||||
74
src/admin.rs
74
src/admin.rs
@@ -1,5 +1,4 @@
|
||||
use crate::pool::BanReason;
|
||||
use crate::stats::pool::PoolStats;
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
use log::{error, info, trace};
|
||||
use nix::sys::signal::{self, Signal};
|
||||
@@ -13,9 +12,9 @@ use tokio::time::Instant;
|
||||
use crate::config::{get_config, reload_config, VERSION};
|
||||
use crate::errors::Error;
|
||||
use crate::messages::*;
|
||||
use crate::pool::ClientServerMap;
|
||||
use crate::pool::{get_all_pools, get_pool};
|
||||
use crate::stats::{get_client_stats, get_server_stats, ClientState, ServerState};
|
||||
use crate::stats::{get_client_stats, get_pool_stats, get_server_stats, ClientState, ServerState};
|
||||
use crate::ClientServerMap;
|
||||
|
||||
pub fn generate_server_info_for_admin() -> BytesMut {
|
||||
let mut server_info = BytesMut::new();
|
||||
@@ -255,12 +254,39 @@ async fn show_pools<T>(stream: &mut T) -> Result<(), Error>
|
||||
where
|
||||
T: tokio::io::AsyncWrite + std::marker::Unpin,
|
||||
{
|
||||
let pool_lookup = PoolStats::construct_pool_lookup();
|
||||
let all_pool_stats = get_pool_stats();
|
||||
|
||||
let columns = vec![
|
||||
("database", DataType::Text),
|
||||
("user", DataType::Text),
|
||||
("pool_mode", DataType::Text),
|
||||
("cl_idle", DataType::Numeric),
|
||||
("cl_active", DataType::Numeric),
|
||||
("cl_waiting", DataType::Numeric),
|
||||
("cl_cancel_req", DataType::Numeric),
|
||||
("sv_active", DataType::Numeric),
|
||||
("sv_idle", DataType::Numeric),
|
||||
("sv_used", DataType::Numeric),
|
||||
("sv_tested", DataType::Numeric),
|
||||
("sv_login", DataType::Numeric),
|
||||
("maxwait", DataType::Numeric),
|
||||
("maxwait_us", DataType::Numeric),
|
||||
];
|
||||
|
||||
let mut res = BytesMut::new();
|
||||
res.put(row_description(&PoolStats::generate_header()));
|
||||
pool_lookup.iter().for_each(|(_identifier, pool_stats)| {
|
||||
res.put(data_row(&pool_stats.generate_row()));
|
||||
});
|
||||
res.put(row_description(&columns));
|
||||
|
||||
for ((_user_pool, _pool), pool_stats) in all_pool_stats {
|
||||
let mut row = vec![
|
||||
pool_stats.database(),
|
||||
pool_stats.user(),
|
||||
pool_stats.pool_mode().to_string(),
|
||||
];
|
||||
pool_stats.populate_row(&mut row);
|
||||
pool_stats.clear_maxwait();
|
||||
res.put(data_row(&row));
|
||||
}
|
||||
|
||||
res.put(command_complete("SHOW"));
|
||||
|
||||
// ReadyForQuery
|
||||
@@ -308,17 +334,17 @@ where
|
||||
let paused = pool.paused();
|
||||
|
||||
res.put(data_row(&vec![
|
||||
address.name(), // name
|
||||
address.host.to_string(), // host
|
||||
address.port.to_string(), // port
|
||||
database_name.to_string(), // database
|
||||
pool_config.user.username.to_string(), // force_user
|
||||
pool_config.user.pool_size.to_string(), // pool_size
|
||||
pool_config.user.min_pool_size.unwrap_or(0).to_string(), // min_pool_size
|
||||
"0".to_string(), // reserve_pool
|
||||
pool_config.pool_mode.to_string(), // pool_mode
|
||||
pool_config.user.pool_size.to_string(), // max_connections
|
||||
pool_state.connections.to_string(), // current_connections
|
||||
address.name(), // name
|
||||
address.host.to_string(), // host
|
||||
address.port.to_string(), // port
|
||||
database_name.to_string(), // database
|
||||
pool_config.user.username.to_string(), // force_user
|
||||
pool_config.user.pool_size.to_string(), // pool_size
|
||||
"0".to_string(), // min_pool_size
|
||||
"0".to_string(), // reserve_pool
|
||||
pool_config.pool_mode.to_string(), // pool_mode
|
||||
pool_config.user.pool_size.to_string(), // max_connections
|
||||
pool_state.connections.to_string(), // current_connections
|
||||
match paused {
|
||||
// paused
|
||||
true => "1".to_string(),
|
||||
@@ -699,8 +725,6 @@ where
|
||||
("bytes_sent", DataType::Numeric),
|
||||
("bytes_received", DataType::Numeric),
|
||||
("age_seconds", DataType::Numeric),
|
||||
("prepare_cache_hit", DataType::Numeric),
|
||||
("prepare_cache_miss", DataType::Numeric),
|
||||
];
|
||||
|
||||
let new_map = get_server_stats();
|
||||
@@ -724,14 +748,6 @@ where
|
||||
.duration_since(server.connect_time())
|
||||
.as_secs()
|
||||
.to_string(),
|
||||
server
|
||||
.prepared_hit_count
|
||||
.load(Ordering::Relaxed)
|
||||
.to_string(),
|
||||
server
|
||||
.prepared_miss_count
|
||||
.load(Ordering::Relaxed)
|
||||
.to_string(),
|
||||
];
|
||||
|
||||
res.put(data_row(&row));
|
||||
|
||||
324
src/client.rs
324
src/client.rs
@@ -3,9 +3,8 @@ use crate::pool::BanReason;
|
||||
/// Handle clients by pretending to be a PostgreSQL server.
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
use log::{debug, error, info, trace, warn};
|
||||
use once_cell::sync::Lazy;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{atomic::AtomicUsize, Arc};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use tokio::io::{split, AsyncReadExt, BufReader, ReadHalf, WriteHalf};
|
||||
use tokio::net::TcpStream;
|
||||
@@ -14,25 +13,17 @@ use tokio::sync::mpsc::Sender;
|
||||
|
||||
use crate::admin::{generate_server_info_for_admin, handle_admin};
|
||||
use crate::auth_passthrough::refetch_auth_hash;
|
||||
use crate::config::{
|
||||
get_config, get_idle_client_in_transaction_timeout, get_prepared_statements, Address, PoolMode,
|
||||
};
|
||||
use crate::config::{get_config, get_idle_client_in_transaction_timeout, Address, PoolMode};
|
||||
use crate::constants::*;
|
||||
use crate::messages::*;
|
||||
use crate::plugins::PluginOutput;
|
||||
use crate::pool::{get_pool, ClientServerMap, ConnectionPool};
|
||||
use crate::query_router::{Command, QueryRouter};
|
||||
use crate::server::Server;
|
||||
use crate::stats::{ClientStats, ServerStats};
|
||||
use crate::stats::{ClientStats, PoolStats, ServerStats};
|
||||
use crate::tls::Tls;
|
||||
|
||||
use tokio_rustls::server::TlsStream;
|
||||
|
||||
/// Incrementally count prepared statements
|
||||
/// to avoid random conflicts in places where the random number generator is weak.
|
||||
pub static PREPARED_STATEMENT_COUNTER: Lazy<Arc<AtomicUsize>> =
|
||||
Lazy::new(|| Arc::new(AtomicUsize::new(0)));
|
||||
|
||||
/// Type of connection received from client.
|
||||
enum ClientConnectionType {
|
||||
Startup,
|
||||
@@ -101,9 +92,6 @@ pub struct Client<S, T> {
|
||||
|
||||
/// Used to notify clients about an impending shutdown
|
||||
shutdown: Receiver<()>,
|
||||
|
||||
/// Prepared statements
|
||||
prepared_statements: HashMap<String, Parse>,
|
||||
}
|
||||
|
||||
/// Client entrypoint.
|
||||
@@ -665,12 +653,24 @@ where
|
||||
ready_for_query(&mut write).await?;
|
||||
|
||||
trace!("Startup OK");
|
||||
let pool_stats = match get_pool(pool_name, username) {
|
||||
Some(pool) => {
|
||||
if !admin {
|
||||
pool.stats
|
||||
} else {
|
||||
Arc::new(PoolStats::default())
|
||||
}
|
||||
}
|
||||
None => Arc::new(PoolStats::default()),
|
||||
};
|
||||
|
||||
let stats = Arc::new(ClientStats::new(
|
||||
process_id,
|
||||
application_name,
|
||||
username,
|
||||
pool_name,
|
||||
tokio::time::Instant::now(),
|
||||
pool_stats,
|
||||
));
|
||||
|
||||
Ok(Client {
|
||||
@@ -693,7 +693,6 @@ where
|
||||
application_name: application_name.to_string(),
|
||||
shutdown,
|
||||
connected_to_server: false,
|
||||
prepared_statements: HashMap::new(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -728,7 +727,6 @@ where
|
||||
application_name: String::from("undefined"),
|
||||
shutdown,
|
||||
connected_to_server: false,
|
||||
prepared_statements: HashMap::new(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -767,13 +765,6 @@ where
|
||||
|
||||
self.stats.register(self.stats.clone());
|
||||
|
||||
// Result returned by one of the plugins.
|
||||
let mut plugin_output = None;
|
||||
|
||||
// Prepared statement being executed
|
||||
let mut prepared_statement = None;
|
||||
let mut will_prepare = false;
|
||||
|
||||
// Our custom protocol loop.
|
||||
// We expect the client to either start a transaction with regular queries
|
||||
// or issue commands for our sharding and server selection protocol.
|
||||
@@ -783,16 +774,13 @@ where
|
||||
self.transaction_mode
|
||||
);
|
||||
|
||||
// Should we rewrite prepared statements and bind messages?
|
||||
let mut prepared_statements_enabled = get_prepared_statements();
|
||||
|
||||
// Read a complete message from the client, which normally would be
|
||||
// either a `Q` (query) or `P` (prepare, extended protocol).
|
||||
// We can parse it here before grabbing a server from the pool,
|
||||
// in case the client is sending some custom protocol messages, e.g.
|
||||
// SET SHARDING KEY TO 'bigint';
|
||||
|
||||
let mut message = tokio::select! {
|
||||
let message = tokio::select! {
|
||||
_ = self.shutdown.recv() => {
|
||||
if !self.admin {
|
||||
error_response_terminal(
|
||||
@@ -820,75 +808,28 @@ where
|
||||
// allocate a connection, we wouldn't be able to send back an error message
|
||||
// to the client so we buffer them and defer the decision to error out or not
|
||||
// to when we get the S message
|
||||
'D' => {
|
||||
if prepared_statements_enabled {
|
||||
let name;
|
||||
(name, message) = self.rewrite_describe(message).await?;
|
||||
|
||||
if let Some(name) = name {
|
||||
prepared_statement = Some(name);
|
||||
}
|
||||
}
|
||||
|
||||
self.buffer.put(&message[..]);
|
||||
continue;
|
||||
}
|
||||
|
||||
'E' => {
|
||||
'D' | 'E' => {
|
||||
self.buffer.put(&message[..]);
|
||||
continue;
|
||||
}
|
||||
|
||||
'Q' => {
|
||||
if query_router.query_parser_enabled() {
|
||||
if let Ok(ast) = QueryRouter::parse(&message) {
|
||||
let plugin_result = query_router.execute_plugins(&ast).await;
|
||||
|
||||
match plugin_result {
|
||||
Ok(PluginOutput::Deny(error)) => {
|
||||
error_response(&mut self.write, &error).await?;
|
||||
continue;
|
||||
}
|
||||
|
||||
Ok(PluginOutput::Intercept(result)) => {
|
||||
write_all(&mut self.write, result).await?;
|
||||
continue;
|
||||
}
|
||||
|
||||
_ => (),
|
||||
};
|
||||
|
||||
let _ = query_router.infer(&ast);
|
||||
}
|
||||
query_router.infer(&message);
|
||||
}
|
||||
}
|
||||
|
||||
'P' => {
|
||||
if prepared_statements_enabled {
|
||||
(prepared_statement, message) = self.rewrite_parse(message)?;
|
||||
will_prepare = true;
|
||||
}
|
||||
|
||||
self.buffer.put(&message[..]);
|
||||
|
||||
if query_router.query_parser_enabled() {
|
||||
if let Ok(ast) = QueryRouter::parse(&message) {
|
||||
if let Ok(output) = query_router.execute_plugins(&ast).await {
|
||||
plugin_output = Some(output);
|
||||
}
|
||||
|
||||
let _ = query_router.infer(&ast);
|
||||
}
|
||||
query_router.infer(&message);
|
||||
}
|
||||
|
||||
continue;
|
||||
}
|
||||
|
||||
'B' => {
|
||||
if prepared_statements_enabled {
|
||||
(prepared_statement, message) = self.rewrite_bind(message).await?;
|
||||
}
|
||||
|
||||
self.buffer.put(&message[..]);
|
||||
|
||||
if query_router.query_parser_enabled() {
|
||||
@@ -916,18 +857,6 @@ where
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check on plugin results.
|
||||
match plugin_output {
|
||||
Some(PluginOutput::Deny(error)) => {
|
||||
self.buffer.clear();
|
||||
error_response(&mut self.write, &error).await?;
|
||||
plugin_output = None;
|
||||
continue;
|
||||
}
|
||||
|
||||
_ => (),
|
||||
};
|
||||
|
||||
// Get a pool instance referenced by the most up-to-date
|
||||
// pointer. This ensures we always read the latest config
|
||||
// when starting a query.
|
||||
@@ -1097,48 +1026,7 @@ where
|
||||
// If the client is in session mode, no more custom protocol
|
||||
// commands will be accepted.
|
||||
loop {
|
||||
// Only check if we should rewrite prepared statements
|
||||
// in session mode. In transaction mode, we check at the beginning of
|
||||
// each transaction.
|
||||
if !self.transaction_mode {
|
||||
prepared_statements_enabled = get_prepared_statements();
|
||||
}
|
||||
|
||||
debug!("Prepared statement active: {:?}", prepared_statement);
|
||||
|
||||
// We are processing a prepared statement.
|
||||
if let Some(ref name) = prepared_statement {
|
||||
debug!("Checking prepared statement is on server");
|
||||
// Get the prepared statement the server expects to see.
|
||||
let statement = match self.prepared_statements.get(name) {
|
||||
Some(statement) => {
|
||||
debug!("Prepared statement `{}` found in cache", name);
|
||||
statement
|
||||
}
|
||||
None => {
|
||||
return Err(Error::ClientError(format!(
|
||||
"prepared statement `{}` not found",
|
||||
name
|
||||
)))
|
||||
}
|
||||
};
|
||||
|
||||
// Since it's already in the buffer, we don't need to prepare it on this server.
|
||||
if will_prepare {
|
||||
server.will_prepare(&statement.name);
|
||||
will_prepare = false;
|
||||
} else {
|
||||
// The statement is not prepared on the server, so we need to prepare it.
|
||||
if server.should_prepare(&statement.name) {
|
||||
server.prepare(statement).await?;
|
||||
}
|
||||
}
|
||||
|
||||
// Done processing the prepared statement.
|
||||
prepared_statement = None;
|
||||
}
|
||||
|
||||
let mut message = match initial_message {
|
||||
let message = match initial_message {
|
||||
None => {
|
||||
trace!("Waiting for message inside transaction or in session mode");
|
||||
|
||||
@@ -1197,27 +1085,6 @@ where
|
||||
match code {
|
||||
// Query
|
||||
'Q' => {
|
||||
if query_router.query_parser_enabled() {
|
||||
if let Ok(ast) = QueryRouter::parse(&message) {
|
||||
let plugin_result = query_router.execute_plugins(&ast).await;
|
||||
|
||||
match plugin_result {
|
||||
Ok(PluginOutput::Deny(error)) => {
|
||||
error_response(&mut self.write, &error).await?;
|
||||
continue;
|
||||
}
|
||||
|
||||
Ok(PluginOutput::Intercept(result)) => {
|
||||
write_all(&mut self.write, result).await?;
|
||||
continue;
|
||||
}
|
||||
|
||||
_ => (),
|
||||
};
|
||||
|
||||
let _ = query_router.infer(&ast);
|
||||
}
|
||||
}
|
||||
debug!("Sending query to server");
|
||||
|
||||
self.send_and_receive_loop(
|
||||
@@ -1257,44 +1124,18 @@ where
|
||||
// Parse
|
||||
// The query with placeholders is here, e.g. `SELECT * FROM users WHERE email = $1 AND active = $2`.
|
||||
'P' => {
|
||||
if prepared_statements_enabled {
|
||||
(prepared_statement, message) = self.rewrite_parse(message)?;
|
||||
will_prepare = true;
|
||||
}
|
||||
|
||||
if query_router.query_parser_enabled() {
|
||||
if let Ok(ast) = QueryRouter::parse(&message) {
|
||||
if let Ok(output) = query_router.execute_plugins(&ast).await {
|
||||
plugin_output = Some(output);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
self.buffer.put(&message[..]);
|
||||
}
|
||||
|
||||
// Bind
|
||||
// The placeholder's replacements are here, e.g. 'user@email.com' and 'true'
|
||||
'B' => {
|
||||
if prepared_statements_enabled {
|
||||
(prepared_statement, message) = self.rewrite_bind(message).await?;
|
||||
}
|
||||
|
||||
self.buffer.put(&message[..]);
|
||||
}
|
||||
|
||||
// Describe
|
||||
// Command a client can issue to describe a previously prepared named statement.
|
||||
'D' => {
|
||||
if prepared_statements_enabled {
|
||||
let name;
|
||||
(name, message) = self.rewrite_describe(message).await?;
|
||||
|
||||
if let Some(name) = name {
|
||||
prepared_statement = Some(name);
|
||||
}
|
||||
}
|
||||
|
||||
self.buffer.put(&message[..]);
|
||||
}
|
||||
|
||||
@@ -1314,30 +1155,12 @@ where
|
||||
'S' => {
|
||||
debug!("Sending query to server");
|
||||
|
||||
match plugin_output {
|
||||
Some(PluginOutput::Deny(error)) => {
|
||||
error_response(&mut self.write, &error).await?;
|
||||
plugin_output = None;
|
||||
self.buffer.clear();
|
||||
continue;
|
||||
}
|
||||
|
||||
Some(PluginOutput::Intercept(result)) => {
|
||||
write_all(&mut self.write, result).await?;
|
||||
plugin_output = None;
|
||||
self.buffer.clear();
|
||||
continue;
|
||||
}
|
||||
|
||||
_ => (),
|
||||
};
|
||||
|
||||
self.buffer.put(&message[..]);
|
||||
|
||||
let first_message_code = (*self.buffer.get(0).unwrap_or(&0)) as char;
|
||||
|
||||
// Almost certainly true
|
||||
if first_message_code == 'P' && !prepared_statements_enabled {
|
||||
if first_message_code == 'P' {
|
||||
// Message layout
|
||||
// P followed by 32 int followed by null-terminated statement name
|
||||
// So message code should be in offset 0 of the buffer, first character
|
||||
@@ -1403,7 +1226,7 @@ where
|
||||
.receive_server_message(server, &address, &pool, &self.stats.clone())
|
||||
.await?;
|
||||
|
||||
match write_all_flush(&mut self.write, &response).await {
|
||||
match write_all_half(&mut self.write, &response).await {
|
||||
Ok(_) => (),
|
||||
Err(err) => {
|
||||
server.mark_bad();
|
||||
@@ -1465,107 +1288,6 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
/// Rewrite Parse (F) message to set the prepared statement name to one we control.
|
||||
/// Save it into the client cache.
|
||||
fn rewrite_parse(&mut self, message: BytesMut) -> Result<(Option<String>, BytesMut), Error> {
|
||||
let parse: Parse = (&message).try_into()?;
|
||||
|
||||
let name = parse.name.clone();
|
||||
|
||||
// Don't rewrite anonymous prepared statements
|
||||
if parse.anonymous() {
|
||||
debug!("Anonymous prepared statement");
|
||||
return Ok((None, message));
|
||||
}
|
||||
|
||||
let parse = parse.rename();
|
||||
|
||||
debug!(
|
||||
"Renamed prepared statement `{}` to `{}` and saved to cache",
|
||||
name, parse.name
|
||||
);
|
||||
|
||||
self.prepared_statements.insert(name.clone(), parse.clone());
|
||||
|
||||
Ok((Some(name), parse.try_into()?))
|
||||
}
|
||||
|
||||
/// Rewrite the Bind (F) message to use the prepared statement name
|
||||
/// saved in the client cache.
|
||||
async fn rewrite_bind(
|
||||
&mut self,
|
||||
message: BytesMut,
|
||||
) -> Result<(Option<String>, BytesMut), Error> {
|
||||
let bind: Bind = (&message).try_into()?;
|
||||
let name = bind.prepared_statement.clone();
|
||||
|
||||
if bind.anonymous() {
|
||||
debug!("Anonymous bind message");
|
||||
return Ok((None, message));
|
||||
}
|
||||
|
||||
match self.prepared_statements.get(&name) {
|
||||
Some(prepared_stmt) => {
|
||||
let bind = bind.reassign(prepared_stmt);
|
||||
|
||||
debug!("Rewrote bind `{}` to `{}`", name, bind.prepared_statement);
|
||||
|
||||
Ok((Some(name), bind.try_into()?))
|
||||
}
|
||||
None => {
|
||||
debug!("Got bind for unknown prepared statement {:?}", bind);
|
||||
|
||||
error_response(
|
||||
&mut self.write,
|
||||
&format!(
|
||||
"prepared statement \"{}\" does not exist",
|
||||
bind.prepared_statement
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Err(Error::ClientError(format!(
|
||||
"Prepared statement `{}` doesn't exist",
|
||||
name
|
||||
)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Rewrite the Describe (F) message to use the prepared statement name
|
||||
/// saved in the client cache.
|
||||
async fn rewrite_describe(
|
||||
&mut self,
|
||||
message: BytesMut,
|
||||
) -> Result<(Option<String>, BytesMut), Error> {
|
||||
let describe: Describe = (&message).try_into()?;
|
||||
let name = describe.statement_name.clone();
|
||||
|
||||
if describe.anonymous() {
|
||||
debug!("Anonymous describe");
|
||||
return Ok((None, message));
|
||||
}
|
||||
|
||||
match self.prepared_statements.get(&name) {
|
||||
Some(prepared_stmt) => {
|
||||
let describe = describe.rename(&prepared_stmt.name);
|
||||
|
||||
debug!(
|
||||
"Rewrote describe `{}` to `{}`",
|
||||
name, describe.statement_name
|
||||
);
|
||||
|
||||
Ok((Some(name), describe.try_into()?))
|
||||
}
|
||||
|
||||
None => {
|
||||
debug!("Got describe for unknown prepared statement {:?}", describe);
|
||||
|
||||
Ok((None, message))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Release the server from the client: it can't cancel its queries anymore.
|
||||
pub fn release(&self) {
|
||||
let mut guard = self.client_server_map.lock();
|
||||
@@ -1599,7 +1321,7 @@ where
|
||||
.receive_server_message(server, address, pool, client_stats)
|
||||
.await?;
|
||||
|
||||
match write_all_flush(&mut self.write, &response).await {
|
||||
match write_all_half(&mut self.write, &response).await {
|
||||
Ok(_) => (),
|
||||
Err(err) => {
|
||||
server.mark_bad();
|
||||
|
||||
196
src/config.rs
196
src/config.rs
@@ -12,7 +12,6 @@ use std::sync::Arc;
|
||||
use tokio::fs::File;
|
||||
use tokio::io::AsyncReadExt;
|
||||
|
||||
use crate::dns_cache::CachedResolver;
|
||||
use crate::errors::Error;
|
||||
use crate::pool::{ClientServerMap, ConnectionPool};
|
||||
use crate::sharding::ShardingFunction;
|
||||
@@ -122,16 +121,6 @@ impl Default for Address {
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for Address {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"[address: {}:{}][database: {}][user: {}]",
|
||||
self.host, self.port, self.database, self.username
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// We need to implement PartialEq by ourselves so we skip stats in the comparison
|
||||
impl PartialEq for Address {
|
||||
fn eq(&self, other: &Self) -> bool {
|
||||
@@ -245,8 +234,6 @@ pub struct General {
|
||||
pub port: u16,
|
||||
|
||||
pub enable_prometheus_exporter: Option<bool>,
|
||||
|
||||
#[serde(default = "General::default_prometheus_exporter_port")]
|
||||
pub prometheus_exporter_port: i16,
|
||||
|
||||
#[serde(default = "General::default_connect_timeout")]
|
||||
@@ -268,12 +255,6 @@ pub struct General {
|
||||
#[serde(default)] // False
|
||||
pub log_client_disconnections: bool,
|
||||
|
||||
#[serde(default)] // False
|
||||
pub dns_cache_enabled: bool,
|
||||
|
||||
#[serde(default = "General::default_dns_max_ttl")]
|
||||
pub dns_max_ttl: u64,
|
||||
|
||||
#[serde(default = "General::default_shutdown_timeout")]
|
||||
pub shutdown_timeout: u64,
|
||||
|
||||
@@ -292,9 +273,6 @@ pub struct General {
|
||||
#[serde(default = "General::default_server_lifetime")]
|
||||
pub server_lifetime: u64,
|
||||
|
||||
#[serde(default = "General::default_server_round_robin")] // False
|
||||
pub server_round_robin: bool,
|
||||
|
||||
#[serde(default = "General::default_worker_threads")]
|
||||
pub worker_threads: usize,
|
||||
|
||||
@@ -313,16 +291,9 @@ pub struct General {
|
||||
pub admin_username: String,
|
||||
pub admin_password: String,
|
||||
|
||||
#[serde(default = "General::default_validate_config")]
|
||||
pub validate_config: bool,
|
||||
|
||||
// Support for auth query
|
||||
pub auth_query: Option<String>,
|
||||
pub auth_query_user: Option<String>,
|
||||
pub auth_query_password: Option<String>,
|
||||
|
||||
#[serde(default)]
|
||||
pub prepared_statements: bool,
|
||||
}
|
||||
|
||||
impl General {
|
||||
@@ -335,7 +306,7 @@ impl General {
|
||||
}
|
||||
|
||||
pub fn default_server_lifetime() -> u64 {
|
||||
1000 * 60 * 60 // 1 hour
|
||||
1000 * 60 * 60 * 24 // 24 hours
|
||||
}
|
||||
|
||||
pub fn default_connect_timeout() -> u64 {
|
||||
@@ -358,17 +329,13 @@ impl General {
|
||||
}
|
||||
|
||||
pub fn default_idle_timeout() -> u64 {
|
||||
600000 // 10 minutes
|
||||
60000 // 10 minutes
|
||||
}
|
||||
|
||||
pub fn default_shutdown_timeout() -> u64 {
|
||||
60000
|
||||
}
|
||||
|
||||
pub fn default_dns_max_ttl() -> u64 {
|
||||
30
|
||||
}
|
||||
|
||||
pub fn default_healthcheck_timeout() -> u64 {
|
||||
1000
|
||||
}
|
||||
@@ -388,18 +355,6 @@ impl General {
|
||||
pub fn default_idle_client_in_transaction_timeout() -> u64 {
|
||||
0
|
||||
}
|
||||
|
||||
pub fn default_validate_config() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
pub fn default_prometheus_exporter_port() -> i16 {
|
||||
9930
|
||||
}
|
||||
|
||||
pub fn default_server_round_robin() -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for General {
|
||||
@@ -423,8 +378,6 @@ impl Default for General {
|
||||
log_client_connections: false,
|
||||
log_client_disconnections: false,
|
||||
autoreload: None,
|
||||
dns_cache_enabled: false,
|
||||
dns_max_ttl: Self::default_dns_max_ttl(),
|
||||
tls_certificate: None,
|
||||
tls_private_key: None,
|
||||
server_tls: false,
|
||||
@@ -434,10 +387,7 @@ impl Default for General {
|
||||
auth_query: None,
|
||||
auth_query_user: None,
|
||||
auth_query_password: None,
|
||||
server_lifetime: Self::default_server_lifetime(),
|
||||
server_round_robin: false,
|
||||
validate_config: true,
|
||||
prepared_statements: false,
|
||||
server_lifetime: 1000 * 3600 * 24, // 24 hours,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -490,7 +440,6 @@ pub struct Pool {
|
||||
#[serde(default = "Pool::default_load_balancing_mode")]
|
||||
pub load_balancing_mode: LoadBalancingMode,
|
||||
|
||||
#[serde(default = "Pool::default_default_role")]
|
||||
pub default_role: String,
|
||||
|
||||
#[serde(default)] // False
|
||||
@@ -499,18 +448,12 @@ pub struct Pool {
|
||||
#[serde(default)] // False
|
||||
pub primary_reads_enabled: bool,
|
||||
|
||||
/// Maximum time to allow for establishing a new server connection.
|
||||
pub connect_timeout: Option<u64>,
|
||||
|
||||
/// Close idle connections that have been opened for longer than this.
|
||||
pub idle_timeout: Option<u64>,
|
||||
|
||||
/// Close server connections that have been opened for longer than this.
|
||||
/// Only applied to idle connections. If the connection is actively used for
|
||||
/// longer than this period, the pool will not interrupt it.
|
||||
pub server_lifetime: Option<u64>,
|
||||
|
||||
#[serde(default = "Pool::default_sharding_function")]
|
||||
pub sharding_function: ShardingFunction,
|
||||
|
||||
#[serde(default = "Pool::default_automatic_sharding_key")]
|
||||
@@ -524,10 +467,6 @@ pub struct Pool {
|
||||
pub auth_query_user: Option<String>,
|
||||
pub auth_query_password: Option<String>,
|
||||
|
||||
#[serde(default = "Pool::default_cleanup_server_connections")]
|
||||
pub cleanup_server_connections: bool,
|
||||
|
||||
pub plugins: Option<Plugins>,
|
||||
pub shards: BTreeMap<String, Shard>,
|
||||
pub users: BTreeMap<String, User>,
|
||||
// Note, don't put simple fields below these configs. There's a compatibility issue with TOML that makes it
|
||||
@@ -560,18 +499,6 @@ impl Pool {
|
||||
None
|
||||
}
|
||||
|
||||
pub fn default_default_role() -> String {
|
||||
"any".into()
|
||||
}
|
||||
|
||||
pub fn default_sharding_function() -> ShardingFunction {
|
||||
ShardingFunction::PgBigintHash
|
||||
}
|
||||
|
||||
pub fn default_cleanup_server_connections() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
pub fn validate(&mut self) -> Result<(), Error> {
|
||||
match self.default_role.as_ref() {
|
||||
"any" => (),
|
||||
@@ -660,8 +587,6 @@ impl Default for Pool {
|
||||
auth_query_user: None,
|
||||
auth_query_password: None,
|
||||
server_lifetime: None,
|
||||
plugins: None,
|
||||
cleanup_server_connections: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -740,76 +665,6 @@ impl Default for Shard {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Hash, Eq)]
|
||||
pub struct Plugins {
|
||||
pub intercept: Option<Intercept>,
|
||||
pub table_access: Option<TableAccess>,
|
||||
pub query_logger: Option<QueryLogger>,
|
||||
pub prewarmer: Option<Prewarmer>,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for Plugins {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"interceptor: {}, table_access: {}, query_logger: {}, prewarmer: {}",
|
||||
self.intercept.is_some(),
|
||||
self.table_access.is_some(),
|
||||
self.query_logger.is_some(),
|
||||
self.prewarmer.is_some(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Hash, Eq)]
|
||||
pub struct Intercept {
|
||||
pub enabled: bool,
|
||||
pub queries: BTreeMap<String, Query>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Hash, Eq)]
|
||||
pub struct TableAccess {
|
||||
pub enabled: bool,
|
||||
pub tables: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Hash, Eq)]
|
||||
pub struct QueryLogger {
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Hash, Eq)]
|
||||
pub struct Prewarmer {
|
||||
pub enabled: bool,
|
||||
pub queries: Vec<String>,
|
||||
}
|
||||
|
||||
impl Intercept {
|
||||
pub fn substitute(&mut self, db: &str, user: &str) {
|
||||
for (_, query) in self.queries.iter_mut() {
|
||||
query.substitute(db, user);
|
||||
query.query = query.query.to_ascii_lowercase();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Default, Hash, Eq)]
|
||||
pub struct Query {
|
||||
pub query: String,
|
||||
pub schema: Vec<Vec<String>>,
|
||||
pub result: Vec<Vec<String>>,
|
||||
}
|
||||
|
||||
impl Query {
|
||||
pub fn substitute(&mut self, db: &str, user: &str) {
|
||||
for col in self.result.iter_mut() {
|
||||
for i in 0..col.len() {
|
||||
col[i] = col[i].replace("${USER}", user).replace("${DATABASE}", db);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration wrapper.
|
||||
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
|
||||
pub struct Config {
|
||||
@@ -827,13 +682,7 @@ pub struct Config {
|
||||
#[serde(default = "Config::default_path")]
|
||||
pub path: String,
|
||||
|
||||
// General and global settings.
|
||||
pub general: General,
|
||||
|
||||
// Plugins that should run in all pools.
|
||||
pub plugins: Option<Plugins>,
|
||||
|
||||
// Connection pools.
|
||||
pub pools: HashMap<String, Pool>,
|
||||
}
|
||||
|
||||
@@ -871,7 +720,6 @@ impl Default for Config {
|
||||
path: Self::default_path(),
|
||||
general: General::default(),
|
||||
pools: HashMap::default(),
|
||||
plugins: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -995,7 +843,6 @@ impl Config {
|
||||
"Default max server lifetime: {}ms",
|
||||
self.general.server_lifetime
|
||||
);
|
||||
info!("Sever round robin: {}", self.general.server_round_robin);
|
||||
match self.general.tls_certificate.clone() {
|
||||
Some(tls_certificate) => {
|
||||
info!("TLS certificate: {}", tls_certificate);
|
||||
@@ -1019,14 +866,6 @@ impl Config {
|
||||
"Server TLS certificate verification: {}",
|
||||
self.general.verify_server_certificate
|
||||
);
|
||||
info!("Prepared statements: {}", self.general.prepared_statements);
|
||||
info!(
|
||||
"Plugins: {}",
|
||||
match self.plugins {
|
||||
Some(ref plugins) => plugins.to_string(),
|
||||
None => "not configured".into(),
|
||||
}
|
||||
);
|
||||
|
||||
for (pool_name, pool_config) in &self.pools {
|
||||
// TODO: Make this output prettier (maybe a table?)
|
||||
@@ -1093,18 +932,6 @@ impl Config {
|
||||
None => "default".to_string(),
|
||||
}
|
||||
);
|
||||
info!(
|
||||
"[pool: {}] Cleanup server connections: {}",
|
||||
pool_name, pool_config.cleanup_server_connections
|
||||
);
|
||||
info!(
|
||||
"[pool: {}] Plugins: {}",
|
||||
pool_name,
|
||||
match pool_config.plugins {
|
||||
Some(ref plugins) => plugins.to_string(),
|
||||
None => "not configured".into(),
|
||||
}
|
||||
);
|
||||
|
||||
for user in &pool_config.users {
|
||||
info!(
|
||||
@@ -1244,10 +1071,6 @@ pub fn get_idle_client_in_transaction_timeout() -> u64 {
|
||||
.idle_client_in_transaction_timeout
|
||||
}
|
||||
|
||||
pub fn get_prepared_statements() -> bool {
|
||||
(*(*CONFIG.load())).general.prepared_statements
|
||||
}
|
||||
|
||||
/// Parse the configuration file located at the path.
|
||||
pub async fn parse(path: &str) -> Result<(), Error> {
|
||||
let mut contents = String::new();
|
||||
@@ -1288,7 +1111,6 @@ pub async fn parse(path: &str) -> Result<(), Error> {
|
||||
|
||||
pub async fn reload_config(client_server_map: ClientServerMap) -> Result<bool, Error> {
|
||||
let old_config = get_config();
|
||||
|
||||
match parse(&old_config.path).await {
|
||||
Ok(()) => (),
|
||||
Err(err) => {
|
||||
@@ -1296,18 +1118,14 @@ pub async fn reload_config(client_server_map: ClientServerMap) -> Result<bool, E
|
||||
return Err(Error::BadConfig);
|
||||
}
|
||||
};
|
||||
|
||||
let new_config = get_config();
|
||||
|
||||
match CachedResolver::from_config().await {
|
||||
Ok(_) => (),
|
||||
Err(err) => error!("DNS cache reinitialization error: {:?}", err),
|
||||
};
|
||||
|
||||
if old_config != new_config {
|
||||
info!("Config changed, reloading");
|
||||
if old_config.pools != new_config.pools {
|
||||
info!("Pool configuration changed");
|
||||
ConnectionPool::from_config(client_server_map).await?;
|
||||
Ok(true)
|
||||
} else if old_config != new_config {
|
||||
Ok(true)
|
||||
} else {
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
410
src/dns_cache.rs
410
src/dns_cache.rs
@@ -1,410 +0,0 @@
|
||||
use crate::config::get_config;
|
||||
use crate::errors::Error;
|
||||
use arc_swap::ArcSwap;
|
||||
use log::{debug, error, info, warn};
|
||||
use once_cell::sync::Lazy;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::io;
|
||||
use std::net::IpAddr;
|
||||
use std::sync::Arc;
|
||||
use std::sync::RwLock;
|
||||
use tokio::time::{sleep, Duration};
|
||||
use trust_dns_resolver::error::{ResolveError, ResolveResult};
|
||||
use trust_dns_resolver::lookup_ip::LookupIp;
|
||||
use trust_dns_resolver::TokioAsyncResolver;
|
||||
|
||||
/// Cached Resolver Globally available
|
||||
pub static CACHED_RESOLVER: Lazy<ArcSwap<CachedResolver>> =
|
||||
Lazy::new(|| ArcSwap::from_pointee(CachedResolver::default()));
|
||||
|
||||
// Ip addressed are returned as a set of addresses
|
||||
// so we can compare.
|
||||
#[derive(Clone, PartialEq, Debug)]
|
||||
pub struct AddrSet {
|
||||
set: HashSet<IpAddr>,
|
||||
}
|
||||
|
||||
impl AddrSet {
|
||||
fn new() -> AddrSet {
|
||||
AddrSet {
|
||||
set: HashSet::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<LookupIp> for AddrSet {
|
||||
fn from(lookup_ip: LookupIp) -> Self {
|
||||
let mut addr_set = AddrSet::new();
|
||||
for address in lookup_ip.iter() {
|
||||
addr_set.set.insert(address);
|
||||
}
|
||||
addr_set
|
||||
}
|
||||
}
|
||||
|
||||
///
|
||||
/// A CachedResolver is a DNS resolution cache mechanism with customizable expiration time.
|
||||
///
|
||||
/// The system works as follows:
|
||||
///
|
||||
/// When a host is to be resolved, if we have not resolved it before, a new resolution is
|
||||
/// executed and stored in the internal cache. Concurrently, every `dns_max_ttl` time, the
|
||||
/// cache is refreshed.
|
||||
///
|
||||
/// # Example:
|
||||
///
|
||||
/// ```
|
||||
/// use pgcat::dns_cache::{CachedResolverConfig, CachedResolver};
|
||||
///
|
||||
/// # tokio_test::block_on(async {
|
||||
/// let config = CachedResolverConfig::default();
|
||||
/// let resolver = CachedResolver::new(config, None).await.unwrap();
|
||||
/// let addrset = resolver.lookup_ip("www.example.com.").await.unwrap();
|
||||
/// # })
|
||||
/// ```
|
||||
///
|
||||
/// // Now the ip resolution is stored in local cache and subsequent
|
||||
/// // calls will be returned from cache. Also, the cache is refreshed
|
||||
/// // and updated every 10 seconds.
|
||||
///
|
||||
/// // You can now check if an 'old' lookup differs from what it's currently
|
||||
/// // store in cache by using `has_changed`.
|
||||
/// resolver.has_changed("www.example.com.", addrset)
|
||||
#[derive(Default)]
|
||||
pub struct CachedResolver {
|
||||
// The configuration of the cached_resolver.
|
||||
config: CachedResolverConfig,
|
||||
|
||||
// This is the hash that contains the hash.
|
||||
data: Option<RwLock<HashMap<String, AddrSet>>>,
|
||||
|
||||
// The resolver to be used for DNS queries.
|
||||
resolver: Option<TokioAsyncResolver>,
|
||||
|
||||
// The RefreshLoop
|
||||
refresh_loop: RwLock<Option<tokio::task::JoinHandle<()>>>,
|
||||
}
|
||||
|
||||
///
|
||||
/// Configuration
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub struct CachedResolverConfig {
|
||||
/// Amount of time in secods that a resolved dns address is considered stale.
|
||||
dns_max_ttl: u64,
|
||||
|
||||
/// Enabled or disabled? (this is so we can reload config)
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
impl CachedResolverConfig {
|
||||
fn new(dns_max_ttl: u64, enabled: bool) -> Self {
|
||||
CachedResolverConfig {
|
||||
dns_max_ttl,
|
||||
enabled,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<crate::config::Config> for CachedResolverConfig {
|
||||
fn from(config: crate::config::Config) -> Self {
|
||||
CachedResolverConfig::new(config.general.dns_max_ttl, config.general.dns_cache_enabled)
|
||||
}
|
||||
}
|
||||
|
||||
impl CachedResolver {
|
||||
///
|
||||
/// Returns a new Arc<CachedResolver> based on passed configuration.
|
||||
/// It also starts the loop that will refresh cache entries.
|
||||
///
|
||||
/// # Arguments:
|
||||
///
|
||||
/// * `config` - The `CachedResolverConfig` to be used to create the resolver.
|
||||
///
|
||||
/// # Example:
|
||||
///
|
||||
/// ```
|
||||
/// use pgcat::dns_cache::{CachedResolverConfig, CachedResolver};
|
||||
///
|
||||
/// # tokio_test::block_on(async {
|
||||
/// let config = CachedResolverConfig::default();
|
||||
/// let resolver = CachedResolver::new(config, None).await.unwrap();
|
||||
/// # })
|
||||
/// ```
|
||||
///
|
||||
pub async fn new(
|
||||
config: CachedResolverConfig,
|
||||
data: Option<HashMap<String, AddrSet>>,
|
||||
) -> Result<Arc<Self>, io::Error> {
|
||||
// Construct a new Resolver with default configuration options
|
||||
let resolver = Some(TokioAsyncResolver::tokio_from_system_conf()?);
|
||||
|
||||
let data = if let Some(hash) = data {
|
||||
Some(RwLock::new(hash))
|
||||
} else {
|
||||
Some(RwLock::new(HashMap::new()))
|
||||
};
|
||||
|
||||
let instance = Arc::new(Self {
|
||||
config,
|
||||
resolver,
|
||||
data,
|
||||
refresh_loop: RwLock::new(None),
|
||||
});
|
||||
|
||||
if instance.enabled() {
|
||||
info!("Scheduling DNS refresh loop");
|
||||
let refresh_loop = tokio::task::spawn({
|
||||
let instance = instance.clone();
|
||||
async move {
|
||||
instance.refresh_dns_entries_loop().await;
|
||||
}
|
||||
});
|
||||
*(instance.refresh_loop.write().unwrap()) = Some(refresh_loop);
|
||||
}
|
||||
|
||||
Ok(instance)
|
||||
}
|
||||
|
||||
pub fn enabled(&self) -> bool {
|
||||
self.config.enabled
|
||||
}
|
||||
|
||||
// Schedules the refresher
|
||||
async fn refresh_dns_entries_loop(&self) {
|
||||
let resolver = TokioAsyncResolver::tokio_from_system_conf().unwrap();
|
||||
let interval = Duration::from_secs(self.config.dns_max_ttl);
|
||||
loop {
|
||||
debug!("Begin refreshing cached DNS addresses.");
|
||||
// To minimize the time we hold the lock, we first create
|
||||
// an array with keys.
|
||||
let mut hostnames: Vec<String> = Vec::new();
|
||||
{
|
||||
if let Some(ref data) = self.data {
|
||||
for hostname in data.read().unwrap().keys() {
|
||||
hostnames.push(hostname.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for hostname in hostnames.iter() {
|
||||
let addrset = self
|
||||
.fetch_from_cache(hostname.as_str())
|
||||
.expect("Could not obtain expected address from cache, this should not happen");
|
||||
|
||||
match resolver.lookup_ip(hostname).await {
|
||||
Ok(lookup_ip) => {
|
||||
let new_addrset = AddrSet::from(lookup_ip);
|
||||
debug!(
|
||||
"Obtained address for host ({}) -> ({:?})",
|
||||
hostname, new_addrset
|
||||
);
|
||||
|
||||
if addrset != new_addrset {
|
||||
debug!(
|
||||
"Addr changed from {:?} to {:?} updating cache.",
|
||||
addrset, new_addrset
|
||||
);
|
||||
self.store_in_cache(hostname, new_addrset);
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
error!(
|
||||
"There was an error trying to resolv {}: ({}).",
|
||||
hostname, err
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
debug!("Finished refreshing cached DNS addresses.");
|
||||
sleep(interval).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns a `AddrSet` given the specified hostname.
|
||||
///
|
||||
/// This method first tries to fetch the value from the cache, if it misses
|
||||
/// then it is resolved and stored in the cache. TTL from records is ignored.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `host` - A string slice referencing the hostname to be resolved.
|
||||
///
|
||||
/// # Example:
|
||||
///
|
||||
/// ```
|
||||
/// use pgcat::dns_cache::{CachedResolverConfig, CachedResolver};
|
||||
///
|
||||
/// # tokio_test::block_on(async {
|
||||
/// let config = CachedResolverConfig::default();
|
||||
/// let resolver = CachedResolver::new(config, None).await.unwrap();
|
||||
/// let response = resolver.lookup_ip("www.google.com.");
|
||||
/// # })
|
||||
/// ```
|
||||
///
|
||||
pub async fn lookup_ip(&self, host: &str) -> ResolveResult<AddrSet> {
|
||||
debug!("Lookup up {} in cache", host);
|
||||
match self.fetch_from_cache(host) {
|
||||
Some(addr_set) => {
|
||||
debug!("Cache hit!");
|
||||
Ok(addr_set)
|
||||
}
|
||||
None => {
|
||||
debug!("Not found, executing a dns query!");
|
||||
if let Some(ref resolver) = self.resolver {
|
||||
let addr_set = AddrSet::from(resolver.lookup_ip(host).await?);
|
||||
debug!("Obtained: {:?}", addr_set);
|
||||
self.store_in_cache(host, addr_set.clone());
|
||||
Ok(addr_set)
|
||||
} else {
|
||||
Err(ResolveError::from("No resolver available"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Returns true if the stored host resolution differs from the AddrSet passed.
|
||||
pub fn has_changed(&self, host: &str, addr_set: &AddrSet) -> bool {
|
||||
if let Some(fetched_addr_set) = self.fetch_from_cache(host) {
|
||||
return fetched_addr_set != *addr_set;
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
// Fetches an AddrSet from the inner cache adquiring the read lock.
|
||||
fn fetch_from_cache(&self, key: &str) -> Option<AddrSet> {
|
||||
if let Some(ref hash) = self.data {
|
||||
if let Some(addr_set) = hash.read().unwrap().get(key) {
|
||||
return Some(addr_set.clone());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
// Sets up the global CACHED_RESOLVER static variable so we can globally use DNS
|
||||
// cache.
|
||||
pub async fn from_config() -> Result<(), Error> {
|
||||
let cached_resolver = CACHED_RESOLVER.load();
|
||||
let desired_config = CachedResolverConfig::from(get_config());
|
||||
|
||||
if cached_resolver.config != desired_config {
|
||||
if let Some(ref refresh_loop) = *(cached_resolver.refresh_loop.write().unwrap()) {
|
||||
warn!("Killing Dnscache refresh loop as its configuration is being reloaded");
|
||||
refresh_loop.abort()
|
||||
}
|
||||
let new_resolver = if let Some(ref data) = cached_resolver.data {
|
||||
let data = Some(data.read().unwrap().clone());
|
||||
CachedResolver::new(desired_config, data).await
|
||||
} else {
|
||||
CachedResolver::new(desired_config, None).await
|
||||
};
|
||||
|
||||
match new_resolver {
|
||||
Ok(ok) => {
|
||||
CACHED_RESOLVER.store(ok);
|
||||
Ok(())
|
||||
}
|
||||
Err(err) => {
|
||||
let message = format!("Error setting up cached_resolver. Error: {:?}, will continue without this feature.", err);
|
||||
Err(Error::DNSCachedError(message))
|
||||
}
|
||||
}
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// Stores the AddrSet in cache adquiring the write lock.
|
||||
fn store_in_cache(&self, host: &str, addr_set: AddrSet) {
|
||||
if let Some(ref data) = self.data {
|
||||
data.write().unwrap().insert(host.to_string(), addr_set);
|
||||
} else {
|
||||
error!("Could not insert, Hash not initialized");
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use trust_dns_resolver::error::ResolveError;
|
||||
|
||||
#[tokio::test]
|
||||
async fn new() {
|
||||
let config = CachedResolverConfig {
|
||||
dns_max_ttl: 10,
|
||||
enabled: true,
|
||||
};
|
||||
let resolver = CachedResolver::new(config, None).await;
|
||||
assert!(resolver.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn lookup_ip() {
|
||||
let config = CachedResolverConfig {
|
||||
dns_max_ttl: 10,
|
||||
enabled: true,
|
||||
};
|
||||
let resolver = CachedResolver::new(config, None).await.unwrap();
|
||||
let response = resolver.lookup_ip("www.google.com.").await;
|
||||
assert!(response.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn has_changed() {
|
||||
let config = CachedResolverConfig {
|
||||
dns_max_ttl: 10,
|
||||
enabled: true,
|
||||
};
|
||||
let resolver = CachedResolver::new(config, None).await.unwrap();
|
||||
let hostname = "www.google.com.";
|
||||
let response = resolver.lookup_ip(hostname).await;
|
||||
let addr_set = response.unwrap();
|
||||
assert!(!resolver.has_changed(hostname, &addr_set));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn unknown_host() {
|
||||
let config = CachedResolverConfig {
|
||||
dns_max_ttl: 10,
|
||||
enabled: true,
|
||||
};
|
||||
let resolver = CachedResolver::new(config, None).await.unwrap();
|
||||
let hostname = "www.idontexists.";
|
||||
let response = resolver.lookup_ip(hostname).await;
|
||||
assert!(matches!(response, Err(ResolveError { .. })));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn incorrect_address() {
|
||||
let config = CachedResolverConfig {
|
||||
dns_max_ttl: 10,
|
||||
enabled: true,
|
||||
};
|
||||
let resolver = CachedResolver::new(config, None).await.unwrap();
|
||||
let hostname = "w ww.idontexists.";
|
||||
let response = resolver.lookup_ip(hostname).await;
|
||||
assert!(matches!(response, Err(ResolveError { .. })));
|
||||
assert!(!resolver.has_changed(hostname, &AddrSet::new()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
// Ok, this test is based on the fact that google does DNS RR
|
||||
// and does not responds with every available ip everytime, so
|
||||
// if I cache here, it will miss after one cache iteration or two.
|
||||
async fn thread() {
|
||||
let config = CachedResolverConfig {
|
||||
dns_max_ttl: 10,
|
||||
enabled: true,
|
||||
};
|
||||
let resolver = CachedResolver::new(config, None).await.unwrap();
|
||||
let hostname = "www.google.com.";
|
||||
let response = resolver.lookup_ip(hostname).await;
|
||||
let addr_set = response.unwrap();
|
||||
assert!(!resolver.has_changed(hostname, &addr_set));
|
||||
let resolver_for_refresher = resolver.clone();
|
||||
let _thread_handle = tokio::task::spawn(async move {
|
||||
resolver_for_refresher.refresh_dns_entries_loop().await;
|
||||
});
|
||||
assert!(!resolver.has_changed(hostname, &addr_set));
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,7 @@
|
||||
//! Errors.
|
||||
|
||||
/// Various errors.
|
||||
#[derive(Debug, PartialEq, Clone)]
|
||||
#[derive(Debug, PartialEq)]
|
||||
pub enum Error {
|
||||
SocketError(String),
|
||||
ClientSocketError(String, ClientIdentifier),
|
||||
@@ -19,14 +19,10 @@ pub enum Error {
|
||||
ClientError(String),
|
||||
TlsError,
|
||||
StatementTimeout,
|
||||
DNSCachedError(String),
|
||||
ShuttingDown,
|
||||
ParseBytesError(String),
|
||||
AuthError(String),
|
||||
AuthPassthroughError(String),
|
||||
UnsupportedStatement,
|
||||
QueryRouterParserError(String),
|
||||
QueryRouterError(String),
|
||||
}
|
||||
|
||||
#[derive(Clone, PartialEq, Debug)]
|
||||
@@ -122,9 +118,3 @@ impl std::fmt::Display for Error {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<std::ffi::NulError> for Error {
|
||||
fn from(err: std::ffi::NulError) -> Self {
|
||||
Error::QueryRouterError(err.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,17 +1,11 @@
|
||||
pub mod admin;
|
||||
pub mod auth_passthrough;
|
||||
pub mod client;
|
||||
pub mod config;
|
||||
pub mod constants;
|
||||
pub mod dns_cache;
|
||||
pub mod errors;
|
||||
pub mod messages;
|
||||
pub mod mirrors;
|
||||
pub mod multi_logger;
|
||||
pub mod plugins;
|
||||
pub mod pool;
|
||||
pub mod prometheus;
|
||||
pub mod query_router;
|
||||
pub mod scram;
|
||||
pub mod server;
|
||||
pub mod sharding;
|
||||
|
||||
46
src/main.rs
46
src/main.rs
@@ -36,7 +36,6 @@ extern crate sqlparser;
|
||||
extern crate tokio;
|
||||
extern crate tokio_rustls;
|
||||
extern crate toml;
|
||||
extern crate trust_dns_resolver;
|
||||
|
||||
#[cfg(not(target_env = "msvc"))]
|
||||
use jemallocator::Jemalloc;
|
||||
@@ -61,19 +60,36 @@ use std::str::FromStr;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use pgcat::config::{get_config, reload_config, VERSION};
|
||||
use pgcat::dns_cache;
|
||||
use pgcat::messages::configure_socket;
|
||||
use pgcat::pool::{ClientServerMap, ConnectionPool};
|
||||
use pgcat::prometheus::start_metric_server;
|
||||
use pgcat::stats::{Collector, Reporter, REPORTER};
|
||||
mod admin;
|
||||
mod auth_passthrough;
|
||||
mod client;
|
||||
mod config;
|
||||
mod constants;
|
||||
mod errors;
|
||||
mod messages;
|
||||
mod mirrors;
|
||||
mod multi_logger;
|
||||
mod pool;
|
||||
mod prometheus;
|
||||
mod query_router;
|
||||
mod scram;
|
||||
mod server;
|
||||
mod sharding;
|
||||
mod stats;
|
||||
mod tls;
|
||||
|
||||
use crate::config::{get_config, reload_config, VERSION};
|
||||
use crate::messages::configure_socket;
|
||||
use crate::pool::{ClientServerMap, ConnectionPool};
|
||||
use crate::prometheus::start_metric_server;
|
||||
use crate::stats::{Collector, Reporter, REPORTER};
|
||||
|
||||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
pgcat::multi_logger::MultiLogger::init().unwrap();
|
||||
multi_logger::MultiLogger::init().unwrap();
|
||||
|
||||
info!("Welcome to PgCat! Meow. (Version {})", VERSION);
|
||||
|
||||
if !pgcat::query_router::QueryRouter::setup() {
|
||||
if !query_router::QueryRouter::setup() {
|
||||
error!("Could not setup query router");
|
||||
std::process::exit(exitcode::CONFIG);
|
||||
}
|
||||
@@ -91,7 +107,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let runtime = Builder::new_multi_thread().worker_threads(1).build()?;
|
||||
|
||||
runtime.block_on(async {
|
||||
match pgcat::config::parse(&config_file).await {
|
||||
match config::parse(&config_file).await {
|
||||
Ok(_) => (),
|
||||
Err(err) => {
|
||||
error!("Config parse error: {:?}", err);
|
||||
@@ -150,12 +166,6 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
// Statistics reporting.
|
||||
REPORTER.store(Arc::new(Reporter::default()));
|
||||
|
||||
// Starts (if enabled) dns cache before pools initialization
|
||||
match dns_cache::CachedResolver::from_config().await {
|
||||
Ok(_) => (),
|
||||
Err(err) => error!("DNS cache initialization error: {:?}", err),
|
||||
};
|
||||
|
||||
// Connection pool that allows to query all shards and replicas.
|
||||
match ConnectionPool::from_config(client_server_map.clone()).await {
|
||||
Ok(_) => (),
|
||||
@@ -285,7 +295,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
tokio::task::spawn(async move {
|
||||
let start = chrono::offset::Utc::now().naive_utc();
|
||||
|
||||
match pgcat::client::client_entrypoint(
|
||||
match client::client_entrypoint(
|
||||
socket,
|
||||
client_server_map,
|
||||
shutdown_rx,
|
||||
@@ -316,7 +326,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
|
||||
Err(err) => {
|
||||
match err {
|
||||
pgcat::errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err),
|
||||
errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err),
|
||||
_ => warn!("Client disconnected with error {:?}", err),
|
||||
}
|
||||
|
||||
|
||||
328
src/messages.rs
328
src/messages.rs
@@ -7,15 +7,11 @@ use socket2::{SockRef, TcpKeepalive};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpStream;
|
||||
|
||||
use crate::client::PREPARED_STATEMENT_COUNTER;
|
||||
use crate::config::get_config;
|
||||
use crate::errors::Error;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::ffi::CString;
|
||||
use std::io::{BufRead, Cursor};
|
||||
use std::mem;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Postgres data type mappings
|
||||
@@ -24,10 +20,6 @@ pub enum DataType {
|
||||
Text,
|
||||
Int4,
|
||||
Numeric,
|
||||
Bool,
|
||||
Oid,
|
||||
AnyArray,
|
||||
Any,
|
||||
}
|
||||
|
||||
impl From<&DataType> for i32 {
|
||||
@@ -36,10 +28,6 @@ impl From<&DataType> for i32 {
|
||||
DataType::Text => 25,
|
||||
DataType::Int4 => 23,
|
||||
DataType::Numeric => 1700,
|
||||
DataType::Bool => 16,
|
||||
DataType::Oid => 26,
|
||||
DataType::AnyArray => 2277,
|
||||
DataType::Any => 2276,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -455,10 +443,6 @@ pub fn row_description(columns: &Vec<(&str, DataType)>) -> BytesMut {
|
||||
DataType::Text => -1,
|
||||
DataType::Int4 => 4,
|
||||
DataType::Numeric => -1,
|
||||
DataType::Bool => 1,
|
||||
DataType::Oid => 4,
|
||||
DataType::AnyArray => -1,
|
||||
DataType::Any => -1,
|
||||
};
|
||||
|
||||
row_desc.put_i16(type_size);
|
||||
@@ -497,29 +481,6 @@ pub fn data_row(row: &Vec<String>) -> BytesMut {
|
||||
res
|
||||
}
|
||||
|
||||
pub fn data_row_nullable(row: &Vec<Option<String>>) -> BytesMut {
|
||||
let mut res = BytesMut::new();
|
||||
let mut data_row = BytesMut::new();
|
||||
|
||||
data_row.put_i16(row.len() as i16);
|
||||
|
||||
for column in row {
|
||||
if let Some(column) = column {
|
||||
let column = column.as_bytes();
|
||||
data_row.put_i32(column.len() as i32);
|
||||
data_row.put_slice(column);
|
||||
} else {
|
||||
data_row.put_i32(-1 as i32);
|
||||
}
|
||||
}
|
||||
|
||||
res.put_u8(b'D');
|
||||
res.put_i32(data_row.len() as i32 + 4);
|
||||
res.put(data_row);
|
||||
|
||||
res
|
||||
}
|
||||
|
||||
/// Create a CommandComplete message.
|
||||
pub fn command_complete(command: &str) -> BytesMut {
|
||||
let cmd = BytesMut::from(format!("{}\0", command).as_bytes());
|
||||
@@ -530,13 +491,6 @@ pub fn command_complete(command: &str) -> BytesMut {
|
||||
res
|
||||
}
|
||||
|
||||
pub fn flush() -> BytesMut {
|
||||
let mut bytes = BytesMut::new();
|
||||
bytes.put_u8(b'H');
|
||||
bytes.put_i32(4);
|
||||
bytes
|
||||
}
|
||||
|
||||
/// Write all data in the buffer to the TcpStream.
|
||||
pub async fn write_all<S>(stream: &mut S, buf: BytesMut) -> Result<(), Error>
|
||||
where
|
||||
@@ -700,285 +654,3 @@ impl BytesMutReader for Cursor<&BytesMut> {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse (F) message.
|
||||
/// See: <https://www.postgresql.org/docs/current/protocol-message-formats.html>
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Parse {
|
||||
code: char,
|
||||
#[allow(dead_code)]
|
||||
len: i32,
|
||||
pub name: String,
|
||||
pub generated_name: String,
|
||||
query: String,
|
||||
num_params: i16,
|
||||
param_types: Vec<i32>,
|
||||
}
|
||||
|
||||
impl TryFrom<&BytesMut> for Parse {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(buf: &BytesMut) -> Result<Parse, Error> {
|
||||
let mut cursor = Cursor::new(buf);
|
||||
let code = cursor.get_u8() as char;
|
||||
let len = cursor.get_i32();
|
||||
let name = cursor.read_string()?;
|
||||
let query = cursor.read_string()?;
|
||||
let num_params = cursor.get_i16();
|
||||
let mut param_types = Vec::new();
|
||||
|
||||
for _ in 0..num_params {
|
||||
param_types.push(cursor.get_i32());
|
||||
}
|
||||
|
||||
Ok(Parse {
|
||||
code,
|
||||
len,
|
||||
name,
|
||||
generated_name: prepared_statement_name(),
|
||||
query,
|
||||
num_params,
|
||||
param_types,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<Parse> for BytesMut {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(parse: Parse) -> Result<BytesMut, Error> {
|
||||
let mut bytes = BytesMut::new();
|
||||
|
||||
let name_binding = CString::new(parse.name)?;
|
||||
let name = name_binding.as_bytes_with_nul();
|
||||
|
||||
let query_binding = CString::new(parse.query)?;
|
||||
let query = query_binding.as_bytes_with_nul();
|
||||
|
||||
// Recompute length of the message.
|
||||
let len = 4 // self
|
||||
+ name.len()
|
||||
+ query.len()
|
||||
+ 2
|
||||
+ 4 * parse.num_params as usize;
|
||||
|
||||
bytes.put_u8(parse.code as u8);
|
||||
bytes.put_i32(len as i32);
|
||||
bytes.put_slice(name);
|
||||
bytes.put_slice(query);
|
||||
bytes.put_i16(parse.num_params);
|
||||
for param in parse.param_types {
|
||||
bytes.put_i32(param);
|
||||
}
|
||||
|
||||
Ok(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<&Parse> for BytesMut {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(parse: &Parse) -> Result<BytesMut, Error> {
|
||||
parse.clone().try_into()
|
||||
}
|
||||
}
|
||||
|
||||
impl Parse {
|
||||
pub fn rename(mut self) -> Self {
|
||||
self.name = self.generated_name.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn anonymous(&self) -> bool {
|
||||
self.name.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
/// Bind (B) message.
|
||||
/// See: <https://www.postgresql.org/docs/current/protocol-message-formats.html>
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Bind {
|
||||
code: char,
|
||||
#[allow(dead_code)]
|
||||
len: i64,
|
||||
portal: String,
|
||||
pub prepared_statement: String,
|
||||
num_param_format_codes: i16,
|
||||
param_format_codes: Vec<i16>,
|
||||
num_param_values: i16,
|
||||
param_values: Vec<(i32, BytesMut)>,
|
||||
num_result_column_format_codes: i16,
|
||||
result_columns_format_codes: Vec<i16>,
|
||||
}
|
||||
|
||||
impl TryFrom<&BytesMut> for Bind {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(buf: &BytesMut) -> Result<Bind, Error> {
|
||||
let mut cursor = Cursor::new(buf);
|
||||
let code = cursor.get_u8() as char;
|
||||
let len = cursor.get_i32();
|
||||
let portal = cursor.read_string()?;
|
||||
let prepared_statement = cursor.read_string()?;
|
||||
let num_param_format_codes = cursor.get_i16();
|
||||
let mut param_format_codes = Vec::new();
|
||||
|
||||
for _ in 0..num_param_format_codes {
|
||||
param_format_codes.push(cursor.get_i16());
|
||||
}
|
||||
|
||||
let num_param_values = cursor.get_i16();
|
||||
let mut param_values = Vec::new();
|
||||
|
||||
for _ in 0..num_param_values {
|
||||
let param_len = cursor.get_i32();
|
||||
let mut param = BytesMut::with_capacity(param_len as usize);
|
||||
param.resize(param_len as usize, b'0');
|
||||
cursor.copy_to_slice(&mut param);
|
||||
param_values.push((param_len, param));
|
||||
}
|
||||
|
||||
let num_result_column_format_codes = cursor.get_i16();
|
||||
let mut result_columns_format_codes = Vec::new();
|
||||
|
||||
for _ in 0..num_result_column_format_codes {
|
||||
result_columns_format_codes.push(cursor.get_i16());
|
||||
}
|
||||
|
||||
Ok(Bind {
|
||||
code,
|
||||
len: len as i64,
|
||||
portal,
|
||||
prepared_statement,
|
||||
num_param_format_codes,
|
||||
param_format_codes,
|
||||
num_param_values,
|
||||
param_values,
|
||||
num_result_column_format_codes,
|
||||
result_columns_format_codes,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<Bind> for BytesMut {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(bind: Bind) -> Result<BytesMut, Error> {
|
||||
let mut bytes = BytesMut::new();
|
||||
|
||||
let portal_binding = CString::new(bind.portal)?;
|
||||
let portal = portal_binding.as_bytes_with_nul();
|
||||
|
||||
let prepared_statement_binding = CString::new(bind.prepared_statement)?;
|
||||
let prepared_statement = prepared_statement_binding.as_bytes_with_nul();
|
||||
|
||||
let mut len = 4 // self
|
||||
+ portal.len()
|
||||
+ prepared_statement.len()
|
||||
+ 2 // num_param_format_codes
|
||||
+ 2 * bind.num_param_format_codes as usize // num_param_format_codes
|
||||
+ 2; // num_param_values
|
||||
|
||||
for (param_len, _) in &bind.param_values {
|
||||
len += 4 + *param_len as usize;
|
||||
}
|
||||
len += 2; // num_result_column_format_codes
|
||||
len += 2 * bind.num_result_column_format_codes as usize;
|
||||
|
||||
bytes.put_u8(bind.code as u8);
|
||||
bytes.put_i32(len as i32);
|
||||
bytes.put_slice(portal);
|
||||
bytes.put_slice(prepared_statement);
|
||||
bytes.put_i16(bind.num_param_format_codes);
|
||||
for param_format_code in bind.param_format_codes {
|
||||
bytes.put_i16(param_format_code);
|
||||
}
|
||||
bytes.put_i16(bind.num_param_values);
|
||||
for (param_len, param) in bind.param_values {
|
||||
bytes.put_i32(param_len);
|
||||
bytes.put_slice(¶m);
|
||||
}
|
||||
bytes.put_i16(bind.num_result_column_format_codes);
|
||||
for result_column_format_code in bind.result_columns_format_codes {
|
||||
bytes.put_i16(result_column_format_code);
|
||||
}
|
||||
|
||||
Ok(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
impl Bind {
|
||||
pub fn reassign(mut self, parse: &Parse) -> Self {
|
||||
self.prepared_statement = parse.name.clone();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn anonymous(&self) -> bool {
|
||||
self.prepared_statement.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Describe {
|
||||
code: char,
|
||||
|
||||
#[allow(dead_code)]
|
||||
len: i32,
|
||||
target: char,
|
||||
pub statement_name: String,
|
||||
}
|
||||
|
||||
impl TryFrom<&BytesMut> for Describe {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(bytes: &BytesMut) -> Result<Describe, Error> {
|
||||
let mut cursor = Cursor::new(bytes);
|
||||
let code = cursor.get_u8() as char;
|
||||
let len = cursor.get_i32();
|
||||
let target = cursor.get_u8() as char;
|
||||
let statement_name = cursor.read_string()?;
|
||||
|
||||
Ok(Describe {
|
||||
code,
|
||||
len,
|
||||
target,
|
||||
statement_name,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<Describe> for BytesMut {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(describe: Describe) -> Result<BytesMut, Error> {
|
||||
let mut bytes = BytesMut::new();
|
||||
let statement_name_binding = CString::new(describe.statement_name)?;
|
||||
let statement_name = statement_name_binding.as_bytes_with_nul();
|
||||
let len = 4 + 1 + statement_name.len();
|
||||
|
||||
bytes.put_u8(describe.code as u8);
|
||||
bytes.put_i32(len as i32);
|
||||
bytes.put_u8(describe.target as u8);
|
||||
bytes.put_slice(statement_name);
|
||||
|
||||
Ok(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
impl Describe {
|
||||
pub fn rename(mut self, name: &str) -> Self {
|
||||
self.statement_name = name.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
pub fn anonymous(&self) -> bool {
|
||||
self.statement_name.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn prepared_statement_name() -> String {
|
||||
format!(
|
||||
"P_{}",
|
||||
PREPARED_STATEMENT_COUNTER.fetch_add(1, Ordering::SeqCst)
|
||||
)
|
||||
}
|
||||
|
||||
@@ -7,7 +7,8 @@ use bytes::{Bytes, BytesMut};
|
||||
use parking_lot::RwLock;
|
||||
|
||||
use crate::config::{get_config, Address, Role, User};
|
||||
use crate::pool::{ClientServerMap, ServerPool};
|
||||
use crate::pool::{ClientServerMap, PoolIdentifier, ServerPool};
|
||||
use crate::stats::PoolStats;
|
||||
use log::{error, info, trace, warn};
|
||||
use tokio::sync::mpsc::{channel, Receiver, Sender};
|
||||
|
||||
@@ -23,7 +24,7 @@ impl MirroredClient {
|
||||
async fn create_pool(&self) -> Pool<ServerPool> {
|
||||
let config = get_config();
|
||||
let default = std::time::Duration::from_millis(10_000).as_millis() as u64;
|
||||
let (connection_timeout, idle_timeout, _cfg) =
|
||||
let (connection_timeout, idle_timeout, cfg) =
|
||||
match config.pools.get(&self.address.pool_name) {
|
||||
Some(cfg) => (
|
||||
cfg.connect_timeout.unwrap_or(default),
|
||||
@@ -33,14 +34,15 @@ impl MirroredClient {
|
||||
None => (default, default, crate::config::Pool::default()),
|
||||
};
|
||||
|
||||
let identifier = PoolIdentifier::new(&self.database, &self.user.username);
|
||||
|
||||
let manager = ServerPool::new(
|
||||
self.address.clone(),
|
||||
self.user.clone(),
|
||||
self.database.as_str(),
|
||||
ClientServerMap::default(),
|
||||
Arc::new(PoolStats::new(identifier, cfg.clone())),
|
||||
Arc::new(RwLock::new(None)),
|
||||
None,
|
||||
true,
|
||||
);
|
||||
|
||||
Pool::builder()
|
||||
|
||||
@@ -1,120 +0,0 @@
|
||||
//! The intercept plugin.
|
||||
//!
|
||||
//! It intercepts queries and returns fake results.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use bytes::{BufMut, BytesMut};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sqlparser::ast::Statement;
|
||||
|
||||
use log::debug;
|
||||
|
||||
use crate::{
|
||||
config::Intercept as InterceptConfig,
|
||||
errors::Error,
|
||||
messages::{command_complete, data_row_nullable, row_description, DataType},
|
||||
plugins::{Plugin, PluginOutput},
|
||||
query_router::QueryRouter,
|
||||
};
|
||||
|
||||
// TODO: use these structs for deserialization
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub struct Rule {
|
||||
query: String,
|
||||
schema: Vec<Column>,
|
||||
result: Vec<Vec<String>>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
pub struct Column {
|
||||
name: String,
|
||||
data_type: String,
|
||||
}
|
||||
|
||||
/// The intercept plugin.
|
||||
pub struct Intercept<'a> {
|
||||
pub enabled: bool,
|
||||
pub config: &'a InterceptConfig,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<'a> Plugin for Intercept<'a> {
|
||||
async fn run(
|
||||
&mut self,
|
||||
query_router: &QueryRouter,
|
||||
ast: &Vec<Statement>,
|
||||
) -> Result<PluginOutput, Error> {
|
||||
if !self.enabled || ast.is_empty() {
|
||||
return Ok(PluginOutput::Allow);
|
||||
}
|
||||
|
||||
let mut config = self.config.clone();
|
||||
config.substitute(
|
||||
&query_router.pool_settings().db,
|
||||
&query_router.pool_settings().user.username,
|
||||
);
|
||||
|
||||
let mut result = BytesMut::new();
|
||||
|
||||
for q in ast {
|
||||
// Normalization
|
||||
let q = q.to_string().to_ascii_lowercase();
|
||||
|
||||
for (_, target) in config.queries.iter() {
|
||||
if target.query.as_str() == q {
|
||||
debug!("Intercepting query: {}", q);
|
||||
|
||||
let rd = target
|
||||
.schema
|
||||
.iter()
|
||||
.map(|row| {
|
||||
let name = &row[0];
|
||||
let data_type = &row[1];
|
||||
(
|
||||
name.as_str(),
|
||||
match data_type.as_str() {
|
||||
"text" => DataType::Text,
|
||||
"anyarray" => DataType::AnyArray,
|
||||
"oid" => DataType::Oid,
|
||||
"bool" => DataType::Bool,
|
||||
"int4" => DataType::Int4,
|
||||
_ => DataType::Any,
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect::<Vec<(&str, DataType)>>();
|
||||
|
||||
result.put(row_description(&rd));
|
||||
|
||||
target.result.iter().for_each(|row| {
|
||||
let row = row
|
||||
.iter()
|
||||
.map(|s| {
|
||||
let s = s.as_str().to_string();
|
||||
|
||||
if s == "" {
|
||||
None
|
||||
} else {
|
||||
Some(s)
|
||||
}
|
||||
})
|
||||
.collect::<Vec<Option<String>>>();
|
||||
result.put(data_row_nullable(&row));
|
||||
});
|
||||
|
||||
result.put(command_complete("SELECT"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !result.is_empty() {
|
||||
result.put_u8(b'Z');
|
||||
result.put_i32(5);
|
||||
result.put_u8(b'I');
|
||||
|
||||
return Ok(PluginOutput::Intercept(result));
|
||||
} else {
|
||||
Ok(PluginOutput::Allow)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
//! The plugin ecosystem.
|
||||
//!
|
||||
//! Currently plugins only grant access or deny access to the database for a particual query.
|
||||
//! Example use cases:
|
||||
//! - block known bad queries
|
||||
//! - block access to system catalogs
|
||||
//! - block dangerous modifications like `DROP TABLE`
|
||||
//! - etc
|
||||
//!
|
||||
|
||||
pub mod intercept;
|
||||
pub mod prewarmer;
|
||||
pub mod query_logger;
|
||||
pub mod table_access;
|
||||
|
||||
use crate::{errors::Error, query_router::QueryRouter};
|
||||
use async_trait::async_trait;
|
||||
use bytes::BytesMut;
|
||||
use sqlparser::ast::Statement;
|
||||
|
||||
pub use intercept::Intercept;
|
||||
pub use query_logger::QueryLogger;
|
||||
pub use table_access::TableAccess;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum PluginOutput {
|
||||
Allow,
|
||||
Deny(String),
|
||||
Overwrite(Vec<Statement>),
|
||||
Intercept(BytesMut),
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait Plugin {
|
||||
// Run before the query is sent to the server.
|
||||
async fn run(
|
||||
&mut self,
|
||||
query_router: &QueryRouter,
|
||||
ast: &Vec<Statement>,
|
||||
) -> Result<PluginOutput, Error>;
|
||||
|
||||
// TODO: run after the result is returned
|
||||
// async fn callback(&mut self, query_router: &QueryRouter);
|
||||
}
|
||||
@@ -1,28 +0,0 @@
|
||||
//! Prewarm new connections before giving them to the client.
|
||||
use crate::{errors::Error, server::Server};
|
||||
use log::info;
|
||||
|
||||
pub struct Prewarmer<'a> {
|
||||
pub enabled: bool,
|
||||
pub server: &'a mut Server,
|
||||
pub queries: &'a Vec<String>,
|
||||
}
|
||||
|
||||
impl<'a> Prewarmer<'a> {
|
||||
pub async fn run(&mut self) -> Result<(), Error> {
|
||||
if !self.enabled {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
for query in self.queries {
|
||||
info!(
|
||||
"{} Prewarning with query: `{}`",
|
||||
self.server.address(),
|
||||
query
|
||||
);
|
||||
self.server.query(&query).await?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,38 +0,0 @@
|
||||
//! Log all queries to stdout (or somewhere else, why not).
|
||||
|
||||
use crate::{
|
||||
errors::Error,
|
||||
plugins::{Plugin, PluginOutput},
|
||||
query_router::QueryRouter,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use log::info;
|
||||
use sqlparser::ast::Statement;
|
||||
|
||||
pub struct QueryLogger<'a> {
|
||||
pub enabled: bool,
|
||||
pub user: &'a str,
|
||||
pub db: &'a str,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<'a> Plugin for QueryLogger<'a> {
|
||||
async fn run(
|
||||
&mut self,
|
||||
_query_router: &QueryRouter,
|
||||
ast: &Vec<Statement>,
|
||||
) -> Result<PluginOutput, Error> {
|
||||
if !self.enabled {
|
||||
return Ok(PluginOutput::Allow);
|
||||
}
|
||||
|
||||
let query = ast
|
||||
.iter()
|
||||
.map(|q| q.to_string())
|
||||
.collect::<Vec<String>>()
|
||||
.join("; ");
|
||||
info!("[pool: {}][user: {}] {}", self.user, self.db, query);
|
||||
|
||||
Ok(PluginOutput::Allow)
|
||||
}
|
||||
}
|
||||
@@ -1,59 +0,0 @@
|
||||
//! This query router plugin will check if the user can access a particular
|
||||
//! table as part of their query. If they can't, the query will not be routed.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use sqlparser::ast::{visit_relations, Statement};
|
||||
|
||||
use crate::{
|
||||
errors::Error,
|
||||
plugins::{Plugin, PluginOutput},
|
||||
query_router::QueryRouter,
|
||||
};
|
||||
|
||||
use log::debug;
|
||||
|
||||
use core::ops::ControlFlow;
|
||||
|
||||
pub struct TableAccess<'a> {
|
||||
pub enabled: bool,
|
||||
pub tables: &'a Vec<String>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<'a> Plugin for TableAccess<'a> {
|
||||
async fn run(
|
||||
&mut self,
|
||||
_query_router: &QueryRouter,
|
||||
ast: &Vec<Statement>,
|
||||
) -> Result<PluginOutput, Error> {
|
||||
if !self.enabled {
|
||||
return Ok(PluginOutput::Allow);
|
||||
}
|
||||
|
||||
let mut found = None;
|
||||
|
||||
visit_relations(ast, |relation| {
|
||||
let relation = relation.to_string();
|
||||
let parts = relation.split(".").collect::<Vec<&str>>();
|
||||
let table_name = parts.last().unwrap();
|
||||
|
||||
if self.tables.contains(&table_name.to_string()) {
|
||||
found = Some(table_name.to_string());
|
||||
ControlFlow::<()>::Break(())
|
||||
} else {
|
||||
ControlFlow::<()>::Continue(())
|
||||
}
|
||||
});
|
||||
|
||||
if let Some(found) = found {
|
||||
debug!("Blocking access to table \"{}\"", found);
|
||||
|
||||
Ok(PluginOutput::Deny(format!(
|
||||
"permission for table \"{}\" denied",
|
||||
found
|
||||
)))
|
||||
} else {
|
||||
Ok(PluginOutput::Allow)
|
||||
}
|
||||
}
|
||||
}
|
||||
149
src/pool.rs
149
src/pool.rs
@@ -1,6 +1,6 @@
|
||||
use arc_swap::ArcSwap;
|
||||
use async_trait::async_trait;
|
||||
use bb8::{ManageConnection, Pool, PooledConnection, QueueStrategy};
|
||||
use bb8::{ManageConnection, Pool, PooledConnection};
|
||||
use bytes::{BufMut, BytesMut};
|
||||
use chrono::naive::NaiveDateTime;
|
||||
use log::{debug, error, info, warn};
|
||||
@@ -10,7 +10,6 @@ use rand::seq::SliceRandom;
|
||||
use rand::thread_rng;
|
||||
use regex::Regex;
|
||||
use std::collections::HashMap;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
@@ -18,16 +17,13 @@ use std::sync::{
|
||||
use std::time::Instant;
|
||||
use tokio::sync::Notify;
|
||||
|
||||
use crate::config::{
|
||||
get_config, Address, General, LoadBalancingMode, Plugins, PoolMode, Role, User,
|
||||
};
|
||||
use crate::config::{get_config, Address, General, LoadBalancingMode, PoolMode, Role, User};
|
||||
use crate::errors::Error;
|
||||
|
||||
use crate::auth_passthrough::AuthPassthrough;
|
||||
use crate::plugins::prewarmer;
|
||||
use crate::server::Server;
|
||||
use crate::sharding::ShardingFunction;
|
||||
use crate::stats::{AddressStats, ClientStats, ServerStats};
|
||||
use crate::stats::{AddressStats, ClientStats, PoolStats, ServerStats};
|
||||
|
||||
pub type ProcessId = i32;
|
||||
pub type SecretKey = i32;
|
||||
@@ -65,8 +61,6 @@ pub struct PoolIdentifier {
|
||||
pub user: String,
|
||||
}
|
||||
|
||||
static POOL_REAPER_RATE: u64 = 30_000; // 30 seconds by default
|
||||
|
||||
impl PoolIdentifier {
|
||||
/// Create a new user/pool identifier.
|
||||
pub fn new(db: &str, user: &str) -> PoolIdentifier {
|
||||
@@ -77,12 +71,6 @@ impl PoolIdentifier {
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for PoolIdentifier {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}@{}", self.user, self.db)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<&Address> for PoolIdentifier {
|
||||
fn from(address: &Address) -> PoolIdentifier {
|
||||
PoolIdentifier::new(&address.database, &address.username)
|
||||
@@ -103,7 +91,6 @@ pub struct PoolSettings {
|
||||
|
||||
// Connecting user.
|
||||
pub user: User,
|
||||
pub db: String,
|
||||
|
||||
// Default server role to connect to.
|
||||
pub default_role: Option<Role>,
|
||||
@@ -142,9 +129,6 @@ pub struct PoolSettings {
|
||||
pub auth_query: Option<String>,
|
||||
pub auth_query_user: Option<String>,
|
||||
pub auth_query_password: Option<String>,
|
||||
|
||||
/// Plugins
|
||||
pub plugins: Option<Plugins>,
|
||||
}
|
||||
|
||||
impl Default for PoolSettings {
|
||||
@@ -154,7 +138,6 @@ impl Default for PoolSettings {
|
||||
load_balancing_mode: LoadBalancingMode::Random,
|
||||
shards: 1,
|
||||
user: User::default(),
|
||||
db: String::default(),
|
||||
default_role: None,
|
||||
query_parser_enabled: false,
|
||||
primary_reads_enabled: true,
|
||||
@@ -169,7 +152,6 @@ impl Default for PoolSettings {
|
||||
auth_query: None,
|
||||
auth_query_user: None,
|
||||
auth_query_password: None,
|
||||
plugins: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -209,6 +191,8 @@ pub struct ConnectionPool {
|
||||
paused: Arc<AtomicBool>,
|
||||
paused_waiter: Arc<Notify>,
|
||||
|
||||
pub stats: Arc<PoolStats>,
|
||||
|
||||
/// AuthInfo
|
||||
pub auth_hash: Arc<RwLock<Option<String>>>,
|
||||
}
|
||||
@@ -258,6 +242,10 @@ impl ConnectionPool {
|
||||
.clone()
|
||||
.into_keys()
|
||||
.collect::<Vec<String>>();
|
||||
let pool_stats = Arc::new(PoolStats::new(identifier, pool_config.clone()));
|
||||
|
||||
// Allow the pool to be seen in statistics
|
||||
pool_stats.register(pool_stats.clone());
|
||||
|
||||
// Sort by shard number to ensure consistency.
|
||||
shard_ids.sort_by_key(|k| k.parse::<i64>().unwrap());
|
||||
@@ -358,12 +346,8 @@ impl ConnectionPool {
|
||||
user.clone(),
|
||||
&shard.database,
|
||||
client_server_map.clone(),
|
||||
pool_stats.clone(),
|
||||
pool_auth_hash.clone(),
|
||||
match pool_config.plugins {
|
||||
Some(ref plugins) => Some(plugins.clone()),
|
||||
None => config.plugins.clone(),
|
||||
},
|
||||
pool_config.cleanup_server_connections,
|
||||
);
|
||||
|
||||
let connect_timeout = match pool_config.connect_timeout {
|
||||
@@ -384,36 +368,15 @@ impl ConnectionPool {
|
||||
},
|
||||
};
|
||||
|
||||
let reaper_rate = *vec![idle_timeout, server_lifetime, POOL_REAPER_RATE]
|
||||
.iter()
|
||||
.min()
|
||||
.unwrap();
|
||||
|
||||
let queue_strategy = match config.general.server_round_robin {
|
||||
true => QueueStrategy::Fifo,
|
||||
false => QueueStrategy::Lifo,
|
||||
};
|
||||
|
||||
debug!(
|
||||
"[pool: {}][user: {}] Pool reaper rate: {}ms",
|
||||
pool_name, user.username, reaper_rate
|
||||
);
|
||||
|
||||
let pool = Pool::builder()
|
||||
.max_size(user.pool_size)
|
||||
.min_idle(user.min_pool_size)
|
||||
.connection_timeout(std::time::Duration::from_millis(connect_timeout))
|
||||
.idle_timeout(Some(std::time::Duration::from_millis(idle_timeout)))
|
||||
.max_lifetime(Some(std::time::Duration::from_millis(server_lifetime)))
|
||||
.reaper_rate(std::time::Duration::from_millis(reaper_rate))
|
||||
.queue_strategy(queue_strategy)
|
||||
.test_on_check_out(false);
|
||||
|
||||
let pool = if config.general.validate_config {
|
||||
pool.build(manager).await?
|
||||
} else {
|
||||
pool.build_unchecked(manager)
|
||||
};
|
||||
.test_on_check_out(false)
|
||||
.build(manager)
|
||||
.await?;
|
||||
|
||||
pools.push(pool);
|
||||
servers.push(address);
|
||||
@@ -434,6 +397,7 @@ impl ConnectionPool {
|
||||
|
||||
let pool = ConnectionPool {
|
||||
databases: shards,
|
||||
stats: pool_stats,
|
||||
addresses,
|
||||
banlist: Arc::new(RwLock::new(banlist)),
|
||||
config_hash: new_pool_hash_value,
|
||||
@@ -448,7 +412,6 @@ impl ConnectionPool {
|
||||
// shards: pool_config.shards.clone(),
|
||||
shards: shard_ids.len(),
|
||||
user: user.clone(),
|
||||
db: pool_name.clone(),
|
||||
default_role: match pool_config.default_role.as_str() {
|
||||
"any" => None,
|
||||
"replica" => Some(Role::Replica),
|
||||
@@ -474,10 +437,6 @@ impl ConnectionPool {
|
||||
auth_query: pool_config.auth_query.clone(),
|
||||
auth_query_user: pool_config.auth_query_user.clone(),
|
||||
auth_query_password: pool_config.auth_query_password.clone(),
|
||||
plugins: match pool_config.plugins {
|
||||
Some(ref plugins) => Some(plugins.clone()),
|
||||
None => config.plugins.clone(),
|
||||
},
|
||||
},
|
||||
validated: Arc::new(AtomicBool::new(false)),
|
||||
paused: Arc::new(AtomicBool::new(false)),
|
||||
@@ -487,12 +446,10 @@ impl ConnectionPool {
|
||||
// Connect to the servers to make sure pool configuration is valid
|
||||
// before setting it globally.
|
||||
// Do this async and somewhere else, we don't have to wait here.
|
||||
if config.general.validate_config {
|
||||
let mut validate_pool = pool.clone();
|
||||
tokio::task::spawn(async move {
|
||||
let _ = validate_pool.validate().await;
|
||||
});
|
||||
}
|
||||
let mut validate_pool = pool.clone();
|
||||
tokio::task::spawn(async move {
|
||||
let _ = validate_pool.validate().await;
|
||||
});
|
||||
|
||||
// There is one pool per database/user pair.
|
||||
new_pools.insert(PoolIdentifier::new(pool_name, &user.username), pool);
|
||||
@@ -614,10 +571,6 @@ impl ConnectionPool {
|
||||
});
|
||||
}
|
||||
|
||||
// Indicate we're waiting on a server connection from a pool.
|
||||
let now = Instant::now();
|
||||
client_stats.waiting();
|
||||
|
||||
while !candidates.is_empty() {
|
||||
// Get the next candidate
|
||||
let address = match candidates.pop() {
|
||||
@@ -636,6 +589,10 @@ impl ConnectionPool {
|
||||
}
|
||||
}
|
||||
|
||||
// Indicate we're waiting on a server connection from a pool.
|
||||
let now = Instant::now();
|
||||
client_stats.waiting();
|
||||
|
||||
// Check if we can connect
|
||||
let mut conn = match self.databases[address.shard][address.address_index]
|
||||
.get()
|
||||
@@ -643,10 +600,7 @@ impl ConnectionPool {
|
||||
{
|
||||
Ok(conn) => conn,
|
||||
Err(err) => {
|
||||
error!(
|
||||
"Connection checkout error for instance {:?}, error: {:?}",
|
||||
address, err
|
||||
);
|
||||
error!("Banning instance {:?}, error: {:?}", address, err);
|
||||
self.ban(address, BanReason::FailedCheckout, Some(client_stats));
|
||||
address.stats.error();
|
||||
client_stats.idle();
|
||||
@@ -673,7 +627,7 @@ impl ConnectionPool {
|
||||
.stats()
|
||||
.checkout_time(checkout_time, client_stats.application_name());
|
||||
server.stats().active(client_stats.application_name());
|
||||
client_stats.active();
|
||||
|
||||
return Ok((conn, address.clone()));
|
||||
}
|
||||
|
||||
@@ -681,19 +635,11 @@ impl ConnectionPool {
|
||||
.run_health_check(address, server, now, client_stats)
|
||||
.await
|
||||
{
|
||||
let checkout_time: u64 = now.elapsed().as_micros() as u64;
|
||||
client_stats.checkout_time(checkout_time);
|
||||
server
|
||||
.stats()
|
||||
.checkout_time(checkout_time, client_stats.application_name());
|
||||
server.stats().active(client_stats.application_name());
|
||||
client_stats.active();
|
||||
return Ok((conn, address.clone()));
|
||||
} else {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
client_stats.idle();
|
||||
Err(Error::AllServersDown)
|
||||
}
|
||||
|
||||
@@ -730,7 +676,7 @@ impl ConnectionPool {
|
||||
// Health check failed.
|
||||
Err(err) => {
|
||||
error!(
|
||||
"Failed health check on instance {:?}, error: {:?}",
|
||||
"Banning instance {:?} because of failed health check, {:?}",
|
||||
address, err
|
||||
);
|
||||
}
|
||||
@@ -739,7 +685,7 @@ impl ConnectionPool {
|
||||
// Health check timed out.
|
||||
Err(err) => {
|
||||
error!(
|
||||
"Health check timeout on instance {:?}, error: {:?}",
|
||||
"Banning instance {:?} because of health check timeout, {:?}",
|
||||
address, err
|
||||
);
|
||||
}
|
||||
@@ -761,16 +707,13 @@ impl ConnectionPool {
|
||||
return;
|
||||
}
|
||||
|
||||
error!("Banning instance {:?}, reason: {:?}", address, reason);
|
||||
|
||||
let now = chrono::offset::Utc::now().naive_utc();
|
||||
let mut guard = self.banlist.write();
|
||||
|
||||
error!("Banning {:?}", address);
|
||||
if let Some(client_info) = client_info {
|
||||
client_info.ban_error();
|
||||
address.stats.error();
|
||||
}
|
||||
|
||||
guard[address.shard].insert(address.clone(), (reason, now));
|
||||
}
|
||||
|
||||
@@ -927,26 +870,12 @@ impl ConnectionPool {
|
||||
|
||||
/// Wrapper for the bb8 connection pool.
|
||||
pub struct ServerPool {
|
||||
/// Server address.
|
||||
address: Address,
|
||||
|
||||
/// Server Postgres user.
|
||||
user: User,
|
||||
|
||||
/// Server database.
|
||||
database: String,
|
||||
|
||||
/// Client/server mapping.
|
||||
client_server_map: ClientServerMap,
|
||||
|
||||
/// Server auth hash (for auth passthrough).
|
||||
stats: Arc<PoolStats>,
|
||||
auth_hash: Arc<RwLock<Option<String>>>,
|
||||
|
||||
/// Server plugins.
|
||||
plugins: Option<Plugins>,
|
||||
|
||||
/// Should we clean up dirty connections before putting them into the pool?
|
||||
cleanup_connections: bool,
|
||||
}
|
||||
|
||||
impl ServerPool {
|
||||
@@ -955,18 +884,16 @@ impl ServerPool {
|
||||
user: User,
|
||||
database: &str,
|
||||
client_server_map: ClientServerMap,
|
||||
stats: Arc<PoolStats>,
|
||||
auth_hash: Arc<RwLock<Option<String>>>,
|
||||
plugins: Option<Plugins>,
|
||||
cleanup_connections: bool,
|
||||
) -> ServerPool {
|
||||
ServerPool {
|
||||
address,
|
||||
user: user.clone(),
|
||||
database: database.to_string(),
|
||||
client_server_map,
|
||||
stats,
|
||||
auth_hash,
|
||||
plugins,
|
||||
cleanup_connections,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -982,6 +909,7 @@ impl ManageConnection for ServerPool {
|
||||
|
||||
let stats = Arc::new(ServerStats::new(
|
||||
self.address.clone(),
|
||||
self.stats.clone(),
|
||||
tokio::time::Instant::now(),
|
||||
));
|
||||
|
||||
@@ -995,23 +923,10 @@ impl ManageConnection for ServerPool {
|
||||
self.client_server_map.clone(),
|
||||
stats.clone(),
|
||||
self.auth_hash.clone(),
|
||||
self.cleanup_connections,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(mut conn) => {
|
||||
if let Some(ref plugins) = self.plugins {
|
||||
if let Some(ref prewarmer) = plugins.prewarmer {
|
||||
let mut prewarmer = prewarmer::Prewarmer {
|
||||
enabled: prewarmer.enabled,
|
||||
server: &mut conn,
|
||||
queries: &prewarmer.queries,
|
||||
};
|
||||
|
||||
prewarmer.run().await?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(conn) => {
|
||||
stats.idle();
|
||||
Ok(conn)
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
use hyper::service::{make_service_fn, service_fn};
|
||||
use hyper::{Body, Method, Request, Response, Server, StatusCode};
|
||||
use log::{debug, error, info};
|
||||
use log::{error, info, warn};
|
||||
use phf::phf_map;
|
||||
use std::collections::HashMap;
|
||||
use std::fmt;
|
||||
@@ -9,9 +9,8 @@ use std::sync::atomic::Ordering;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::config::Address;
|
||||
use crate::pool::{get_all_pools, PoolIdentifier};
|
||||
use crate::stats::pool::PoolStats;
|
||||
use crate::stats::{get_server_stats, ServerStats};
|
||||
use crate::pool::get_all_pools;
|
||||
use crate::stats::{get_pool_stats, get_server_stats, ServerStats};
|
||||
|
||||
struct MetricHelpType {
|
||||
help: &'static str,
|
||||
@@ -234,10 +233,10 @@ impl<Value: fmt::Display> PrometheusMetric<Value> {
|
||||
Self::from_name(&format!("stats_{}", name), value, labels)
|
||||
}
|
||||
|
||||
fn from_pool(pool_id: PoolIdentifier, name: &str, value: u64) -> Option<PrometheusMetric<u64>> {
|
||||
fn from_pool(pool: &(String, String), name: &str, value: u64) -> Option<PrometheusMetric<u64>> {
|
||||
let mut labels = HashMap::new();
|
||||
labels.insert("pool", pool_id.db);
|
||||
labels.insert("user", pool_id.user);
|
||||
labels.insert("pool", pool.0.clone());
|
||||
labels.insert("user", pool.1.clone());
|
||||
|
||||
Self::from_name(&format!("pools_{}", name), value, labels)
|
||||
}
|
||||
@@ -275,7 +274,7 @@ fn push_address_stats(lines: &mut Vec<String>) {
|
||||
{
|
||||
lines.push(prometheus_metric.to_string());
|
||||
} else {
|
||||
debug!("Metric {} not implemented for {}", key, address.name());
|
||||
warn!("Metric {} not implemented for {}", key, address.name());
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -285,15 +284,18 @@ fn push_address_stats(lines: &mut Vec<String>) {
|
||||
|
||||
// Adds relevant metrics shown in a SHOW POOLS admin command.
|
||||
fn push_pool_stats(lines: &mut Vec<String>) {
|
||||
let pool_stats = PoolStats::construct_pool_lookup();
|
||||
for (pool_id, stats) in pool_stats.iter() {
|
||||
let pool_stats = get_pool_stats();
|
||||
for (pool, stats) in pool_stats.iter() {
|
||||
let stats = &**stats;
|
||||
for (name, value) in stats.clone() {
|
||||
if let Some(prometheus_metric) =
|
||||
PrometheusMetric::<u64>::from_pool(pool_id.clone(), &name, value)
|
||||
if let Some(prometheus_metric) = PrometheusMetric::<u64>::from_pool(pool, &name, value)
|
||||
{
|
||||
lines.push(prometheus_metric.to_string());
|
||||
} else {
|
||||
debug!("Metric {} not implemented for ({})", name, *pool_id);
|
||||
warn!(
|
||||
"Metric {} not implemented for ({},{})",
|
||||
name, pool.0, pool.1
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -318,7 +320,7 @@ fn push_database_stats(lines: &mut Vec<String>) {
|
||||
{
|
||||
lines.push(prometheus_metric.to_string());
|
||||
} else {
|
||||
debug!("Metric {} not implemented for {}", key, address.name());
|
||||
warn!("Metric {} not implemented for {}", key, address.name());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,16 +6,13 @@ use once_cell::sync::OnceCell;
|
||||
use regex::{Regex, RegexSet};
|
||||
use sqlparser::ast::Statement::{Query, StartTransaction};
|
||||
use sqlparser::ast::{
|
||||
BinaryOperator, Expr, Ident, JoinConstraint, JoinOperator, SetExpr, Statement, TableFactor,
|
||||
Value,
|
||||
BinaryOperator, Expr, Ident, JoinConstraint, JoinOperator, SetExpr, TableFactor, Value,
|
||||
};
|
||||
use sqlparser::dialect::PostgreSqlDialect;
|
||||
use sqlparser::parser::Parser;
|
||||
|
||||
use crate::config::Role;
|
||||
use crate::errors::Error;
|
||||
use crate::messages::BytesMutReader;
|
||||
use crate::plugins::{Intercept, Plugin, PluginOutput, QueryLogger, TableAccess};
|
||||
use crate::pool::PoolSettings;
|
||||
use crate::sharding::Sharder;
|
||||
|
||||
@@ -132,10 +129,6 @@ impl QueryRouter {
|
||||
self.pool_settings = pool_settings;
|
||||
}
|
||||
|
||||
pub fn pool_settings<'a>(&'a self) -> &'a PoolSettings {
|
||||
&self.pool_settings
|
||||
}
|
||||
|
||||
/// Try to parse a command and execute it.
|
||||
pub fn try_execute_command(&mut self, message_buffer: &BytesMut) -> Option<(Command, String)> {
|
||||
let mut message_cursor = Cursor::new(message_buffer);
|
||||
@@ -331,7 +324,10 @@ impl QueryRouter {
|
||||
Some((command, value))
|
||||
}
|
||||
|
||||
pub fn parse(message: &BytesMut) -> Result<Vec<Statement>, Error> {
|
||||
/// Try to infer which server to connect to based on the contents of the query.
|
||||
pub fn infer(&mut self, message: &BytesMut) -> bool {
|
||||
debug!("Inferring role");
|
||||
|
||||
let mut message_cursor = Cursor::new(message);
|
||||
|
||||
let code = message_cursor.get_u8() as char;
|
||||
@@ -348,39 +344,37 @@ impl QueryRouter {
|
||||
// Parse (prepared statement)
|
||||
'P' => {
|
||||
// Reads statement name
|
||||
let _name = message_cursor.read_string().unwrap();
|
||||
message_cursor.read_string().unwrap();
|
||||
|
||||
// Reads query string
|
||||
let query = message_cursor.read_string().unwrap();
|
||||
|
||||
debug!("Prepared statement: '{}'", query);
|
||||
|
||||
query
|
||||
}
|
||||
|
||||
_ => return Err(Error::UnsupportedStatement),
|
||||
_ => return false,
|
||||
};
|
||||
|
||||
match Parser::parse_sql(&PostgreSqlDialect {}, &query) {
|
||||
Ok(ast) => Ok(ast),
|
||||
let ast = match Parser::parse_sql(&PostgreSqlDialect {}, &query) {
|
||||
Ok(ast) => ast,
|
||||
Err(err) => {
|
||||
// SELECT ... FOR UPDATE won't get parsed correctly.
|
||||
debug!("{}: {}", err, query);
|
||||
Err(Error::QueryRouterParserError(err.to_string()))
|
||||
self.active_role = Some(Role::Primary);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/// Try to infer which server to connect to based on the contents of the query.
|
||||
pub fn infer(&mut self, ast: &Vec<sqlparser::ast::Statement>) -> Result<(), Error> {
|
||||
debug!("Inferring role");
|
||||
debug!("AST: {:?}", ast);
|
||||
|
||||
if ast.is_empty() {
|
||||
// That's weird, no idea, let's go to primary
|
||||
self.active_role = Some(Role::Primary);
|
||||
return Err(Error::QueryRouterParserError("empty query".into()));
|
||||
return false;
|
||||
}
|
||||
|
||||
for q in ast {
|
||||
for q in &ast {
|
||||
match q {
|
||||
// All transactions go to the primary, probably a write.
|
||||
StartTransaction { .. } => {
|
||||
@@ -424,7 +418,7 @@ impl QueryRouter {
|
||||
};
|
||||
}
|
||||
|
||||
Ok(())
|
||||
true
|
||||
}
|
||||
|
||||
/// Parse the shard number from the Bind message
|
||||
@@ -789,52 +783,6 @@ impl QueryRouter {
|
||||
}
|
||||
}
|
||||
|
||||
/// Add your plugins here and execute them.
|
||||
pub async fn execute_plugins(&self, ast: &Vec<Statement>) -> Result<PluginOutput, Error> {
|
||||
let plugins = match self.pool_settings.plugins {
|
||||
Some(ref plugins) => plugins,
|
||||
None => return Ok(PluginOutput::Allow),
|
||||
};
|
||||
|
||||
if let Some(ref query_logger) = plugins.query_logger {
|
||||
let mut query_logger = QueryLogger {
|
||||
enabled: query_logger.enabled,
|
||||
user: &self.pool_settings.user.username,
|
||||
db: &self.pool_settings.db,
|
||||
};
|
||||
|
||||
let _ = query_logger.run(&self, ast).await;
|
||||
}
|
||||
|
||||
if let Some(ref intercept) = plugins.intercept {
|
||||
let mut intercept = Intercept {
|
||||
enabled: intercept.enabled,
|
||||
config: &intercept,
|
||||
};
|
||||
|
||||
let result = intercept.run(&self, ast).await;
|
||||
|
||||
if let Ok(PluginOutput::Intercept(output)) = result {
|
||||
return Ok(PluginOutput::Intercept(output));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref table_access) = plugins.table_access {
|
||||
let mut table_access = TableAccess {
|
||||
enabled: table_access.enabled,
|
||||
tables: &table_access.tables,
|
||||
};
|
||||
|
||||
let result = table_access.run(&self, ast).await;
|
||||
|
||||
if let Ok(PluginOutput::Deny(error)) = result {
|
||||
return Ok(PluginOutput::Deny(error));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(PluginOutput::Allow)
|
||||
}
|
||||
|
||||
fn set_sharding_key(&mut self, sharding_key: i64) -> Option<usize> {
|
||||
let sharder = Sharder::new(
|
||||
self.pool_settings.shards,
|
||||
@@ -862,23 +810,12 @@ impl QueryRouter {
|
||||
/// Should we attempt to parse queries?
|
||||
pub fn query_parser_enabled(&self) -> bool {
|
||||
let enabled = match self.query_parser_enabled {
|
||||
None => {
|
||||
debug!(
|
||||
"Using pool settings, query_parser_enabled: {}",
|
||||
self.pool_settings.query_parser_enabled
|
||||
);
|
||||
self.pool_settings.query_parser_enabled
|
||||
}
|
||||
|
||||
Some(value) => {
|
||||
debug!(
|
||||
"Using query parser override, query_parser_enabled: {}",
|
||||
value
|
||||
);
|
||||
value
|
||||
}
|
||||
None => self.pool_settings.query_parser_enabled,
|
||||
Some(value) => value,
|
||||
};
|
||||
|
||||
debug!("Query parser enabled: {}", enabled);
|
||||
|
||||
enabled
|
||||
}
|
||||
|
||||
@@ -925,7 +862,7 @@ mod test {
|
||||
|
||||
for query in queries {
|
||||
// It's a recognized query
|
||||
assert!(qr.infer(&QueryRouter::parse(&query).unwrap()).is_ok());
|
||||
assert!(qr.infer(&query));
|
||||
assert_eq!(qr.role(), Some(Role::Replica));
|
||||
}
|
||||
}
|
||||
@@ -944,7 +881,7 @@ mod test {
|
||||
|
||||
for query in queries {
|
||||
// It's a recognized query
|
||||
assert!(qr.infer(&QueryRouter::parse(&query).unwrap()).is_ok());
|
||||
assert!(qr.infer(&query));
|
||||
assert_eq!(qr.role(), Some(Role::Primary));
|
||||
}
|
||||
}
|
||||
@@ -956,7 +893,7 @@ mod test {
|
||||
let query = simple_query("SELECT * FROM items WHERE id = 5");
|
||||
assert!(qr.try_execute_command(&simple_query("SET PRIMARY READS TO on")) != None);
|
||||
|
||||
assert!(qr.infer(&QueryRouter::parse(&query).unwrap()).is_ok());
|
||||
assert!(qr.infer(&query));
|
||||
assert_eq!(qr.role(), None);
|
||||
}
|
||||
|
||||
@@ -976,7 +913,7 @@ mod test {
|
||||
res.put(prepared_stmt);
|
||||
res.put_i16(0);
|
||||
|
||||
assert!(qr.infer(&QueryRouter::parse(&res).unwrap()).is_ok());
|
||||
assert!(qr.infer(&res));
|
||||
assert_eq!(qr.role(), Some(Role::Replica));
|
||||
}
|
||||
|
||||
@@ -1140,11 +1077,11 @@ mod test {
|
||||
assert_eq!(qr.role(), None);
|
||||
|
||||
let query = simple_query("INSERT INTO test_table VALUES (1)");
|
||||
assert!(qr.infer(&QueryRouter::parse(&query).unwrap()).is_ok());
|
||||
assert!(qr.infer(&query));
|
||||
assert_eq!(qr.role(), Some(Role::Primary));
|
||||
|
||||
let query = simple_query("SELECT * FROM test_table");
|
||||
assert!(qr.infer(&QueryRouter::parse(&query).unwrap()).is_ok());
|
||||
assert!(qr.infer(&query));
|
||||
assert_eq!(qr.role(), Some(Role::Replica));
|
||||
|
||||
assert!(qr.query_parser_enabled());
|
||||
@@ -1176,8 +1113,6 @@ mod test {
|
||||
auth_query: None,
|
||||
auth_query_password: None,
|
||||
auth_query_user: None,
|
||||
db: "test".to_string(),
|
||||
plugins: None,
|
||||
};
|
||||
let mut qr = QueryRouter::new();
|
||||
assert_eq!(qr.active_role, None);
|
||||
@@ -1207,24 +1142,15 @@ mod test {
|
||||
QueryRouter::setup();
|
||||
|
||||
let mut qr = QueryRouter::new();
|
||||
assert!(qr
|
||||
.infer(&QueryRouter::parse(&simple_query("BEGIN; SELECT 1; COMMIT;")).unwrap())
|
||||
.is_ok());
|
||||
assert!(qr.infer(&simple_query("BEGIN; SELECT 1; COMMIT;")));
|
||||
assert_eq!(qr.role(), Role::Primary);
|
||||
|
||||
assert!(qr
|
||||
.infer(&QueryRouter::parse(&simple_query("SELECT 1; SELECT 2;")).unwrap())
|
||||
.is_ok());
|
||||
assert!(qr.infer(&simple_query("SELECT 1; SELECT 2;")));
|
||||
assert_eq!(qr.role(), Role::Replica);
|
||||
|
||||
assert!(qr
|
||||
.infer(
|
||||
&QueryRouter::parse(&simple_query(
|
||||
"SELECT 123; INSERT INTO t VALUES (5); SELECT 1;"
|
||||
))
|
||||
.unwrap()
|
||||
)
|
||||
.is_ok());
|
||||
assert!(qr.infer(&simple_query(
|
||||
"SELECT 123; INSERT INTO t VALUES (5); SELECT 1;"
|
||||
)));
|
||||
assert_eq!(qr.role(), Role::Primary);
|
||||
}
|
||||
|
||||
@@ -1251,10 +1177,7 @@ mod test {
|
||||
auth_query: None,
|
||||
auth_query_password: None,
|
||||
auth_query_user: None,
|
||||
db: "test".to_string(),
|
||||
plugins: None,
|
||||
};
|
||||
|
||||
let mut qr = QueryRouter::new();
|
||||
qr.update_pool_settings(pool_settings.clone());
|
||||
|
||||
@@ -1285,84 +1208,47 @@ mod test {
|
||||
qr.pool_settings.automatic_sharding_key = Some("data.id".to_string());
|
||||
qr.pool_settings.shards = 3;
|
||||
|
||||
assert!(qr
|
||||
.infer(&QueryRouter::parse(&simple_query("SELECT * FROM data WHERE id = 5")).unwrap())
|
||||
.is_ok());
|
||||
assert!(qr.infer(&simple_query("SELECT * FROM data WHERE id = 5")));
|
||||
assert_eq!(qr.shard(), 2);
|
||||
|
||||
assert!(qr
|
||||
.infer(
|
||||
&QueryRouter::parse(&simple_query(
|
||||
"SELECT one, two, three FROM public.data WHERE id = 6"
|
||||
))
|
||||
.unwrap()
|
||||
)
|
||||
.is_ok());
|
||||
assert!(qr.infer(&simple_query(
|
||||
"SELECT one, two, three FROM public.data WHERE id = 6"
|
||||
)));
|
||||
assert_eq!(qr.shard(), 0);
|
||||
|
||||
assert!(qr
|
||||
.infer(
|
||||
&QueryRouter::parse(&simple_query(
|
||||
"SELECT * FROM data
|
||||
assert!(qr.infer(&simple_query(
|
||||
"SELECT * FROM data
|
||||
INNER JOIN t2 ON data.id = 5
|
||||
AND t2.data_id = data.id
|
||||
WHERE data.id = 5"
|
||||
))
|
||||
.unwrap()
|
||||
)
|
||||
.is_ok());
|
||||
)));
|
||||
assert_eq!(qr.shard(), 2);
|
||||
|
||||
// Shard did not move because we couldn't determine the sharding key since it could be ambiguous
|
||||
// in the query.
|
||||
assert!(qr
|
||||
.infer(
|
||||
&QueryRouter::parse(&simple_query(
|
||||
"SELECT * FROM t2 INNER JOIN data ON id = 6 AND data.id = t2.data_id"
|
||||
))
|
||||
.unwrap()
|
||||
)
|
||||
.is_ok());
|
||||
assert!(qr.infer(&simple_query(
|
||||
"SELECT * FROM t2 INNER JOIN data ON id = 6 AND data.id = t2.data_id"
|
||||
)));
|
||||
assert_eq!(qr.shard(), 2);
|
||||
|
||||
assert!(qr
|
||||
.infer(
|
||||
&QueryRouter::parse(&simple_query(
|
||||
r#"SELECT * FROM "public"."data" WHERE "id" = 6"#
|
||||
))
|
||||
.unwrap()
|
||||
)
|
||||
.is_ok());
|
||||
assert!(qr.infer(&simple_query(
|
||||
r#"SELECT * FROM "public"."data" WHERE "id" = 6"#
|
||||
)));
|
||||
assert_eq!(qr.shard(), 0);
|
||||
|
||||
assert!(qr
|
||||
.infer(
|
||||
&QueryRouter::parse(&simple_query(
|
||||
r#"SELECT * FROM "public"."data" WHERE "data"."id" = 5"#
|
||||
))
|
||||
.unwrap()
|
||||
)
|
||||
.is_ok());
|
||||
assert!(qr.infer(&simple_query(
|
||||
r#"SELECT * FROM "public"."data" WHERE "data"."id" = 5"#
|
||||
)));
|
||||
assert_eq!(qr.shard(), 2);
|
||||
|
||||
// Super unique sharding key
|
||||
qr.pool_settings.automatic_sharding_key = Some("*.unique_enough_column_name".to_string());
|
||||
assert!(qr
|
||||
.infer(
|
||||
&QueryRouter::parse(&simple_query(
|
||||
"SELECT * FROM table_x WHERE unique_enough_column_name = 6"
|
||||
))
|
||||
.unwrap()
|
||||
)
|
||||
.is_ok());
|
||||
assert!(qr.infer(&simple_query(
|
||||
"SELECT * FROM table_x WHERE unique_enough_column_name = 6"
|
||||
)));
|
||||
assert_eq!(qr.shard(), 0);
|
||||
|
||||
assert!(qr
|
||||
.infer(
|
||||
&QueryRouter::parse(&simple_query("SELECT * FROM table_y WHERE another_key = 5"))
|
||||
.unwrap()
|
||||
)
|
||||
.is_ok());
|
||||
assert!(qr.infer(&simple_query("SELECT * FROM table_y WHERE another_key = 5")));
|
||||
assert_eq!(qr.shard(), 0);
|
||||
}
|
||||
|
||||
@@ -1386,61 +1272,11 @@ mod test {
|
||||
qr.pool_settings.automatic_sharding_key = Some("data.id".to_string());
|
||||
qr.pool_settings.shards = 3;
|
||||
|
||||
assert!(qr
|
||||
.infer(&QueryRouter::parse(&simple_query(stmt)).unwrap())
|
||||
.is_ok());
|
||||
assert!(qr.infer(&simple_query(stmt)));
|
||||
assert_eq!(qr.placeholders.len(), 1);
|
||||
|
||||
assert!(qr.infer_shard_from_bind(&bind));
|
||||
assert_eq!(qr.shard(), 2);
|
||||
assert!(qr.placeholders.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_table_access_plugin() {
|
||||
use crate::config::{Plugins, TableAccess};
|
||||
let table_access = TableAccess {
|
||||
enabled: true,
|
||||
tables: vec![String::from("pg_database")],
|
||||
};
|
||||
let plugins = Plugins {
|
||||
table_access: Some(table_access),
|
||||
intercept: None,
|
||||
query_logger: None,
|
||||
prewarmer: None,
|
||||
};
|
||||
|
||||
QueryRouter::setup();
|
||||
let mut pool_settings = PoolSettings::default();
|
||||
pool_settings.query_parser_enabled = true;
|
||||
pool_settings.plugins = Some(plugins);
|
||||
|
||||
let mut qr = QueryRouter::new();
|
||||
qr.update_pool_settings(pool_settings);
|
||||
|
||||
let query = simple_query("SELECT * FROM pg_database");
|
||||
let ast = QueryRouter::parse(&query).unwrap();
|
||||
|
||||
let res = qr.execute_plugins(&ast).await;
|
||||
|
||||
assert_eq!(
|
||||
res,
|
||||
Ok(PluginOutput::Deny(
|
||||
"permission for table \"pg_database\" denied".to_string()
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_plugins_disabled_by_defaault() {
|
||||
QueryRouter::setup();
|
||||
let qr = QueryRouter::new();
|
||||
|
||||
let query = simple_query("SELECT * FROM pg_database");
|
||||
let ast = QueryRouter::parse(&query).unwrap();
|
||||
|
||||
let res = qr.execute_plugins(&ast).await;
|
||||
|
||||
assert_eq!(res, Ok(PluginOutput::Allow));
|
||||
}
|
||||
}
|
||||
|
||||
182
src/server.rs
182
src/server.rs
@@ -5,9 +5,8 @@ use fallible_iterator::FallibleIterator;
|
||||
use log::{debug, error, info, trace, warn};
|
||||
use parking_lot::{Mutex, RwLock};
|
||||
use postgres_protocol::message;
|
||||
use std::collections::{BTreeSet, HashMap};
|
||||
use std::collections::HashMap;
|
||||
use std::io::Read;
|
||||
use std::net::IpAddr;
|
||||
use std::sync::Arc;
|
||||
use std::time::SystemTime;
|
||||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, BufStream};
|
||||
@@ -17,7 +16,6 @@ use tokio_rustls::{client::TlsStream, TlsConnector};
|
||||
|
||||
use crate::config::{get_config, Address, User};
|
||||
use crate::constants::*;
|
||||
use crate::dns_cache::{AddrSet, CACHED_RESOLVER};
|
||||
use crate::errors::{Error, ServerIdentifier};
|
||||
use crate::messages::*;
|
||||
use crate::mirrors::MirroringManager;
|
||||
@@ -103,48 +101,6 @@ impl StreamInner {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone)]
|
||||
struct CleanupState {
|
||||
/// If server connection requires DISCARD ALL before checkin because of set statement
|
||||
needs_cleanup_set: bool,
|
||||
|
||||
/// If server connection requires DISCARD ALL before checkin because of prepare statement
|
||||
needs_cleanup_prepare: bool,
|
||||
}
|
||||
|
||||
impl CleanupState {
|
||||
fn new() -> Self {
|
||||
CleanupState {
|
||||
needs_cleanup_set: false,
|
||||
needs_cleanup_prepare: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn needs_cleanup(&self) -> bool {
|
||||
self.needs_cleanup_set || self.needs_cleanup_prepare
|
||||
}
|
||||
|
||||
fn set_true(&mut self) {
|
||||
self.needs_cleanup_set = true;
|
||||
self.needs_cleanup_prepare = true;
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
self.needs_cleanup_set = false;
|
||||
self.needs_cleanup_prepare = false;
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for CleanupState {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"SET: {}, PREPARE: {}",
|
||||
self.needs_cleanup_set, self.needs_cleanup_prepare
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Server state.
|
||||
pub struct Server {
|
||||
/// Server host, e.g. localhost,
|
||||
@@ -173,8 +129,8 @@ pub struct Server {
|
||||
/// Is the server broken? We'll remote it from the pool if so.
|
||||
bad: bool,
|
||||
|
||||
/// If server connection requires DISCARD ALL before checkin
|
||||
cleanup_state: CleanupState,
|
||||
/// If server connection requires a DISCARD ALL before checkin
|
||||
needs_cleanup: bool,
|
||||
|
||||
/// Mapping of clients and servers used for query cancellation.
|
||||
client_server_map: ClientServerMap,
|
||||
@@ -188,19 +144,10 @@ pub struct Server {
|
||||
/// Application name using the server at the moment.
|
||||
application_name: String,
|
||||
|
||||
/// Last time that a successful server send or response happened
|
||||
// Last time that a successful server send or response happened
|
||||
last_activity: SystemTime,
|
||||
|
||||
mirror_manager: Option<MirroringManager>,
|
||||
|
||||
/// Associated addresses used
|
||||
addr_set: Option<AddrSet>,
|
||||
|
||||
/// Should clean up dirty connections?
|
||||
cleanup_connections: bool,
|
||||
|
||||
/// Prepared statements
|
||||
prepared_statements: BTreeSet<String>,
|
||||
}
|
||||
|
||||
impl Server {
|
||||
@@ -213,26 +160,7 @@ impl Server {
|
||||
client_server_map: ClientServerMap,
|
||||
stats: Arc<ServerStats>,
|
||||
auth_hash: Arc<RwLock<Option<String>>>,
|
||||
cleanup_connections: bool,
|
||||
) -> Result<Server, Error> {
|
||||
let cached_resolver = CACHED_RESOLVER.load();
|
||||
let mut addr_set: Option<AddrSet> = None;
|
||||
|
||||
// If we are caching addresses and hostname is not an IP
|
||||
if cached_resolver.enabled() && address.host.parse::<IpAddr>().is_err() {
|
||||
debug!("Resolving {}", &address.host);
|
||||
addr_set = match cached_resolver.lookup_ip(&address.host).await {
|
||||
Ok(ok) => {
|
||||
debug!("Obtained: {:?}", ok);
|
||||
Some(ok)
|
||||
}
|
||||
Err(err) => {
|
||||
warn!("Error trying to resolve {}, ({:?})", &address.host, err);
|
||||
None
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let mut stream =
|
||||
match TcpStream::connect(&format!("{}:{}", &address.host, address.port)).await {
|
||||
Ok(stream) => stream,
|
||||
@@ -679,9 +607,8 @@ impl Server {
|
||||
in_transaction: false,
|
||||
data_available: false,
|
||||
bad: false,
|
||||
cleanup_state: CleanupState::new(),
|
||||
needs_cleanup: false,
|
||||
client_server_map,
|
||||
addr_set,
|
||||
connected_at: chrono::offset::Utc::now().naive_utc(),
|
||||
stats,
|
||||
application_name: String::new(),
|
||||
@@ -694,8 +621,6 @@ impl Server {
|
||||
address.mirrors.clone(),
|
||||
)),
|
||||
},
|
||||
cleanup_connections,
|
||||
prepared_statements: BTreeSet::new(),
|
||||
};
|
||||
|
||||
server.set_name("pgcat").await?;
|
||||
@@ -756,10 +681,7 @@ impl Server {
|
||||
Ok(())
|
||||
}
|
||||
Err(err) => {
|
||||
error!(
|
||||
"Terminating server {:?} because of: {:?}",
|
||||
self.address, err
|
||||
);
|
||||
error!("Terminating server because of: {:?}", err);
|
||||
self.bad = true;
|
||||
Err(err)
|
||||
}
|
||||
@@ -774,10 +696,7 @@ impl Server {
|
||||
let mut message = match read_message(&mut self.stream).await {
|
||||
Ok(message) => message,
|
||||
Err(err) => {
|
||||
error!(
|
||||
"Terminating server {:?} because of: {:?}",
|
||||
self.address, err
|
||||
);
|
||||
error!("Terminating server because of: {:?}", err);
|
||||
self.bad = true;
|
||||
return Err(err);
|
||||
}
|
||||
@@ -844,12 +763,12 @@ impl Server {
|
||||
// This will reduce amount of discard statements sent
|
||||
if !self.in_transaction {
|
||||
debug!("Server connection marked for clean up");
|
||||
self.cleanup_state.needs_cleanup_set = true;
|
||||
self.needs_cleanup = true;
|
||||
}
|
||||
}
|
||||
"PREPARE\0" => {
|
||||
debug!("Server connection marked for clean up");
|
||||
self.cleanup_state.needs_cleanup_prepare = true;
|
||||
self.needs_cleanup = true;
|
||||
}
|
||||
_ => (),
|
||||
}
|
||||
@@ -914,43 +833,6 @@ impl Server {
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
pub fn will_prepare(&mut self, name: &str) {
|
||||
debug!("Will prepare `{}`", name);
|
||||
|
||||
self.prepared_statements.insert(name.to_string());
|
||||
}
|
||||
|
||||
pub fn should_prepare(&self, name: &str) -> bool {
|
||||
let should_prepare = !self.prepared_statements.contains(name);
|
||||
|
||||
debug!("Should prepare `{}`: {}", name, should_prepare);
|
||||
|
||||
if should_prepare {
|
||||
self.stats.prepared_cache_miss();
|
||||
} else {
|
||||
self.stats.prepared_cache_hit();
|
||||
}
|
||||
|
||||
should_prepare
|
||||
}
|
||||
|
||||
pub async fn prepare(&mut self, parse: &Parse) -> Result<(), Error> {
|
||||
debug!("Preparing `{}`", parse.name);
|
||||
|
||||
let bytes: BytesMut = parse.try_into()?;
|
||||
self.send(&bytes).await?;
|
||||
self.send(&flush()).await?;
|
||||
|
||||
// Read and discard ParseComplete (B)
|
||||
let _ = read_message(&mut self.stream).await?;
|
||||
|
||||
self.prepared_statements.insert(parse.name.to_string());
|
||||
|
||||
debug!("Prepared `{}`", parse.name);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// If the server is still inside a transaction.
|
||||
/// If the client disconnects while the server is in a transaction, we will clean it up.
|
||||
pub fn in_transaction(&self) -> bool {
|
||||
@@ -967,23 +849,7 @@ impl Server {
|
||||
/// Server & client are out of sync, we must discard this connection.
|
||||
/// This happens with clients that misbehave.
|
||||
pub fn is_bad(&self) -> bool {
|
||||
if self.bad {
|
||||
return self.bad;
|
||||
};
|
||||
let cached_resolver = CACHED_RESOLVER.load();
|
||||
if cached_resolver.enabled() {
|
||||
if let Some(addr_set) = &self.addr_set {
|
||||
if cached_resolver.has_changed(self.address.host.as_str(), addr_set) {
|
||||
warn!(
|
||||
"DNS changed for {}, it was {:?}. Dropping server connection.",
|
||||
self.address.host.as_str(),
|
||||
addr_set
|
||||
);
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
self.bad
|
||||
}
|
||||
|
||||
/// Get server startup information to forward it to the client.
|
||||
@@ -1016,8 +882,6 @@ impl Server {
|
||||
/// It will use the simple query protocol.
|
||||
/// Result will not be returned, so this is useful for things like `SET` or `ROLLBACK`.
|
||||
pub async fn query(&mut self, query: &str) -> Result<(), Error> {
|
||||
debug!("Running `{}` on server {:?}", query, self.address);
|
||||
|
||||
let query = simple_query(query);
|
||||
|
||||
self.send(&query).await?;
|
||||
@@ -1050,11 +914,10 @@ impl Server {
|
||||
// to avoid leaking state between clients. For performance reasons we only
|
||||
// send `DISCARD ALL` if we think the session is altered instead of just sending
|
||||
// it before each checkin.
|
||||
if self.cleanup_state.needs_cleanup() && self.cleanup_connections {
|
||||
warn!("Server returned with session state altered, discarding state ({}) for application {}", self.cleanup_state, self.application_name);
|
||||
if self.needs_cleanup {
|
||||
warn!("Server returned with session state altered, discarding state");
|
||||
self.query("DISCARD ALL").await?;
|
||||
self.query("RESET ROLE").await?;
|
||||
self.cleanup_state.reset();
|
||||
self.needs_cleanup = false;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
@@ -1066,12 +929,12 @@ impl Server {
|
||||
self.application_name = name.to_string();
|
||||
// We don't want `SET application_name` to mark the server connection
|
||||
// as needing cleanup
|
||||
let needs_cleanup_before = self.cleanup_state;
|
||||
let needs_cleanup_before = self.needs_cleanup;
|
||||
|
||||
let result = Ok(self
|
||||
.query(&format!("SET application_name = '{}'", name))
|
||||
.await?);
|
||||
self.cleanup_state = needs_cleanup_before;
|
||||
self.needs_cleanup = needs_cleanup_before;
|
||||
result
|
||||
} else {
|
||||
Ok(())
|
||||
@@ -1096,7 +959,7 @@ impl Server {
|
||||
|
||||
// Marks a connection as needing DISCARD ALL at checkin
|
||||
pub fn mark_dirty(&mut self) {
|
||||
self.cleanup_state.set_true();
|
||||
self.needs_cleanup = true;
|
||||
}
|
||||
|
||||
pub fn mirror_send(&mut self, bytes: &BytesMut) {
|
||||
@@ -1130,7 +993,6 @@ impl Server {
|
||||
client_server_map,
|
||||
Arc::new(ServerStats::default()),
|
||||
Arc::new(RwLock::new(None)),
|
||||
true,
|
||||
)
|
||||
.await?;
|
||||
debug!("Connected!, sending query.");
|
||||
@@ -1233,18 +1095,14 @@ impl Drop for Server {
|
||||
_ => debug!("Dirty shutdown"),
|
||||
};
|
||||
|
||||
// Should not matter.
|
||||
self.bad = true;
|
||||
|
||||
let now = chrono::offset::Utc::now().naive_utc();
|
||||
let duration = now - self.connected_at;
|
||||
|
||||
let message = if self.bad {
|
||||
"Server connection terminated"
|
||||
} else {
|
||||
"Server connection closed"
|
||||
};
|
||||
|
||||
info!(
|
||||
"{} {:?}, session duration: {}",
|
||||
message,
|
||||
"Server connection closed {:?}, session duration: {}",
|
||||
self.address,
|
||||
crate::format_duration(&duration)
|
||||
);
|
||||
|
||||
37
src/stats.rs
37
src/stats.rs
@@ -1,3 +1,4 @@
|
||||
use crate::pool::PoolIdentifier;
|
||||
/// Statistics and reporting.
|
||||
use arc_swap::ArcSwap;
|
||||
|
||||
@@ -15,11 +16,13 @@ pub mod pool;
|
||||
pub mod server;
|
||||
pub use address::AddressStats;
|
||||
pub use client::{ClientState, ClientStats};
|
||||
pub use pool::PoolStats;
|
||||
pub use server::{ServerState, ServerStats};
|
||||
|
||||
/// Convenience types for various stats
|
||||
type ClientStatesLookup = HashMap<i32, Arc<ClientStats>>;
|
||||
type ServerStatesLookup = HashMap<i32, Arc<ServerStats>>;
|
||||
type PoolStatsLookup = HashMap<(String, String), Arc<PoolStats>>;
|
||||
|
||||
/// Stats for individual client connections
|
||||
/// Used in SHOW CLIENTS.
|
||||
@@ -31,6 +34,11 @@ static CLIENT_STATS: Lazy<Arc<RwLock<ClientStatesLookup>>> =
|
||||
static SERVER_STATS: Lazy<Arc<RwLock<ServerStatesLookup>>> =
|
||||
Lazy::new(|| Arc::new(RwLock::new(ServerStatesLookup::default())));
|
||||
|
||||
/// Aggregate stats for each pool (a pool is identified by database name and username)
|
||||
/// Used in SHOW POOLS.
|
||||
static POOL_STATS: Lazy<Arc<RwLock<PoolStatsLookup>>> =
|
||||
Lazy::new(|| Arc::new(RwLock::new(PoolStatsLookup::default())));
|
||||
|
||||
/// The statistics reporter. An instance is given to each possible source of statistics,
|
||||
/// e.g. client stats, server stats, connection pool stats.
|
||||
pub static REPORTER: Lazy<ArcSwap<Reporter>> =
|
||||
@@ -72,6 +80,13 @@ impl Reporter {
|
||||
fn server_disconnecting(&self, server_id: i32) {
|
||||
SERVER_STATS.write().remove(&server_id);
|
||||
}
|
||||
|
||||
/// Register a pool with the stats system.
|
||||
fn pool_register(&self, identifier: PoolIdentifier, stats: Arc<PoolStats>) {
|
||||
POOL_STATS
|
||||
.write()
|
||||
.insert((identifier.db, identifier.user), stats);
|
||||
}
|
||||
}
|
||||
|
||||
/// The statistics collector which used for calculating averages
|
||||
@@ -92,20 +107,8 @@ impl Collector {
|
||||
loop {
|
||||
interval.tick().await;
|
||||
|
||||
// Hold read lock for duration of update to retain all server stats
|
||||
let server_stats = SERVER_STATS.read();
|
||||
|
||||
for stats in server_stats.values() {
|
||||
if !stats.check_address_stat_average_is_updated_status() {
|
||||
stats.address_stats().update_averages();
|
||||
stats.address_stats().reset_current_counts();
|
||||
stats.set_address_stat_average_is_updated_status(true);
|
||||
}
|
||||
}
|
||||
|
||||
// Reset to false for next update
|
||||
for stats in server_stats.values() {
|
||||
stats.set_address_stat_average_is_updated_status(false);
|
||||
for stats in SERVER_STATS.read().values() {
|
||||
stats.address_stats().update_averages();
|
||||
}
|
||||
}
|
||||
});
|
||||
@@ -124,6 +127,12 @@ pub fn get_server_stats() -> ServerStatesLookup {
|
||||
SERVER_STATS.read().clone()
|
||||
}
|
||||
|
||||
/// Get a snapshot of pool statistics.
|
||||
/// by the `Collector`.
|
||||
pub fn get_pool_stats() -> PoolStatsLookup {
|
||||
POOL_STATS.read().clone()
|
||||
}
|
||||
|
||||
/// Get the statistics reporter used to update stats across the pools/clients.
|
||||
pub fn get_reporter() -> Reporter {
|
||||
(*(*REPORTER.load())).clone()
|
||||
|
||||
@@ -1,29 +1,26 @@
|
||||
use log::warn;
|
||||
use std::sync::atomic::*;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
struct AddressStatFields {
|
||||
xact_count: Arc<AtomicU64>,
|
||||
query_count: Arc<AtomicU64>,
|
||||
bytes_received: Arc<AtomicU64>,
|
||||
bytes_sent: Arc<AtomicU64>,
|
||||
xact_time: Arc<AtomicU64>,
|
||||
query_time: Arc<AtomicU64>,
|
||||
wait_time: Arc<AtomicU64>,
|
||||
errors: Arc<AtomicU64>,
|
||||
}
|
||||
|
||||
/// Internal address stats
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct AddressStats {
|
||||
total: AddressStatFields,
|
||||
|
||||
current: AddressStatFields,
|
||||
|
||||
averages: AddressStatFields,
|
||||
|
||||
// Determines if the averages have been updated since the last time they were reported
|
||||
pub averages_updated: Arc<AtomicBool>,
|
||||
pub total_xact_count: Arc<AtomicU64>,
|
||||
pub total_query_count: Arc<AtomicU64>,
|
||||
pub total_received: Arc<AtomicU64>,
|
||||
pub total_sent: Arc<AtomicU64>,
|
||||
pub total_xact_time: Arc<AtomicU64>,
|
||||
pub total_query_time: Arc<AtomicU64>,
|
||||
pub total_wait_time: Arc<AtomicU64>,
|
||||
pub total_errors: Arc<AtomicU64>,
|
||||
pub avg_query_count: Arc<AtomicU64>,
|
||||
pub avg_query_time: Arc<AtomicU64>,
|
||||
pub avg_recv: Arc<AtomicU64>,
|
||||
pub avg_sent: Arc<AtomicU64>,
|
||||
pub avg_errors: Arc<AtomicU64>,
|
||||
pub avg_xact_time: Arc<AtomicU64>,
|
||||
pub avg_xact_count: Arc<AtomicU64>,
|
||||
pub avg_wait_time: Arc<AtomicU64>,
|
||||
}
|
||||
|
||||
impl IntoIterator for AddressStats {
|
||||
@@ -34,67 +31,67 @@ impl IntoIterator for AddressStats {
|
||||
vec![
|
||||
(
|
||||
"total_xact_count".to_string(),
|
||||
self.total.xact_count.load(Ordering::Relaxed),
|
||||
self.total_xact_count.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"total_query_count".to_string(),
|
||||
self.total.query_count.load(Ordering::Relaxed),
|
||||
self.total_query_count.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"total_received".to_string(),
|
||||
self.total.bytes_received.load(Ordering::Relaxed),
|
||||
self.total_received.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"total_sent".to_string(),
|
||||
self.total.bytes_sent.load(Ordering::Relaxed),
|
||||
self.total_sent.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"total_xact_time".to_string(),
|
||||
self.total.xact_time.load(Ordering::Relaxed),
|
||||
self.total_xact_time.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"total_query_time".to_string(),
|
||||
self.total.query_time.load(Ordering::Relaxed),
|
||||
self.total_query_time.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"total_wait_time".to_string(),
|
||||
self.total.wait_time.load(Ordering::Relaxed),
|
||||
self.total_wait_time.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"total_errors".to_string(),
|
||||
self.total.errors.load(Ordering::Relaxed),
|
||||
self.total_errors.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"avg_xact_count".to_string(),
|
||||
self.averages.xact_count.load(Ordering::Relaxed),
|
||||
self.avg_xact_count.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"avg_query_count".to_string(),
|
||||
self.averages.query_count.load(Ordering::Relaxed),
|
||||
self.avg_query_count.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"avg_recv".to_string(),
|
||||
self.averages.bytes_received.load(Ordering::Relaxed),
|
||||
self.avg_recv.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"avg_sent".to_string(),
|
||||
self.averages.bytes_sent.load(Ordering::Relaxed),
|
||||
self.avg_sent.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"avg_errors".to_string(),
|
||||
self.averages.errors.load(Ordering::Relaxed),
|
||||
self.avg_errors.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"avg_xact_time".to_string(),
|
||||
self.averages.xact_time.load(Ordering::Relaxed),
|
||||
self.avg_xact_time.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"avg_query_time".to_string(),
|
||||
self.averages.query_time.load(Ordering::Relaxed),
|
||||
self.avg_query_time.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"avg_wait_time".to_string(),
|
||||
self.averages.wait_time.load(Ordering::Relaxed),
|
||||
self.avg_wait_time.load(Ordering::Relaxed),
|
||||
),
|
||||
]
|
||||
.into_iter()
|
||||
@@ -102,120 +99,22 @@ impl IntoIterator for AddressStats {
|
||||
}
|
||||
|
||||
impl AddressStats {
|
||||
pub fn xact_count_add(&self) {
|
||||
self.total.xact_count.fetch_add(1, Ordering::Relaxed);
|
||||
self.current.xact_count.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn query_count_add(&self) {
|
||||
self.total.query_count.fetch_add(1, Ordering::Relaxed);
|
||||
self.current.query_count.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn bytes_received_add(&self, bytes: u64) {
|
||||
self.total
|
||||
.bytes_received
|
||||
.fetch_add(bytes, Ordering::Relaxed);
|
||||
self.current
|
||||
.bytes_received
|
||||
.fetch_add(bytes, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn bytes_sent_add(&self, bytes: u64) {
|
||||
self.total.bytes_sent.fetch_add(bytes, Ordering::Relaxed);
|
||||
self.current.bytes_sent.fetch_add(bytes, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn xact_time_add(&self, time: u64) {
|
||||
self.total.xact_time.fetch_add(time, Ordering::Relaxed);
|
||||
self.current.xact_time.fetch_add(time, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn query_time_add(&self, time: u64) {
|
||||
self.total.query_time.fetch_add(time, Ordering::Relaxed);
|
||||
self.current.query_time.fetch_add(time, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn wait_time_add(&self, time: u64) {
|
||||
self.total.wait_time.fetch_add(time, Ordering::Relaxed);
|
||||
self.current.wait_time.fetch_add(time, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn error(&self) {
|
||||
self.total.errors.fetch_add(1, Ordering::Relaxed);
|
||||
self.current.errors.fetch_add(1, Ordering::Relaxed);
|
||||
self.total_errors.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn update_averages(&self) {
|
||||
let stat_period_per_second = crate::stats::STAT_PERIOD / 1_000;
|
||||
|
||||
// xact_count
|
||||
let current_xact_count = self.current.xact_count.load(Ordering::Relaxed);
|
||||
let current_xact_time = self.current.xact_time.load(Ordering::Relaxed);
|
||||
self.averages.xact_count.store(
|
||||
current_xact_count / stat_period_per_second,
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
if current_xact_count == 0 {
|
||||
self.averages.xact_time.store(0, Ordering::Relaxed);
|
||||
} else {
|
||||
self.averages
|
||||
.xact_time
|
||||
.store(current_xact_time / current_xact_count, Ordering::Relaxed);
|
||||
let (totals, averages) = self.fields_iterators();
|
||||
for data in totals.iter().zip(averages.iter()) {
|
||||
let (total, average) = data;
|
||||
if let Err(err) = average.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |avg| {
|
||||
let total = total.load(Ordering::Relaxed);
|
||||
let avg = (total - avg) / (crate::stats::STAT_PERIOD / 1_000); // Avg / second
|
||||
Some(avg)
|
||||
}) {
|
||||
warn!("Could not update averages for addresses stats, {:?}", err);
|
||||
}
|
||||
}
|
||||
|
||||
// query_count
|
||||
let current_query_count = self.current.query_count.load(Ordering::Relaxed);
|
||||
let current_query_time = self.current.query_time.load(Ordering::Relaxed);
|
||||
self.averages.query_count.store(
|
||||
current_query_count / stat_period_per_second,
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
if current_query_count == 0 {
|
||||
self.averages.query_time.store(0, Ordering::Relaxed);
|
||||
} else {
|
||||
self.averages
|
||||
.query_time
|
||||
.store(current_query_time / current_query_count, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
// bytes_received
|
||||
let current_bytes_received = self.current.bytes_received.load(Ordering::Relaxed);
|
||||
self.averages.bytes_received.store(
|
||||
current_bytes_received / stat_period_per_second,
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
|
||||
// bytes_sent
|
||||
let current_bytes_sent = self.current.bytes_sent.load(Ordering::Relaxed);
|
||||
self.averages.bytes_sent.store(
|
||||
current_bytes_sent / stat_period_per_second,
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
|
||||
// wait_time
|
||||
let current_wait_time = self.current.wait_time.load(Ordering::Relaxed);
|
||||
self.averages.wait_time.store(
|
||||
current_wait_time / stat_period_per_second,
|
||||
Ordering::Relaxed,
|
||||
);
|
||||
|
||||
// errors
|
||||
let current_errors = self.current.errors.load(Ordering::Relaxed);
|
||||
self.averages
|
||||
.errors
|
||||
.store(current_errors / stat_period_per_second, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn reset_current_counts(&self) {
|
||||
self.current.xact_count.store(0, Ordering::Relaxed);
|
||||
self.current.xact_time.store(0, Ordering::Relaxed);
|
||||
self.current.query_count.store(0, Ordering::Relaxed);
|
||||
self.current.query_time.store(0, Ordering::Relaxed);
|
||||
self.current.bytes_received.store(0, Ordering::Relaxed);
|
||||
self.current.bytes_sent.store(0, Ordering::Relaxed);
|
||||
self.current.wait_time.store(0, Ordering::Relaxed);
|
||||
self.current.errors.store(0, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub fn populate_row(&self, row: &mut Vec<String>) {
|
||||
@@ -223,4 +122,28 @@ impl AddressStats {
|
||||
row.push(value.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
fn fields_iterators(&self) -> (Vec<Arc<AtomicU64>>, Vec<Arc<AtomicU64>>) {
|
||||
let mut totals: Vec<Arc<AtomicU64>> = Vec::new();
|
||||
let mut averages: Vec<Arc<AtomicU64>> = Vec::new();
|
||||
|
||||
totals.push(self.total_xact_count.clone());
|
||||
averages.push(self.avg_xact_count.clone());
|
||||
totals.push(self.total_query_count.clone());
|
||||
averages.push(self.avg_query_count.clone());
|
||||
totals.push(self.total_received.clone());
|
||||
averages.push(self.avg_recv.clone());
|
||||
totals.push(self.total_sent.clone());
|
||||
averages.push(self.avg_sent.clone());
|
||||
totals.push(self.total_xact_time.clone());
|
||||
averages.push(self.avg_xact_time.clone());
|
||||
totals.push(self.total_query_time.clone());
|
||||
averages.push(self.avg_query_time.clone());
|
||||
totals.push(self.total_wait_time.clone());
|
||||
averages.push(self.avg_wait_time.clone());
|
||||
totals.push(self.total_errors.clone());
|
||||
averages.push(self.avg_errors.clone());
|
||||
|
||||
(totals, averages)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use super::PoolStats;
|
||||
use super::{get_reporter, Reporter};
|
||||
use atomic_enum::atomic_enum;
|
||||
use std::sync::atomic::*;
|
||||
@@ -33,14 +34,12 @@ pub struct ClientStats {
|
||||
pool_name: String,
|
||||
connect_time: Instant,
|
||||
|
||||
pool_stats: Arc<PoolStats>,
|
||||
reporter: Reporter,
|
||||
|
||||
/// Total time spent waiting for a connection from pool, measures in microseconds
|
||||
pub total_wait_time: Arc<AtomicU64>,
|
||||
|
||||
/// Maximum time spent waiting for a connection from pool, measures in microseconds
|
||||
pub max_wait_time: Arc<AtomicU64>,
|
||||
|
||||
/// Current state of the client
|
||||
pub state: Arc<AtomicClientState>,
|
||||
|
||||
@@ -62,8 +61,8 @@ impl Default for ClientStats {
|
||||
application_name: String::new(),
|
||||
username: String::new(),
|
||||
pool_name: String::new(),
|
||||
pool_stats: Arc::new(PoolStats::default()),
|
||||
total_wait_time: Arc::new(AtomicU64::new(0)),
|
||||
max_wait_time: Arc::new(AtomicU64::new(0)),
|
||||
state: Arc::new(AtomicClientState::new(ClientState::Idle)),
|
||||
transaction_count: Arc::new(AtomicU64::new(0)),
|
||||
query_count: Arc::new(AtomicU64::new(0)),
|
||||
@@ -80,9 +79,11 @@ impl ClientStats {
|
||||
username: &str,
|
||||
pool_name: &str,
|
||||
connect_time: Instant,
|
||||
pool_stats: Arc<PoolStats>,
|
||||
) -> Self {
|
||||
Self {
|
||||
client_id,
|
||||
pool_stats,
|
||||
connect_time,
|
||||
application_name: application_name.to_string(),
|
||||
username: username.to_string(),
|
||||
@@ -95,6 +96,8 @@ impl ClientStats {
|
||||
/// update metrics on the corresponding pool.
|
||||
pub fn disconnect(&self) {
|
||||
self.reporter.client_disconnecting(self.client_id);
|
||||
self.pool_stats
|
||||
.client_disconnect(self.state.load(Ordering::Relaxed))
|
||||
}
|
||||
|
||||
/// Register a client with the stats system. The stats system uses client_id
|
||||
@@ -102,20 +105,27 @@ impl ClientStats {
|
||||
pub fn register(&self, stats: Arc<ClientStats>) {
|
||||
self.reporter.client_register(self.client_id, stats);
|
||||
self.state.store(ClientState::Idle, Ordering::Relaxed);
|
||||
self.pool_stats.cl_idle.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Reports a client is done querying the server and is no longer assigned a server connection
|
||||
pub fn idle(&self) {
|
||||
self.pool_stats
|
||||
.client_idle(self.state.load(Ordering::Relaxed));
|
||||
self.state.store(ClientState::Idle, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Reports a client is waiting for a connection
|
||||
pub fn waiting(&self) {
|
||||
self.pool_stats
|
||||
.client_waiting(self.state.load(Ordering::Relaxed));
|
||||
self.state.store(ClientState::Waiting, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Reports a client is done waiting for a connection and is about to query the server.
|
||||
pub fn active(&self) {
|
||||
self.pool_stats
|
||||
.client_active(self.state.load(Ordering::Relaxed));
|
||||
self.state.store(ClientState::Active, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
@@ -134,8 +144,6 @@ impl ClientStats {
|
||||
pub fn checkout_time(&self, microseconds: u64) {
|
||||
self.total_wait_time
|
||||
.fetch_add(microseconds, Ordering::Relaxed);
|
||||
self.max_wait_time
|
||||
.fetch_max(microseconds, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Report a query executed by a client against a server
|
||||
|
||||
@@ -1,131 +1,36 @@
|
||||
use log::debug;
|
||||
|
||||
use super::{ClientState, ServerState};
|
||||
use crate::{config::PoolMode, messages::DataType, pool::PoolIdentifier};
|
||||
use std::collections::HashMap;
|
||||
use crate::config::Pool;
|
||||
use crate::config::PoolMode;
|
||||
use crate::pool::PoolIdentifier;
|
||||
use std::sync::atomic::*;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::pool::get_all_pools;
|
||||
use super::get_reporter;
|
||||
use super::Reporter;
|
||||
use super::{ClientState, ServerState};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
#[derive(Debug, Clone, Default)]
|
||||
/// A struct that holds information about a Pool .
|
||||
pub struct PoolStats {
|
||||
pub identifier: PoolIdentifier,
|
||||
pub mode: PoolMode,
|
||||
pub cl_idle: u64,
|
||||
pub cl_active: u64,
|
||||
pub cl_waiting: u64,
|
||||
pub cl_cancel_req: u64,
|
||||
pub sv_active: u64,
|
||||
pub sv_idle: u64,
|
||||
pub sv_used: u64,
|
||||
pub sv_tested: u64,
|
||||
pub sv_login: u64,
|
||||
pub maxwait: u64,
|
||||
}
|
||||
impl PoolStats {
|
||||
pub fn new(identifier: PoolIdentifier, mode: PoolMode) -> Self {
|
||||
PoolStats {
|
||||
identifier,
|
||||
mode,
|
||||
cl_idle: 0,
|
||||
cl_active: 0,
|
||||
cl_waiting: 0,
|
||||
cl_cancel_req: 0,
|
||||
sv_active: 0,
|
||||
sv_idle: 0,
|
||||
sv_used: 0,
|
||||
sv_tested: 0,
|
||||
sv_login: 0,
|
||||
maxwait: 0,
|
||||
}
|
||||
}
|
||||
// Pool identifier, cannot be changed after creating the instance
|
||||
identifier: PoolIdentifier,
|
||||
|
||||
pub fn construct_pool_lookup() -> HashMap<PoolIdentifier, PoolStats> {
|
||||
let mut map: HashMap<PoolIdentifier, PoolStats> = HashMap::new();
|
||||
let client_map = super::get_client_stats();
|
||||
let server_map = super::get_server_stats();
|
||||
// Pool Config, cannot be changed after creating the instance
|
||||
config: Pool,
|
||||
|
||||
for (identifier, pool) in get_all_pools() {
|
||||
map.insert(
|
||||
identifier.clone(),
|
||||
PoolStats::new(identifier, pool.settings.pool_mode),
|
||||
);
|
||||
}
|
||||
// A reference to the global reporter.
|
||||
reporter: Reporter,
|
||||
|
||||
for client in client_map.values() {
|
||||
match map.get_mut(&PoolIdentifier {
|
||||
db: client.pool_name(),
|
||||
user: client.username(),
|
||||
}) {
|
||||
Some(pool_stats) => {
|
||||
match client.state.load(Ordering::Relaxed) {
|
||||
ClientState::Active => pool_stats.cl_active += 1,
|
||||
ClientState::Idle => pool_stats.cl_idle += 1,
|
||||
ClientState::Waiting => pool_stats.cl_waiting += 1,
|
||||
}
|
||||
let max_wait = client.max_wait_time.load(Ordering::Relaxed);
|
||||
pool_stats.maxwait = std::cmp::max(pool_stats.maxwait, max_wait);
|
||||
}
|
||||
None => debug!("Client from an obselete pool"),
|
||||
}
|
||||
}
|
||||
|
||||
for server in server_map.values() {
|
||||
match map.get_mut(&PoolIdentifier {
|
||||
db: server.pool_name(),
|
||||
user: server.username(),
|
||||
}) {
|
||||
Some(pool_stats) => match server.state.load(Ordering::Relaxed) {
|
||||
ServerState::Active => pool_stats.sv_active += 1,
|
||||
ServerState::Idle => pool_stats.sv_idle += 1,
|
||||
ServerState::Login => pool_stats.sv_login += 1,
|
||||
ServerState::Tested => pool_stats.sv_tested += 1,
|
||||
},
|
||||
None => debug!("Server from an obselete pool"),
|
||||
}
|
||||
}
|
||||
|
||||
return map;
|
||||
}
|
||||
|
||||
pub fn generate_header() -> Vec<(&'static str, DataType)> {
|
||||
return vec![
|
||||
("database", DataType::Text),
|
||||
("user", DataType::Text),
|
||||
("pool_mode", DataType::Text),
|
||||
("cl_idle", DataType::Numeric),
|
||||
("cl_active", DataType::Numeric),
|
||||
("cl_waiting", DataType::Numeric),
|
||||
("cl_cancel_req", DataType::Numeric),
|
||||
("sv_active", DataType::Numeric),
|
||||
("sv_idle", DataType::Numeric),
|
||||
("sv_used", DataType::Numeric),
|
||||
("sv_tested", DataType::Numeric),
|
||||
("sv_login", DataType::Numeric),
|
||||
("maxwait", DataType::Numeric),
|
||||
("maxwait_us", DataType::Numeric),
|
||||
];
|
||||
}
|
||||
|
||||
pub fn generate_row(&self) -> Vec<String> {
|
||||
return vec![
|
||||
self.identifier.db.clone(),
|
||||
self.identifier.user.clone(),
|
||||
self.mode.to_string(),
|
||||
self.cl_idle.to_string(),
|
||||
self.cl_active.to_string(),
|
||||
self.cl_waiting.to_string(),
|
||||
self.cl_cancel_req.to_string(),
|
||||
self.sv_active.to_string(),
|
||||
self.sv_idle.to_string(),
|
||||
self.sv_used.to_string(),
|
||||
self.sv_tested.to_string(),
|
||||
self.sv_login.to_string(),
|
||||
(self.maxwait / 1_000_000).to_string(),
|
||||
(self.maxwait % 1_000_000).to_string(),
|
||||
];
|
||||
}
|
||||
/// Counters (atomics)
|
||||
pub cl_idle: Arc<AtomicU64>,
|
||||
pub cl_active: Arc<AtomicU64>,
|
||||
pub cl_waiting: Arc<AtomicU64>,
|
||||
pub cl_cancel_req: Arc<AtomicU64>,
|
||||
pub sv_active: Arc<AtomicU64>,
|
||||
pub sv_idle: Arc<AtomicU64>,
|
||||
pub sv_used: Arc<AtomicU64>,
|
||||
pub sv_tested: Arc<AtomicU64>,
|
||||
pub sv_login: Arc<AtomicU64>,
|
||||
pub maxwait: Arc<AtomicU64>,
|
||||
}
|
||||
|
||||
impl IntoIterator for PoolStats {
|
||||
@@ -134,18 +39,236 @@ impl IntoIterator for PoolStats {
|
||||
|
||||
fn into_iter(self) -> Self::IntoIter {
|
||||
vec![
|
||||
("cl_idle".to_string(), self.cl_idle),
|
||||
("cl_active".to_string(), self.cl_active),
|
||||
("cl_waiting".to_string(), self.cl_waiting),
|
||||
("cl_cancel_req".to_string(), self.cl_cancel_req),
|
||||
("sv_active".to_string(), self.sv_active),
|
||||
("sv_idle".to_string(), self.sv_idle),
|
||||
("sv_used".to_string(), self.sv_used),
|
||||
("sv_tested".to_string(), self.sv_tested),
|
||||
("sv_login".to_string(), self.sv_login),
|
||||
("maxwait".to_string(), self.maxwait / 1_000_000),
|
||||
("maxwait_us".to_string(), self.maxwait % 1_000_000),
|
||||
("cl_idle".to_string(), self.cl_idle.load(Ordering::Relaxed)),
|
||||
(
|
||||
"cl_active".to_string(),
|
||||
self.cl_active.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"cl_waiting".to_string(),
|
||||
self.cl_waiting.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"cl_cancel_req".to_string(),
|
||||
self.cl_cancel_req.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"sv_active".to_string(),
|
||||
self.sv_active.load(Ordering::Relaxed),
|
||||
),
|
||||
("sv_idle".to_string(), self.sv_idle.load(Ordering::Relaxed)),
|
||||
("sv_used".to_string(), self.sv_used.load(Ordering::Relaxed)),
|
||||
(
|
||||
"sv_tested".to_string(),
|
||||
self.sv_tested.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"sv_login".to_string(),
|
||||
self.sv_login.load(Ordering::Relaxed),
|
||||
),
|
||||
(
|
||||
"maxwait".to_string(),
|
||||
self.maxwait.load(Ordering::Relaxed) / 1_000_000,
|
||||
),
|
||||
(
|
||||
"maxwait_us".to_string(),
|
||||
self.maxwait.load(Ordering::Relaxed) % 1_000_000,
|
||||
),
|
||||
]
|
||||
.into_iter()
|
||||
}
|
||||
}
|
||||
|
||||
impl PoolStats {
|
||||
pub fn new(identifier: PoolIdentifier, config: Pool) -> Self {
|
||||
Self {
|
||||
identifier,
|
||||
config,
|
||||
reporter: get_reporter(),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
// Getters
|
||||
pub fn register(&self, stats: Arc<PoolStats>) {
|
||||
self.reporter.pool_register(self.identifier.clone(), stats);
|
||||
}
|
||||
|
||||
pub fn database(&self) -> String {
|
||||
self.identifier.db.clone()
|
||||
}
|
||||
|
||||
pub fn user(&self) -> String {
|
||||
self.identifier.user.clone()
|
||||
}
|
||||
|
||||
pub fn pool_mode(&self) -> PoolMode {
|
||||
self.config.pool_mode
|
||||
}
|
||||
|
||||
/// Populates an array of strings with counters (used by admin in show pools)
|
||||
pub fn populate_row(&self, row: &mut Vec<String>) {
|
||||
for (_key, value) in self.clone() {
|
||||
row.push(value.to_string());
|
||||
}
|
||||
}
|
||||
|
||||
/// Deletes the maxwait counter, this is done everytime we obtain metrics
|
||||
pub fn clear_maxwait(&self) {
|
||||
self.maxwait.store(0, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Notified when a server of the pool enters login state.
|
||||
///
|
||||
/// Arguments:
|
||||
///
|
||||
/// `from`: The state of the server that notifies.
|
||||
pub fn server_login(&self, from: ServerState) {
|
||||
self.sv_login.fetch_add(1, Ordering::Relaxed);
|
||||
if from != ServerState::Login {
|
||||
self.decrease_from_server_state(from);
|
||||
}
|
||||
}
|
||||
|
||||
/// Notified when a server of the pool become 'active'
|
||||
///
|
||||
/// Arguments:
|
||||
///
|
||||
/// `from`: The state of the server that notifies.
|
||||
pub fn server_active(&self, from: ServerState) {
|
||||
self.sv_active.fetch_add(1, Ordering::Relaxed);
|
||||
if from != ServerState::Active {
|
||||
self.decrease_from_server_state(from);
|
||||
}
|
||||
}
|
||||
|
||||
/// Notified when a server of the pool become 'tested'
|
||||
///
|
||||
/// Arguments:
|
||||
///
|
||||
/// `from`: The state of the server that notifies.
|
||||
pub fn server_tested(&self, from: ServerState) {
|
||||
self.sv_tested.fetch_add(1, Ordering::Relaxed);
|
||||
if from != ServerState::Tested {
|
||||
self.decrease_from_server_state(from);
|
||||
}
|
||||
}
|
||||
|
||||
/// Notified when a server of the pool become 'idle'
|
||||
///
|
||||
/// Arguments:
|
||||
///
|
||||
/// `from`: The state of the server that notifies.
|
||||
pub fn server_idle(&self, from: ServerState) {
|
||||
self.sv_idle.fetch_add(1, Ordering::Relaxed);
|
||||
if from != ServerState::Idle {
|
||||
self.decrease_from_server_state(from);
|
||||
}
|
||||
}
|
||||
|
||||
/// Notified when a client of the pool become 'waiting'
|
||||
///
|
||||
/// Arguments:
|
||||
///
|
||||
/// `from`: The state of the client that notifies.
|
||||
pub fn client_waiting(&self, from: ClientState) {
|
||||
if from != ClientState::Waiting {
|
||||
self.cl_waiting.fetch_add(1, Ordering::Relaxed);
|
||||
self.decrease_from_client_state(from);
|
||||
}
|
||||
}
|
||||
|
||||
/// Notified when a client of the pool become 'active'
|
||||
///
|
||||
/// Arguments:
|
||||
///
|
||||
/// `from`: The state of the client that notifies.
|
||||
pub fn client_active(&self, from: ClientState) {
|
||||
if from != ClientState::Active {
|
||||
self.cl_active.fetch_add(1, Ordering::Relaxed);
|
||||
self.decrease_from_client_state(from);
|
||||
}
|
||||
}
|
||||
|
||||
/// Notified when a client of the pool become 'idle'
|
||||
///
|
||||
/// Arguments:
|
||||
///
|
||||
/// `from`: The state of the client that notifies.
|
||||
pub fn client_idle(&self, from: ClientState) {
|
||||
if from != ClientState::Idle {
|
||||
self.cl_idle.fetch_add(1, Ordering::Relaxed);
|
||||
self.decrease_from_client_state(from);
|
||||
}
|
||||
}
|
||||
|
||||
/// Notified when a client disconnects.
|
||||
///
|
||||
/// Arguments:
|
||||
///
|
||||
/// `from`: The state of the client that notifies.
|
||||
pub fn client_disconnect(&self, from: ClientState) {
|
||||
let counter = match from {
|
||||
ClientState::Idle => &self.cl_idle,
|
||||
ClientState::Waiting => &self.cl_waiting,
|
||||
ClientState::Active => &self.cl_active,
|
||||
};
|
||||
|
||||
Self::decrease_counter(counter.clone());
|
||||
}
|
||||
|
||||
/// Notified when a server disconnects.
|
||||
///
|
||||
/// Arguments:
|
||||
///
|
||||
/// `from`: The state of the client that notifies.
|
||||
pub fn server_disconnect(&self, from: ServerState) {
|
||||
let counter = match from {
|
||||
ServerState::Active => &self.sv_active,
|
||||
ServerState::Idle => &self.sv_idle,
|
||||
ServerState::Login => &self.sv_login,
|
||||
ServerState::Tested => &self.sv_tested,
|
||||
};
|
||||
Self::decrease_counter(counter.clone());
|
||||
}
|
||||
|
||||
// helpers for counter decrease
|
||||
fn decrease_from_server_state(&self, from: ServerState) {
|
||||
let counter = match from {
|
||||
ServerState::Tested => &self.sv_tested,
|
||||
ServerState::Active => &self.sv_active,
|
||||
ServerState::Idle => &self.sv_idle,
|
||||
ServerState::Login => &self.sv_login,
|
||||
};
|
||||
Self::decrease_counter(counter.clone());
|
||||
}
|
||||
|
||||
fn decrease_from_client_state(&self, from: ClientState) {
|
||||
let counter = match from {
|
||||
ClientState::Active => &self.cl_active,
|
||||
ClientState::Idle => &self.cl_idle,
|
||||
ClientState::Waiting => &self.cl_waiting,
|
||||
};
|
||||
Self::decrease_counter(counter.clone());
|
||||
}
|
||||
|
||||
fn decrease_counter(value: Arc<AtomicU64>) {
|
||||
if value.load(Ordering::Relaxed) > 0 {
|
||||
value.fetch_sub(1, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_decrease() {
|
||||
let stat: PoolStats = PoolStats::default();
|
||||
stat.server_login(ServerState::Login);
|
||||
stat.server_idle(ServerState::Login);
|
||||
assert_eq!(stat.sv_login.load(Ordering::Relaxed), 0);
|
||||
assert_eq!(stat.sv_idle.load(Ordering::Relaxed), 1);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use super::AddressStats;
|
||||
use super::PoolStats;
|
||||
use super::{get_reporter, Reporter};
|
||||
use crate::config::Address;
|
||||
use atomic_enum::atomic_enum;
|
||||
@@ -37,6 +38,7 @@ pub struct ServerStats {
|
||||
address: Address,
|
||||
connect_time: Instant,
|
||||
|
||||
pool_stats: Arc<PoolStats>,
|
||||
reporter: Reporter,
|
||||
|
||||
/// Data
|
||||
@@ -47,8 +49,6 @@ pub struct ServerStats {
|
||||
pub transaction_count: Arc<AtomicU64>,
|
||||
pub query_count: Arc<AtomicU64>,
|
||||
pub error_count: Arc<AtomicU64>,
|
||||
pub prepared_hit_count: Arc<AtomicU64>,
|
||||
pub prepared_miss_count: Arc<AtomicU64>,
|
||||
}
|
||||
|
||||
impl Default for ServerStats {
|
||||
@@ -57,6 +57,7 @@ impl Default for ServerStats {
|
||||
server_id: 0,
|
||||
application_name: Arc::new(RwLock::new(String::new())),
|
||||
address: Address::default(),
|
||||
pool_stats: Arc::new(PoolStats::default()),
|
||||
connect_time: Instant::now(),
|
||||
state: Arc::new(AtomicServerState::new(ServerState::Login)),
|
||||
bytes_sent: Arc::new(AtomicU64::new(0)),
|
||||
@@ -65,16 +66,15 @@ impl Default for ServerStats {
|
||||
query_count: Arc::new(AtomicU64::new(0)),
|
||||
error_count: Arc::new(AtomicU64::new(0)),
|
||||
reporter: get_reporter(),
|
||||
prepared_hit_count: Arc::new(AtomicU64::new(0)),
|
||||
prepared_miss_count: Arc::new(AtomicU64::new(0)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ServerStats {
|
||||
pub fn new(address: Address, connect_time: Instant) -> Self {
|
||||
pub fn new(address: Address, pool_stats: Arc<PoolStats>, connect_time: Instant) -> Self {
|
||||
Self {
|
||||
address,
|
||||
pool_stats,
|
||||
connect_time,
|
||||
server_id: rand::random::<i32>(),
|
||||
..Default::default()
|
||||
@@ -96,6 +96,9 @@ impl ServerStats {
|
||||
/// Reports a server connection is no longer assigned to a client
|
||||
/// and is available for the next client to pick it up
|
||||
pub fn idle(&self) {
|
||||
self.pool_stats
|
||||
.server_idle(self.state.load(Ordering::Relaxed));
|
||||
|
||||
self.state.store(ServerState::Idle, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
@@ -103,16 +106,22 @@ impl ServerStats {
|
||||
/// Also updates metrics on the pool regarding server usage.
|
||||
pub fn disconnect(&self) {
|
||||
self.reporter.server_disconnecting(self.server_id);
|
||||
self.pool_stats
|
||||
.server_disconnect(self.state.load(Ordering::Relaxed))
|
||||
}
|
||||
|
||||
/// Reports a server connection is being tested before being given to a client.
|
||||
pub fn tested(&self) {
|
||||
self.set_undefined_application();
|
||||
self.pool_stats
|
||||
.server_tested(self.state.load(Ordering::Relaxed));
|
||||
self.state.store(ServerState::Tested, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Reports a server connection is attempting to login.
|
||||
pub fn login(&self) {
|
||||
self.pool_stats
|
||||
.server_login(self.state.load(Ordering::Relaxed));
|
||||
self.state.store(ServerState::Login, Ordering::Relaxed);
|
||||
self.set_undefined_application();
|
||||
}
|
||||
@@ -120,6 +129,8 @@ impl ServerStats {
|
||||
/// Reports a server connection has been assigned to a client that
|
||||
/// is about to query the server
|
||||
pub fn active(&self, application_name: String) {
|
||||
self.pool_stats
|
||||
.server_active(self.state.load(Ordering::Relaxed));
|
||||
self.state.store(ServerState::Active, Ordering::Relaxed);
|
||||
self.set_application(application_name);
|
||||
}
|
||||
@@ -128,24 +139,13 @@ impl ServerStats {
|
||||
self.address.stats.clone()
|
||||
}
|
||||
|
||||
pub fn check_address_stat_average_is_updated_status(&self) -> bool {
|
||||
self.address.stats.averages_updated.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
pub fn set_address_stat_average_is_updated_status(&self, is_checked: bool) {
|
||||
self.address
|
||||
.stats
|
||||
.averages_updated
|
||||
.store(is_checked, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
// Helper methods for show_servers
|
||||
pub fn pool_name(&self) -> String {
|
||||
self.address.pool_name.clone()
|
||||
self.pool_stats.database()
|
||||
}
|
||||
|
||||
pub fn username(&self) -> String {
|
||||
self.address.username.clone()
|
||||
self.pool_stats.user()
|
||||
}
|
||||
|
||||
pub fn address_name(&self) -> String {
|
||||
@@ -166,17 +166,27 @@ impl ServerStats {
|
||||
}
|
||||
|
||||
pub fn checkout_time(&self, microseconds: u64, application_name: String) {
|
||||
// Update server stats and address aggregation stats
|
||||
// Update server stats and address aggergation stats
|
||||
self.set_application(application_name);
|
||||
self.address.stats.wait_time_add(microseconds);
|
||||
self.address
|
||||
.stats
|
||||
.total_wait_time
|
||||
.fetch_add(microseconds, Ordering::Relaxed);
|
||||
self.pool_stats
|
||||
.maxwait
|
||||
.fetch_max(microseconds, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Report a query executed by a client against a server
|
||||
pub fn query(&self, milliseconds: u64, application_name: &str) {
|
||||
self.set_application(application_name.to_string());
|
||||
self.address.stats.query_count_add();
|
||||
self.address.stats.query_time_add(milliseconds);
|
||||
self.query_count.fetch_add(1, Ordering::Relaxed);
|
||||
let address_stats = self.address_stats();
|
||||
address_stats
|
||||
.total_query_count
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
address_stats
|
||||
.total_query_time
|
||||
.fetch_add(milliseconds, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Report a transaction executed by a client a server
|
||||
@@ -187,30 +197,29 @@ impl ServerStats {
|
||||
self.set_application(application_name.to_string());
|
||||
|
||||
self.transaction_count.fetch_add(1, Ordering::Relaxed);
|
||||
self.address.stats.xact_count_add();
|
||||
self.address
|
||||
.stats
|
||||
.total_xact_count
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Report data sent to a server
|
||||
pub fn data_sent(&self, amount_bytes: usize) {
|
||||
self.bytes_sent
|
||||
.fetch_add(amount_bytes as u64, Ordering::Relaxed);
|
||||
self.address.stats.bytes_sent_add(amount_bytes as u64);
|
||||
self.address
|
||||
.stats
|
||||
.total_sent
|
||||
.fetch_add(amount_bytes as u64, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Report data received from a server
|
||||
pub fn data_received(&self, amount_bytes: usize) {
|
||||
self.bytes_received
|
||||
.fetch_add(amount_bytes as u64, Ordering::Relaxed);
|
||||
self.address.stats.bytes_received_add(amount_bytes as u64);
|
||||
}
|
||||
|
||||
/// Report a prepared statement that already exists on the server.
|
||||
pub fn prepared_cache_hit(&self) {
|
||||
self.prepared_hit_count.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Report a prepared statement that does not exist on the server yet.
|
||||
pub fn prepared_cache_miss(&self) {
|
||||
self.prepared_miss_count.fetch_add(1, Ordering::Relaxed);
|
||||
self.address
|
||||
.stats
|
||||
.total_received
|
||||
.fetch_add(amount_bytes as u64, Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -63,7 +63,6 @@ def cleanup_conn(conn: psycopg2.extensions.connection, cur: psycopg2.extensions.
|
||||
|
||||
|
||||
def test_normal_db_access():
|
||||
pgcat_start()
|
||||
conn, cur = connect_db(autocommit=False)
|
||||
cur.execute("SELECT 1")
|
||||
res = cur.fetchall()
|
||||
|
||||
@@ -11,6 +11,323 @@ describe "Admin" do
|
||||
processes.pgcat.shutdown
|
||||
end
|
||||
|
||||
describe "SHOW STATS" do
|
||||
context "clients connect and make one query" do
|
||||
it "updates *_query_time and *_wait_time" do
|
||||
connection = PG::connect("#{pgcat_conn_str}?application_name=one_query")
|
||||
connection.async_exec("SELECT pg_sleep(0.25)")
|
||||
connection.async_exec("SELECT pg_sleep(0.25)")
|
||||
connection.async_exec("SELECT pg_sleep(0.25)")
|
||||
connection.close
|
||||
|
||||
# wait for averages to be calculated, we shouldn't do this too often
|
||||
sleep(15.5)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW STATS")[0]
|
||||
admin_conn.close
|
||||
expect(results["total_query_time"].to_i).to be_within(200).of(750)
|
||||
expect(results["avg_query_time"].to_i).to_not eq(0)
|
||||
|
||||
expect(results["total_wait_time"].to_i).to_not eq(0)
|
||||
expect(results["avg_wait_time"].to_i).to_not eq(0)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
describe "SHOW POOLS" do
|
||||
context "bad credentials" do
|
||||
it "does not change any stats" do
|
||||
bad_password_url = URI(pgcat_conn_str)
|
||||
bad_password_url.password = "wrong"
|
||||
expect { PG::connect("#{bad_password_url.to_s}?application_name=bad_password") }.to raise_error(PG::ConnectionBad)
|
||||
|
||||
sleep(1)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
|
||||
expect(results["sv_idle"]).to eq("1")
|
||||
end
|
||||
end
|
||||
|
||||
context "bad database name" do
|
||||
it "does not change any stats" do
|
||||
bad_db_url = URI(pgcat_conn_str)
|
||||
bad_db_url.path = "/wrong_db"
|
||||
expect { PG::connect("#{bad_db_url.to_s}?application_name=bad_db") }.to raise_error(PG::ConnectionBad)
|
||||
|
||||
sleep(1)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
|
||||
expect(results["sv_idle"]).to eq("1")
|
||||
end
|
||||
end
|
||||
|
||||
context "client connects but issues no queries" do
|
||||
it "only affects cl_idle stats" do
|
||||
connections = Array.new(20) { PG::connect(pgcat_conn_str) }
|
||||
sleep(1)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["cl_idle"]).to eq("20")
|
||||
expect(results["sv_idle"]).to eq("1")
|
||||
|
||||
connections.map(&:close)
|
||||
sleep(1.1)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_active cl_idle cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["sv_idle"]).to eq("1")
|
||||
end
|
||||
end
|
||||
|
||||
context "clients connect and make one query" do
|
||||
it "only affects cl_idle, sv_idle stats" do
|
||||
connections = Array.new(5) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
Thread.new { c.async_exec("SELECT pg_sleep(2.5)") }
|
||||
end
|
||||
|
||||
sleep(1.1)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_waiting cl_cancel_req sv_idle sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["cl_active"]).to eq("5")
|
||||
expect(results["sv_active"]).to eq("5")
|
||||
|
||||
sleep(3)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["cl_idle"]).to eq("5")
|
||||
expect(results["sv_idle"]).to eq("5")
|
||||
|
||||
connections.map(&:close)
|
||||
sleep(1)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["sv_idle"]).to eq("5")
|
||||
end
|
||||
end
|
||||
|
||||
context "client connects and opens a transaction and closes connection uncleanly" do
|
||||
it "produces correct statistics" do
|
||||
connections = Array.new(5) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
Thread.new do
|
||||
c.async_exec("BEGIN")
|
||||
c.async_exec("SELECT pg_sleep(0.01)")
|
||||
c.close
|
||||
end
|
||||
end
|
||||
|
||||
sleep(1.1)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["sv_idle"]).to eq("5")
|
||||
end
|
||||
end
|
||||
|
||||
context "client fail to checkout connection from the pool" do
|
||||
it "counts clients as idle" do
|
||||
new_configs = processes.pgcat.current_config
|
||||
new_configs["general"]["connect_timeout"] = 500
|
||||
new_configs["general"]["ban_time"] = 1
|
||||
new_configs["general"]["shutdown_timeout"] = 1
|
||||
new_configs["pools"]["sharded_db"]["users"]["0"]["pool_size"] = 1
|
||||
processes.pgcat.update_config(new_configs)
|
||||
processes.pgcat.reload_config
|
||||
|
||||
threads = []
|
||||
connections = Array.new(5) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
threads << Thread.new { c.async_exec("SELECT pg_sleep(1)") rescue PG::SystemError }
|
||||
end
|
||||
|
||||
sleep(2)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["cl_idle"]).to eq("5")
|
||||
expect(results["sv_idle"]).to eq("1")
|
||||
|
||||
threads.map(&:join)
|
||||
connections.map(&:close)
|
||||
end
|
||||
end
|
||||
|
||||
context "clients connects and disconnect normally" do
|
||||
let(:processes) { Helpers::Pgcat.single_instance_setup("sharded_db", 2) }
|
||||
|
||||
it 'shows the same number of clients before and after' do
|
||||
clients_before = clients_connected_to_pool(processes: processes)
|
||||
threads = []
|
||||
connections = Array.new(4) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
threads << Thread.new { c.async_exec("SELECT 1") }
|
||||
end
|
||||
clients_between = clients_connected_to_pool(processes: processes)
|
||||
expect(clients_before).not_to eq(clients_between)
|
||||
connections.each(&:close)
|
||||
clients_after = clients_connected_to_pool(processes: processes)
|
||||
expect(clients_before).to eq(clients_after)
|
||||
end
|
||||
end
|
||||
|
||||
context "clients connects and disconnect abruptly" do
|
||||
let(:processes) { Helpers::Pgcat.single_instance_setup("sharded_db", 10) }
|
||||
|
||||
it 'shows the same number of clients before and after' do
|
||||
threads = []
|
||||
connections = Array.new(2) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
threads << Thread.new { c.async_exec("SELECT 1") }
|
||||
end
|
||||
clients_before = clients_connected_to_pool(processes: processes)
|
||||
random_string = (0...8).map { (65 + rand(26)).chr }.join
|
||||
connection_string = "#{pgcat_conn_str}?application_name=#{random_string}"
|
||||
faulty_client = Process.spawn("psql -Atx #{connection_string} >/dev/null")
|
||||
sleep(1)
|
||||
# psql starts two processes, we only know the pid of the parent, this
|
||||
# ensure both are killed
|
||||
`pkill -9 -f '#{random_string}'`
|
||||
Process.wait(faulty_client)
|
||||
clients_after = clients_connected_to_pool(processes: processes)
|
||||
expect(clients_before).to eq(clients_after)
|
||||
end
|
||||
end
|
||||
|
||||
context "clients overwhelm server pools" do
|
||||
let(:processes) { Helpers::Pgcat.single_instance_setup("sharded_db", 2) }
|
||||
|
||||
it "cl_waiting is updated to show it" do
|
||||
threads = []
|
||||
connections = Array.new(4) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
threads << Thread.new { c.async_exec("SELECT pg_sleep(1.5)") }
|
||||
end
|
||||
|
||||
sleep(1.1) # Allow time for stats to update
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_cancel_req sv_idle sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
|
||||
expect(results["cl_waiting"]).to eq("2")
|
||||
expect(results["cl_active"]).to eq("2")
|
||||
expect(results["sv_active"]).to eq("2")
|
||||
|
||||
sleep(2.5) # Allow time for stats to update
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["cl_idle"]).to eq("4")
|
||||
expect(results["sv_idle"]).to eq("2")
|
||||
|
||||
threads.map(&:join)
|
||||
connections.map(&:close)
|
||||
end
|
||||
|
||||
it "show correct max_wait" do
|
||||
threads = []
|
||||
connections = Array.new(4) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
threads << Thread.new { c.async_exec("SELECT pg_sleep(1.5)") }
|
||||
end
|
||||
|
||||
sleep(2.5) # Allow time for stats to update
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
|
||||
expect(results["maxwait"]).to eq("1")
|
||||
expect(results["maxwait_us"].to_i).to be_within(200_000).of(500_000)
|
||||
|
||||
sleep(4.5) # Allow time for stats to update
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
expect(results["maxwait"]).to eq("0")
|
||||
|
||||
threads.map(&:join)
|
||||
connections.map(&:close)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
describe "SHOW CLIENTS" do
|
||||
it "reports correct number and application names" do
|
||||
conn_str = processes.pgcat.connection_string("sharded_db", "sharding_user")
|
||||
connections = Array.new(20) { |i| PG::connect("#{conn_str}?application_name=app#{i % 5}") }
|
||||
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
sleep(1) # Wait for stats to be updated
|
||||
|
||||
results = admin_conn.async_exec("SHOW CLIENTS")
|
||||
expect(results.count).to eq(21) # count admin clients
|
||||
expect(results.select { |c| c["application_name"] == "app3" || c["application_name"] == "app4" }.count).to eq(8)
|
||||
expect(results.select { |c| c["database"] == "pgcat" }.count).to eq(1)
|
||||
|
||||
connections[0..5].map(&:close)
|
||||
sleep(1) # Wait for stats to be updated
|
||||
results = admin_conn.async_exec("SHOW CLIENTS")
|
||||
expect(results.count).to eq(15)
|
||||
|
||||
connections[6..].map(&:close)
|
||||
sleep(1) # Wait for stats to be updated
|
||||
expect(admin_conn.async_exec("SHOW CLIENTS").count).to eq(1)
|
||||
admin_conn.close
|
||||
end
|
||||
|
||||
it "reports correct number of queries and transactions" do
|
||||
conn_str = processes.pgcat.connection_string("sharded_db", "sharding_user")
|
||||
|
||||
connections = Array.new(2) { |i| PG::connect("#{conn_str}?application_name=app#{i}") }
|
||||
connections.each do |c|
|
||||
c.async_exec("SELECT 1")
|
||||
c.async_exec("SELECT 2")
|
||||
c.async_exec("SELECT 3")
|
||||
c.async_exec("BEGIN")
|
||||
c.async_exec("SELECT 4")
|
||||
c.async_exec("SELECT 5")
|
||||
c.async_exec("COMMIT")
|
||||
end
|
||||
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
sleep(1) # Wait for stats to be updated
|
||||
|
||||
results = admin_conn.async_exec("SHOW CLIENTS")
|
||||
expect(results.count).to eq(3)
|
||||
normal_client_results = results.reject { |r| r["database"] == "pgcat" }
|
||||
expect(normal_client_results[0]["transaction_count"]).to eq("4")
|
||||
expect(normal_client_results[1]["transaction_count"]).to eq("4")
|
||||
expect(normal_client_results[0]["query_count"]).to eq("7")
|
||||
expect(normal_client_results[1]["query_count"]).to eq("7")
|
||||
|
||||
admin_conn.close
|
||||
connections.map(&:close)
|
||||
end
|
||||
end
|
||||
|
||||
describe "Manual Banning" do
|
||||
let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 10) }
|
||||
before do
|
||||
@@ -81,7 +398,7 @@ describe "Admin" do
|
||||
end
|
||||
end
|
||||
|
||||
describe "SHOW USERS" do
|
||||
describe "SHOW users" do
|
||||
it "returns the right users" do
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW USERS")[0]
|
||||
|
||||
@@ -41,24 +41,7 @@ module Helpers
|
||||
"1" => { "database" => "shard1", "servers" => [["localhost", primary1.port.to_s, "primary"]] },
|
||||
"2" => { "database" => "shard2", "servers" => [["localhost", primary2.port.to_s, "primary"]] },
|
||||
},
|
||||
"users" => { "0" => user },
|
||||
"plugins" => {
|
||||
"intercept" => {
|
||||
"enabled" => true,
|
||||
"queries" => {
|
||||
"0" => {
|
||||
"query" => "select current_database() as a, current_schemas(false) as b",
|
||||
"schema" => [
|
||||
["a", "text"],
|
||||
["b", "text"],
|
||||
],
|
||||
"result" => [
|
||||
["${DATABASE}", "{public}"],
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"users" => { "0" => user }
|
||||
}
|
||||
}
|
||||
pgcat.update_config(pgcat_cfg)
|
||||
@@ -118,7 +101,7 @@ module Helpers
|
||||
end
|
||||
end
|
||||
|
||||
def self.single_shard_setup(pool_name, pool_size, pool_mode="transaction", lb_mode="random", log_level="info", pool_settings={})
|
||||
def self.single_shard_setup(pool_name, pool_size, pool_mode="transaction", lb_mode="random", log_level="info")
|
||||
user = {
|
||||
"password" => "sharding_user",
|
||||
"pool_size" => pool_size,
|
||||
@@ -134,32 +117,28 @@ module Helpers
|
||||
replica1 = PgInstance.new(8432, user["username"], user["password"], "shard0")
|
||||
replica2 = PgInstance.new(9432, user["username"], user["password"], "shard0")
|
||||
|
||||
pool_config = {
|
||||
"default_role" => "any",
|
||||
"pool_mode" => pool_mode,
|
||||
"load_balancing_mode" => lb_mode,
|
||||
"primary_reads_enabled" => false,
|
||||
"query_parser_enabled" => false,
|
||||
"sharding_function" => "pg_bigint_hash",
|
||||
"shards" => {
|
||||
"0" => {
|
||||
"database" => "shard0",
|
||||
"servers" => [
|
||||
["localhost", primary.port.to_s, "primary"],
|
||||
["localhost", replica0.port.to_s, "replica"],
|
||||
["localhost", replica1.port.to_s, "replica"],
|
||||
["localhost", replica2.port.to_s, "replica"]
|
||||
]
|
||||
},
|
||||
},
|
||||
"users" => { "0" => user }
|
||||
}
|
||||
|
||||
pool_config = pool_config.merge(pool_settings)
|
||||
|
||||
# Main proxy configs
|
||||
pgcat_cfg["pools"] = {
|
||||
"#{pool_name}" => pool_config,
|
||||
"#{pool_name}" => {
|
||||
"default_role" => "any",
|
||||
"pool_mode" => pool_mode,
|
||||
"load_balancing_mode" => lb_mode,
|
||||
"primary_reads_enabled" => false,
|
||||
"query_parser_enabled" => false,
|
||||
"sharding_function" => "pg_bigint_hash",
|
||||
"shards" => {
|
||||
"0" => {
|
||||
"database" => "shard0",
|
||||
"servers" => [
|
||||
["localhost", primary.port.to_s, "primary"],
|
||||
["localhost", replica0.port.to_s, "replica"],
|
||||
["localhost", replica1.port.to_s, "replica"],
|
||||
["localhost", replica2.port.to_s, "replica"]
|
||||
]
|
||||
},
|
||||
},
|
||||
"users" => { "0" => user }
|
||||
}
|
||||
}
|
||||
pgcat_cfg["general"]["port"] = pgcat.port
|
||||
pgcat.update_config(pgcat_cfg)
|
||||
|
||||
@@ -241,18 +241,6 @@ describe "Miscellaneous" do
|
||||
|
||||
expect(processes.primary.count_query("DISCARD ALL")).to eq(10)
|
||||
end
|
||||
|
||||
it "Resets server roles correctly" do
|
||||
10.times do
|
||||
conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
|
||||
conn.async_exec("SET SERVER ROLE to 'primary'")
|
||||
conn.async_exec("SELECT 1")
|
||||
conn.async_exec("SET statement_timeout to 5000")
|
||||
conn.close
|
||||
end
|
||||
|
||||
expect(processes.primary.count_query("RESET ROLE")).to eq(10)
|
||||
end
|
||||
end
|
||||
|
||||
context "transaction mode" do
|
||||
@@ -320,31 +308,6 @@ describe "Miscellaneous" do
|
||||
expect(processes.primary.count_query("DISCARD ALL")).to eq(0)
|
||||
end
|
||||
end
|
||||
|
||||
context "server cleanup disabled" do
|
||||
let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 1, "transaction", "random", "info", { "cleanup_server_connections" => false }) }
|
||||
|
||||
it "will not clean up connection state" do
|
||||
conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
|
||||
processes.primary.reset_stats
|
||||
conn.async_exec("SET statement_timeout TO 1000")
|
||||
conn.close
|
||||
|
||||
puts processes.pgcat.logs
|
||||
expect(processes.primary.count_query("DISCARD ALL")).to eq(0)
|
||||
end
|
||||
|
||||
it "will not clean up prepared statements" do
|
||||
conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
|
||||
processes.primary.reset_stats
|
||||
conn.async_exec("PREPARE prepared_q (int) AS SELECT $1")
|
||||
|
||||
conn.close
|
||||
|
||||
puts processes.pgcat.logs
|
||||
expect(processes.primary.count_query("DISCARD ALL")).to eq(0)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
describe "Idle client timeout" do
|
||||
|
||||
@@ -1,14 +0,0 @@
|
||||
require_relative 'spec_helper'
|
||||
|
||||
|
||||
describe "Plugins" do
|
||||
let(:processes) { Helpers::Pgcat.three_shard_setup("sharded_db", 5) }
|
||||
|
||||
context "intercept" do
|
||||
it "will intercept an intellij query" do
|
||||
conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
|
||||
res = conn.exec("select current_database() as a, current_schemas(false) as b")
|
||||
expect(res.values).to eq([["sharded_db", "{public}"]])
|
||||
end
|
||||
end
|
||||
end
|
||||
@@ -1,369 +0,0 @@
|
||||
# frozen_string_literal: true
|
||||
require 'open3'
|
||||
require_relative 'spec_helper'
|
||||
|
||||
describe "Stats" do
|
||||
let(:processes) { Helpers::Pgcat.single_instance_setup("sharded_db", 10) }
|
||||
let(:pgcat_conn_str) { processes.pgcat.connection_string("sharded_db", "sharding_user") }
|
||||
|
||||
after do
|
||||
processes.all_databases.map(&:reset)
|
||||
processes.pgcat.shutdown
|
||||
end
|
||||
|
||||
describe "SHOW STATS" do
|
||||
context "clients connect and make one query" do
|
||||
it "updates *_query_time and *_wait_time" do
|
||||
connections = Array.new(3) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
Thread.new { c.async_exec("SELECT pg_sleep(0.25)") }
|
||||
end
|
||||
sleep(1)
|
||||
connections.map(&:close)
|
||||
|
||||
# wait for averages to be calculated, we shouldn't do this too often
|
||||
sleep(15.5)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW STATS")[0]
|
||||
admin_conn.close
|
||||
expect(results["total_query_time"].to_i).to be_within(200).of(750)
|
||||
expect(results["avg_query_time"].to_i).to be_within(50).of(250)
|
||||
|
||||
expect(results["total_wait_time"].to_i).to_not eq(0)
|
||||
expect(results["avg_wait_time"].to_i).to_not eq(0)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
describe "SHOW POOLS" do
|
||||
context "bad credentials" do
|
||||
it "does not change any stats" do
|
||||
bad_password_url = URI(pgcat_conn_str)
|
||||
bad_password_url.password = "wrong"
|
||||
expect { PG::connect("#{bad_password_url.to_s}?application_name=bad_password") }.to raise_error(PG::ConnectionBad)
|
||||
|
||||
sleep(1)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
|
||||
expect(results["sv_idle"]).to eq("1")
|
||||
end
|
||||
end
|
||||
|
||||
context "bad database name" do
|
||||
it "does not change any stats" do
|
||||
bad_db_url = URI(pgcat_conn_str)
|
||||
bad_db_url.path = "/wrong_db"
|
||||
expect { PG::connect("#{bad_db_url.to_s}?application_name=bad_db") }.to raise_error(PG::ConnectionBad)
|
||||
|
||||
sleep(1)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
|
||||
expect(results["sv_idle"]).to eq("1")
|
||||
end
|
||||
end
|
||||
|
||||
context "client connects but issues no queries" do
|
||||
it "only affects cl_idle stats" do
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
|
||||
before_test = admin_conn.async_exec("SHOW POOLS")[0]["sv_idle"]
|
||||
connections = Array.new(20) { PG::connect(pgcat_conn_str) }
|
||||
sleep(1)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["cl_idle"]).to eq("20")
|
||||
expect(results["sv_idle"]).to eq(before_test)
|
||||
|
||||
connections.map(&:close)
|
||||
sleep(1.1)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_active cl_idle cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["sv_idle"]).to eq(before_test)
|
||||
end
|
||||
end
|
||||
|
||||
context "clients connect and make one query" do
|
||||
it "only affects cl_idle, sv_idle stats" do
|
||||
connections = Array.new(5) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
Thread.new { c.async_exec("SELECT pg_sleep(2.5)") }
|
||||
end
|
||||
|
||||
sleep(1.1)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_waiting cl_cancel_req sv_idle sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["cl_active"]).to eq("5")
|
||||
expect(results["sv_active"]).to eq("5")
|
||||
|
||||
sleep(3)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["cl_idle"]).to eq("5")
|
||||
expect(results["sv_idle"]).to eq("5")
|
||||
|
||||
connections.map(&:close)
|
||||
sleep(1)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["sv_idle"]).to eq("5")
|
||||
end
|
||||
end
|
||||
|
||||
context "client connects and opens a transaction and closes connection uncleanly" do
|
||||
it "produces correct statistics" do
|
||||
connections = Array.new(5) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
Thread.new do
|
||||
c.async_exec("BEGIN")
|
||||
c.async_exec("SELECT pg_sleep(0.01)")
|
||||
c.close
|
||||
end
|
||||
end
|
||||
|
||||
sleep(1.1)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["sv_idle"]).to eq("5")
|
||||
end
|
||||
end
|
||||
|
||||
context "client fail to checkout connection from the pool" do
|
||||
it "counts clients as idle" do
|
||||
new_configs = processes.pgcat.current_config
|
||||
new_configs["general"]["connect_timeout"] = 500
|
||||
new_configs["general"]["ban_time"] = 1
|
||||
new_configs["general"]["shutdown_timeout"] = 1
|
||||
new_configs["pools"]["sharded_db"]["users"]["0"]["pool_size"] = 1
|
||||
processes.pgcat.update_config(new_configs)
|
||||
processes.pgcat.reload_config
|
||||
|
||||
threads = []
|
||||
connections = Array.new(5) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
threads << Thread.new { c.async_exec("SELECT pg_sleep(1)") rescue PG::SystemError }
|
||||
end
|
||||
|
||||
sleep(2)
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["cl_idle"]).to eq("5")
|
||||
expect(results["sv_idle"]).to eq("1")
|
||||
|
||||
threads.map(&:join)
|
||||
connections.map(&:close)
|
||||
end
|
||||
end
|
||||
|
||||
context "clients connects and disconnect normally" do
|
||||
let(:processes) { Helpers::Pgcat.single_instance_setup("sharded_db", 2) }
|
||||
|
||||
it 'shows the same number of clients before and after' do
|
||||
clients_before = clients_connected_to_pool(processes: processes)
|
||||
threads = []
|
||||
connections = Array.new(4) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
threads << Thread.new { c.async_exec("SELECT 1") rescue nil }
|
||||
end
|
||||
clients_between = clients_connected_to_pool(processes: processes)
|
||||
expect(clients_before).not_to eq(clients_between)
|
||||
connections.each(&:close)
|
||||
clients_after = clients_connected_to_pool(processes: processes)
|
||||
expect(clients_before).to eq(clients_after)
|
||||
end
|
||||
end
|
||||
|
||||
context "clients connects and disconnect abruptly" do
|
||||
let(:processes) { Helpers::Pgcat.single_instance_setup("sharded_db", 10) }
|
||||
|
||||
it 'shows the same number of clients before and after' do
|
||||
threads = []
|
||||
connections = Array.new(2) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
threads << Thread.new { c.async_exec("SELECT 1") }
|
||||
end
|
||||
clients_before = clients_connected_to_pool(processes: processes)
|
||||
random_string = (0...8).map { (65 + rand(26)).chr }.join
|
||||
connection_string = "#{pgcat_conn_str}?application_name=#{random_string}"
|
||||
faulty_client = Process.spawn("psql -Atx #{connection_string} >/dev/null")
|
||||
sleep(1)
|
||||
# psql starts two processes, we only know the pid of the parent, this
|
||||
# ensure both are killed
|
||||
`pkill -9 -f '#{random_string}'`
|
||||
Process.wait(faulty_client)
|
||||
clients_after = clients_connected_to_pool(processes: processes)
|
||||
expect(clients_before).to eq(clients_after)
|
||||
end
|
||||
end
|
||||
|
||||
context "clients overwhelm server pools" do
|
||||
let(:processes) { Helpers::Pgcat.single_instance_setup("sharded_db", 2) }
|
||||
|
||||
it "cl_waiting is updated to show it" do
|
||||
threads = []
|
||||
connections = Array.new(4) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
threads << Thread.new { c.async_exec("SELECT pg_sleep(1.5)") }
|
||||
end
|
||||
|
||||
sleep(1.1) # Allow time for stats to update
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_cancel_req sv_idle sv_used sv_tested sv_login maxwait].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
|
||||
expect(results["cl_waiting"]).to eq("2")
|
||||
expect(results["cl_active"]).to eq("2")
|
||||
expect(results["sv_active"]).to eq("2")
|
||||
|
||||
sleep(2.5) # Allow time for stats to update
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
expect(results["cl_idle"]).to eq("4")
|
||||
expect(results["sv_idle"]).to eq("2")
|
||||
|
||||
threads.map(&:join)
|
||||
connections.map(&:close)
|
||||
end
|
||||
|
||||
it "show correct max_wait" do
|
||||
threads = []
|
||||
connections = Array.new(4) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
|
||||
connections.each do |c|
|
||||
threads << Thread.new { c.async_exec("SELECT pg_sleep(1.5)") rescue nil }
|
||||
end
|
||||
|
||||
sleep(2.5) # Allow time for stats to update
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
|
||||
expect(results["maxwait"]).to eq("1")
|
||||
expect(results["maxwait_us"].to_i).to be_within(200_000).of(500_000)
|
||||
connections.map(&:close)
|
||||
|
||||
sleep(4.5) # Allow time for stats to update
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
expect(results["maxwait"]).to eq("0")
|
||||
|
||||
threads.map(&:join)
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
describe "SHOW CLIENTS" do
|
||||
it "reports correct number and application names" do
|
||||
conn_str = processes.pgcat.connection_string("sharded_db", "sharding_user")
|
||||
connections = Array.new(20) { |i| PG::connect("#{conn_str}?application_name=app#{i % 5}") }
|
||||
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
sleep(1) # Wait for stats to be updated
|
||||
|
||||
results = admin_conn.async_exec("SHOW CLIENTS")
|
||||
expect(results.count).to eq(21) # count admin clients
|
||||
expect(results.select { |c| c["application_name"] == "app3" || c["application_name"] == "app4" }.count).to eq(8)
|
||||
expect(results.select { |c| c["database"] == "pgcat" }.count).to eq(1)
|
||||
|
||||
connections[0..5].map(&:close)
|
||||
sleep(1) # Wait for stats to be updated
|
||||
results = admin_conn.async_exec("SHOW CLIENTS")
|
||||
expect(results.count).to eq(15)
|
||||
|
||||
connections[6..].map(&:close)
|
||||
sleep(1) # Wait for stats to be updated
|
||||
expect(admin_conn.async_exec("SHOW CLIENTS").count).to eq(1)
|
||||
admin_conn.close
|
||||
end
|
||||
|
||||
it "reports correct number of queries and transactions" do
|
||||
conn_str = processes.pgcat.connection_string("sharded_db", "sharding_user")
|
||||
|
||||
connections = Array.new(2) { |i| PG::connect("#{conn_str}?application_name=app#{i}") }
|
||||
connections.each do |c|
|
||||
c.async_exec("SELECT 1")
|
||||
c.async_exec("SELECT 2")
|
||||
c.async_exec("SELECT 3")
|
||||
c.async_exec("BEGIN")
|
||||
c.async_exec("SELECT 4")
|
||||
c.async_exec("SELECT 5")
|
||||
c.async_exec("COMMIT")
|
||||
end
|
||||
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
sleep(1) # Wait for stats to be updated
|
||||
|
||||
results = admin_conn.async_exec("SHOW CLIENTS")
|
||||
expect(results.count).to eq(3)
|
||||
normal_client_results = results.reject { |r| r["database"] == "pgcat" }
|
||||
expect(normal_client_results[0]["transaction_count"]).to eq("4")
|
||||
expect(normal_client_results[1]["transaction_count"]).to eq("4")
|
||||
expect(normal_client_results[0]["query_count"]).to eq("7")
|
||||
expect(normal_client_results[1]["query_count"]).to eq("7")
|
||||
|
||||
admin_conn.close
|
||||
connections.map(&:close)
|
||||
end
|
||||
end
|
||||
|
||||
|
||||
describe "Query Storm" do
|
||||
context "when the proxy receives overwhelmingly large number of short quick queries" do
|
||||
it "should not have lingering clients or active servers" do
|
||||
new_configs = processes.pgcat.current_config
|
||||
|
||||
new_configs["general"]["connect_timeout"] = 500
|
||||
new_configs["general"]["ban_time"] = 1
|
||||
new_configs["general"]["shutdown_timeout"] = 1
|
||||
new_configs["pools"]["sharded_db"]["users"]["0"]["pool_size"] = 1
|
||||
processes.pgcat.update_config(new_configs)
|
||||
processes.pgcat.reload_config
|
||||
|
||||
Array.new(40) do
|
||||
Thread.new do
|
||||
conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
|
||||
conn.async_exec("SELECT pg_sleep(0.1)")
|
||||
rescue PG::SystemError
|
||||
ensure
|
||||
conn.close
|
||||
end
|
||||
end.each(&:join)
|
||||
|
||||
sleep 1
|
||||
|
||||
admin_conn = PG::connect(processes.pgcat.admin_connection_string)
|
||||
results = admin_conn.async_exec("SHOW POOLS")[0]
|
||||
%w[cl_idle cl_waiting cl_cancel_req sv_used sv_tested sv_login].each do |s|
|
||||
raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
|
||||
end
|
||||
|
||||
admin_conn.close
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
1
tests/rust/.gitignore
vendored
1
tests/rust/.gitignore
vendored
@@ -1 +0,0 @@
|
||||
target/
|
||||
1322
tests/rust/Cargo.lock
generated
1322
tests/rust/Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
@@ -1,10 +0,0 @@
|
||||
[package]
|
||||
name = "rust"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||
|
||||
[dependencies]
|
||||
sqlx = { version = "0.6.2", features = [ "runtime-tokio-rustls", "postgres", "json", "tls", "migrate", "time", "uuid", "ipnetwork"] }
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
@@ -1,29 +0,0 @@
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
test_prepared_statements().await;
|
||||
}
|
||||
|
||||
async fn test_prepared_statements() {
|
||||
let pool = sqlx::postgres::PgPoolOptions::new()
|
||||
.max_connections(5)
|
||||
.connect("postgres://sharding_user:sharding_user@127.0.0.1:6432/sharded_db")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let mut handles = Vec::new();
|
||||
|
||||
for _ in 0..5 {
|
||||
let pool = pool.clone();
|
||||
let handle = tokio::task::spawn(async move {
|
||||
for _ in 0..1000 {
|
||||
sqlx::query("SELECT 1").fetch_all(&pool).await.unwrap();
|
||||
}
|
||||
});
|
||||
|
||||
handles.push(handle);
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
handle.await.unwrap();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user