diff options
Diffstat (limited to 'lib/sqlalchemy')
| -rw-r--r-- | lib/sqlalchemy/dialects/postgresql/asyncpg.py | 17 | ||||
| -rw-r--r-- | lib/sqlalchemy/dialects/postgresql/pg8000.py | 4 | ||||
| -rw-r--r-- | lib/sqlalchemy/dialects/postgresql/psycopg.py | 4 | ||||
| -rw-r--r-- | lib/sqlalchemy/orm/evaluator.py | 68 | ||||
| -rw-r--r-- | lib/sqlalchemy/testing/suite/test_dialect.py | 11 |
5 files changed, 85 insertions, 19 deletions
diff --git a/lib/sqlalchemy/dialects/postgresql/asyncpg.py b/lib/sqlalchemy/dialects/postgresql/asyncpg.py index a84bece4f..4cc04d20a 100644 --- a/lib/sqlalchemy/dialects/postgresql/asyncpg.py +++ b/lib/sqlalchemy/dialects/postgresql/asyncpg.py @@ -741,6 +741,12 @@ class AsyncAdapt_asyncpg_connection(AdaptedConnection): else: self.isolation_level = self._isolation_setting + def ping(self): + try: + _ = self.await_(self._connection.fetchrow(";")) + except Exception as error: + self._handle_exception(error) + def set_isolation_level(self, level): if self._started: self.rollback() @@ -996,6 +1002,17 @@ class PGDialect_asyncpg(PGDialect): util.coerce_kw_type(opts, "port", int) return ([], opts) + def do_ping(self, dbapi_connection): + try: + dbapi_connection.ping() + except self.dbapi.Error as err: + if self.is_disconnect(err, dbapi_connection, None): + return False + else: + raise + else: + return True + @classmethod def get_pool_class(cls, url): diff --git a/lib/sqlalchemy/dialects/postgresql/pg8000.py b/lib/sqlalchemy/dialects/postgresql/pg8000.py index ce9a3bb6c..d4b159b73 100644 --- a/lib/sqlalchemy/dialects/postgresql/pg8000.py +++ b/lib/sqlalchemy/dialects/postgresql/pg8000.py @@ -559,5 +559,9 @@ class PGDialect_pg8000(PGDialect): else: return None + @util.memoized_property + def _dialect_specific_select_one(self): + return ";" + dialect = PGDialect_pg8000 diff --git a/lib/sqlalchemy/dialects/postgresql/psycopg.py b/lib/sqlalchemy/dialects/postgresql/psycopg.py index 633357a74..371bf2bc2 100644 --- a/lib/sqlalchemy/dialects/postgresql/psycopg.py +++ b/lib/sqlalchemy/dialects/postgresql/psycopg.py @@ -505,6 +505,10 @@ class PGDialect_psycopg(_PGDialect_common_psycopg): else: self.do_commit(connection.connection) + @util.memoized_property + def _dialect_specific_select_one(self): + return ";" + class AsyncAdapt_psycopg_cursor: __slots__ = ("_cursor", "await_", "_rows") diff --git a/lib/sqlalchemy/orm/evaluator.py b/lib/sqlalchemy/orm/evaluator.py index 72936d1ab..b3129afdd 100644 --- a/lib/sqlalchemy/orm/evaluator.py +++ b/lib/sqlalchemy/orm/evaluator.py @@ -16,6 +16,8 @@ from .. import inspect from .. import util from ..sql import and_ from ..sql import operators +from ..sql.sqltypes import Integer +from ..sql.sqltypes import Numeric class UnevaluatableError(exc.InvalidRequestError): @@ -120,7 +122,7 @@ class EvaluatorCompiler: dispatch = f"visit_{clause.operator.__name__.rstrip('_')}_binary_op" meth = getattr(self, dispatch, None) if meth: - return meth(clause.operator, eval_left, eval_right) + return meth(clause.operator, eval_left, eval_right, clause) else: raise UnevaluatableError( f"Cannot evaluate {type(clause).__name__} with " @@ -165,9 +167,13 @@ class EvaluatorCompiler: return evaluate - def visit_custom_op_binary_op(self, operator, eval_left, eval_right): + def visit_custom_op_binary_op( + self, operator, eval_left, eval_right, clause + ): if operator.python_impl: - return self._straight_evaluate(operator, eval_left, eval_right) + return self._straight_evaluate( + operator, eval_left, eval_right, clause + ) else: raise UnevaluatableError( f"Custom operator {operator.opstring!r} can't be evaluated " @@ -175,19 +181,19 @@ class EvaluatorCompiler: "`.python_impl`." ) - def visit_is_binary_op(self, operator, eval_left, eval_right): + def visit_is_binary_op(self, operator, eval_left, eval_right, clause): def evaluate(obj): return eval_left(obj) == eval_right(obj) return evaluate - def visit_is_not_binary_op(self, operator, eval_left, eval_right): + def visit_is_not_binary_op(self, operator, eval_left, eval_right, clause): def evaluate(obj): return eval_left(obj) != eval_right(obj) return evaluate - def _straight_evaluate(self, operator, eval_left, eval_right): + def _straight_evaluate(self, operator, eval_left, eval_right, clause): def evaluate(obj): left_val = eval_left(obj) right_val = eval_right(obj) @@ -197,11 +203,25 @@ class EvaluatorCompiler: return evaluate - visit_add_binary_op = _straight_evaluate - visit_mul_binary_op = _straight_evaluate - visit_sub_binary_op = _straight_evaluate - visit_mod_binary_op = _straight_evaluate - visit_truediv_binary_op = _straight_evaluate + def _straight_evaluate_numeric_only( + self, operator, eval_left, eval_right, clause + ): + if clause.left.type._type_affinity not in ( + Numeric, + Integer, + ) or clause.right.type._type_affinity not in (Numeric, Integer): + raise UnevaluatableError( + f'Cannot evaluate math operator "{operator.__name__}" for ' + f"datatypes {clause.left.type}, {clause.right.type}" + ) + + return self._straight_evaluate(operator, eval_left, eval_right, clause) + + visit_add_binary_op = _straight_evaluate_numeric_only + visit_mul_binary_op = _straight_evaluate_numeric_only + visit_sub_binary_op = _straight_evaluate_numeric_only + visit_mod_binary_op = _straight_evaluate_numeric_only + visit_truediv_binary_op = _straight_evaluate_numeric_only visit_lt_binary_op = _straight_evaluate visit_le_binary_op = _straight_evaluate visit_ne_binary_op = _straight_evaluate @@ -209,33 +229,43 @@ class EvaluatorCompiler: visit_ge_binary_op = _straight_evaluate visit_eq_binary_op = _straight_evaluate - def visit_in_op_binary_op(self, operator, eval_left, eval_right): + def visit_in_op_binary_op(self, operator, eval_left, eval_right, clause): return self._straight_evaluate( lambda a, b: a in b if a is not _NO_OBJECT else None, eval_left, eval_right, + clause, ) - def visit_not_in_op_binary_op(self, operator, eval_left, eval_right): + def visit_not_in_op_binary_op( + self, operator, eval_left, eval_right, clause + ): return self._straight_evaluate( lambda a, b: a not in b if a is not _NO_OBJECT else None, eval_left, eval_right, + clause, ) - def visit_concat_op_binary_op(self, operator, eval_left, eval_right): + def visit_concat_op_binary_op( + self, operator, eval_left, eval_right, clause + ): return self._straight_evaluate( - lambda a, b: a + b, eval_left, eval_right + lambda a, b: a + b, eval_left, eval_right, clause ) - def visit_startswith_op_binary_op(self, operator, eval_left, eval_right): + def visit_startswith_op_binary_op( + self, operator, eval_left, eval_right, clause + ): return self._straight_evaluate( - lambda a, b: a.startswith(b), eval_left, eval_right + lambda a, b: a.startswith(b), eval_left, eval_right, clause ) - def visit_endswith_op_binary_op(self, operator, eval_left, eval_right): + def visit_endswith_op_binary_op( + self, operator, eval_left, eval_right, clause + ): return self._straight_evaluate( - lambda a, b: a.endswith(b), eval_left, eval_right + lambda a, b: a.endswith(b), eval_left, eval_right, clause ) def visit_unary(self, clause): diff --git a/lib/sqlalchemy/testing/suite/test_dialect.py b/lib/sqlalchemy/testing/suite/test_dialect.py index 55276e21b..bb2dd6574 100644 --- a/lib/sqlalchemy/testing/suite/test_dialect.py +++ b/lib/sqlalchemy/testing/suite/test_dialect.py @@ -8,6 +8,7 @@ from .. import config from .. import engines from .. import eq_ from .. import fixtures +from .. import is_true from .. import ne_ from .. import provide_metadata from ..assertions import expect_raises_message @@ -24,6 +25,16 @@ from ... import select from ... import String +class PingTest(fixtures.TestBase): + __backend__ = True + + def test_do_ping(self): + with testing.db.connect() as conn: + is_true( + testing.db.dialect.do_ping(conn.connection.dbapi_connection) + ) + + class ExceptionTest(fixtures.TablesTest): """Test basic exception wrapping. |