common_meta/error/retry_hint/
postgres.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use common_error::ext::RetryHint;
16
17/// Converts a tokio-postgres error into a conservative retry hint.
18pub fn retry_hint_from_postgres_error(error: &tokio_postgres::Error) -> RetryHint {
19    if error.is_closed() {
20        return RetryHint::Retryable;
21    }
22
23    retry_hint_from_postgres_sql_state(error.code())
24}
25
26/// Converts a Postgres SQLSTATE into a conservative retry hint.
27fn retry_hint_from_postgres_sql_state(
28    state: Option<&tokio_postgres::error::SqlState>,
29) -> RetryHint {
30    use tokio_postgres::error::SqlState;
31
32    let Some(state) = state else {
33        return RetryHint::NonRetryable;
34    };
35
36    // PostgreSQL SQLSTATE reference:
37    // https://www.postgresql.org/docs/current/errcodes-appendix.html
38    match state {
39        // 40001 serialization_failure: concurrent transaction serialization conflict.
40        &SqlState::T_R_SERIALIZATION_FAILURE
41        // 40P01 deadlock_detected: transaction deadlock.
42        | &SqlState::T_R_DEADLOCK_DETECTED
43        // 55P03 lock_not_available: lock could not be acquired now.
44        | &SqlState::LOCK_NOT_AVAILABLE
45        // 53300 too_many_connections: backend connection capacity exhausted.
46        | &SqlState::TOO_MANY_CONNECTIONS
47        // 57P01 admin_shutdown: server is shutting down by administrator request.
48        | &SqlState::ADMIN_SHUTDOWN
49        // 57P02 crash_shutdown: server is shutting down after crash.
50        | &SqlState::CRASH_SHUTDOWN
51        // 57P03 cannot_connect_now: server is not accepting connections now.
52        | &SqlState::CANNOT_CONNECT_NOW
53        // 08000 connection_exception: generic connection exception.
54        | &SqlState::CONNECTION_EXCEPTION
55        // 08001 sqlclient_unable_to_establish_sqlconnection: client could not establish connection.
56        | &SqlState::SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION
57        // 08003 connection_does_not_exist: connection does not exist.
58        | &SqlState::CONNECTION_DOES_NOT_EXIST
59        // 08004 sqlserver_rejected_establishment_of_sqlconnection: server rejected connection establishment.
60        | &SqlState::SQLSERVER_REJECTED_ESTABLISHMENT_OF_SQLCONNECTION
61        // 08006 connection_failure: connection failure.
62        | &SqlState::CONNECTION_FAILURE => RetryHint::Retryable,
63        _ => RetryHint::NonRetryable,
64    }
65}
66
67/// Converts a deadpool Postgres pool error into a conservative retry hint.
68pub fn retry_hint_from_postgres_pool_error(
69    error: &deadpool::managed::PoolError<tokio_postgres::Error>,
70) -> RetryHint {
71    match error {
72        deadpool::managed::PoolError::Timeout(_) => RetryHint::Retryable,
73        deadpool::managed::PoolError::Backend(error) => retry_hint_from_postgres_error(error),
74        deadpool::managed::PoolError::PostCreateHook(error) => match error {
75            deadpool::managed::HookError::Backend(error) => retry_hint_from_postgres_error(error),
76            deadpool::managed::HookError::Message(_) => RetryHint::NonRetryable,
77        },
78        deadpool::managed::PoolError::Closed | deadpool::managed::PoolError::NoRuntimeSpecified => {
79            RetryHint::NonRetryable
80        }
81    }
82}
83
84pub fn is_postgres_serialization_error(error: &tokio_postgres::Error) -> bool {
85    is_postgres_serialization_state(error.code())
86}
87
88fn is_postgres_serialization_state(state: Option<&tokio_postgres::error::SqlState>) -> bool {
89    use tokio_postgres::error::SqlState;
90
91    matches!(
92        state,
93        Some(&SqlState::T_R_SERIALIZATION_FAILURE | &SqlState::T_R_DEADLOCK_DETECTED)
94    )
95}
96
97#[cfg(test)]
98mod tests {
99    use common_error::ext::RetryHint;
100    use tokio_postgres::error::SqlState;
101
102    use super::*;
103
104    #[test]
105    fn test_postgres_sql_state_retry_hint() {
106        let retryable_states = [
107            &SqlState::T_R_SERIALIZATION_FAILURE,
108            &SqlState::T_R_DEADLOCK_DETECTED,
109            &SqlState::LOCK_NOT_AVAILABLE,
110            &SqlState::TOO_MANY_CONNECTIONS,
111            &SqlState::ADMIN_SHUTDOWN,
112            &SqlState::CRASH_SHUTDOWN,
113            &SqlState::CANNOT_CONNECT_NOW,
114            &SqlState::CONNECTION_EXCEPTION,
115            &SqlState::SQLCLIENT_UNABLE_TO_ESTABLISH_SQLCONNECTION,
116            &SqlState::CONNECTION_DOES_NOT_EXIST,
117            &SqlState::SQLSERVER_REJECTED_ESTABLISHMENT_OF_SQLCONNECTION,
118            &SqlState::CONNECTION_FAILURE,
119        ];
120
121        for state in retryable_states {
122            assert_eq!(
123                retry_hint_from_postgres_sql_state(Some(state)),
124                RetryHint::Retryable,
125                "SQLSTATE {} should be retryable",
126                state.code()
127            );
128        }
129
130        assert_eq!(
131            retry_hint_from_postgres_sql_state(Some(&SqlState::UNDEFINED_TABLE)),
132            RetryHint::NonRetryable
133        );
134        assert_eq!(
135            retry_hint_from_postgres_sql_state(None),
136            RetryHint::NonRetryable
137        );
138    }
139
140    #[test]
141    fn test_postgres_serialization_state() {
142        assert!(is_postgres_serialization_state(Some(
143            &SqlState::T_R_SERIALIZATION_FAILURE
144        )));
145        assert!(is_postgres_serialization_state(Some(
146            &SqlState::T_R_DEADLOCK_DETECTED
147        )));
148        assert!(!is_postgres_serialization_state(Some(
149            &SqlState::UNDEFINED_TABLE
150        )));
151        assert!(!is_postgres_serialization_state(None));
152    }
153}
common_meta/error/retry_hint/postgres.rs

common_meta/error/retry_hint/
postgres.rs
