diff --git a/drivine/connection/postgres/postgres_connection.py b/drivine/connection/postgres/postgres_connection.py index e8e1ae28..725875c9 100644 --- a/drivine/connection/postgres/postgres_connection.py +++ b/drivine/connection/postgres/postgres_connection.py @@ -10,6 +10,7 @@ from drivine.backend.sql.sql_json import to_json from drivine.connection.connection import Connection from drivine.connection.postgres.postgres_cursor import PostgresCursor +from drivine.connection.postgres.postgres_non_finite import read_non_finite from drivine.connection.postgres.postgres_result_mapper import PostgresResultMapper from drivine.connection.postgres.postgres_spec_compiler import PostgresSpecCompiler from drivine.cursor.cursor_specification import CursorSpecification @@ -30,7 +31,9 @@ class PostgresConnection(Connection): A connection to Postgres, through psycopg. Outside a transaction each statement commits itself, as a statement does on the graph engines. A statement without parameters is sent as it is, so it may hold several statements, or a block. A Cypher statement is compiled to SQL over the graph's tables (see SqlCypherCompiler), and gives a row as a - graph engine gives it: its one item, or a list of its items. + graph engine gives it: its one item, or a list of its items. A compiled statement that says it was given NaN or an + infinity has its cursor read the numbers that stand for them (see postgres_non_finite); every other is read as + psycopg reads it. """ def __init__(self, connection: Any, release: Any, session_id: str, graph: str = DEFAULT_GRAPH) -> None: @@ -48,6 +51,8 @@ async def query(self, spec: QuerySpecification[T]) -> list[T]: compiled = PostgresSpecCompiler(sql).compile() start = time.perf_counter() async with self.connection.cursor() as cursor: + if cypher is not None and cypher.non_finite: + read_non_finite(cursor) if prelude is not None: await cursor.execute(prelude) await cursor.execute(compiled.statement, _bindable(compiled.parameters)) @@ -74,6 +79,8 @@ async def open_cursor(self, spec: CursorSpecification[T]) -> PostgresCursor[T]: cursor = self.connection.cursor() else: cursor = self.connection.cursor(name=f"drivine_{next(_CURSORS)}", withhold=not self._in_transaction) + if cypher is not None and cypher.non_finite: + read_non_finite(cursor) await cursor.execute(compiled.statement, _bindable(compiled.parameters)) return PostgresCursor(finalized, cursor, self.result_mapper, compiled=sql is not finalized) diff --git a/drivine/connection/postgres/postgres_non_finite.py b/drivine/connection/postgres/postgres_non_finite.py new file mode 100644 index 00000000..9aec4018 --- /dev/null +++ b/drivine/connection/postgres/postgres_non_finite.py @@ -0,0 +1,81 @@ +""" +How the rows of a statement given NaN or an infinity are read on Postgres. `jsonb` holds neither, so drivine4j's +compiler writes each as a number past any double's: NaN as `to_jsonb(2e309 + 0.5)`, the infinities as +`to_jsonb(1e309 + 0.5)` and `to_jsonb(-1e309 - 0.5)`. Read as `json.loads` reads them, NaN would be an infinity. Such a +statement says so (`CompiledCypher.non_finite`), and its rows alone are read here, comparing each number's decimal text +exactly, as drivine4j's `PostgresCypherCompiler.read` does (liberation-data/drivine#673). Drivine's own compiler says it +of no statement. +""" + +from __future__ import annotations + +import json +import math +from decimal import Decimal, InvalidOperation +from typing import Any + +NOT_A_NUMBER = Decimal("2" + "0" * 309 + ".5") +"""The number that stands for NaN, 2e309 + 0.5, as drivine4j's compiler writes it on Postgres.""" + +INFINITE = Decimal("1e309") +"""Where the numbers that stand for the infinities begin, past any double's: `±(1e309 + 0.5)` are written for them.""" + +_LONGEST_FINITE_WHOLE = 309 +"""The most digits a whole number under INFINITE, 1 and 309 zeros, has: one with more may stand for an infinity.""" + + +def non_finite_of(text: str) -> float | None: + """ + The double the JSON number written `text`, in the rows of a statement given NaN or an infinity, stands for: NaN + where it is exactly the number that stands for NaN, an infinity where it is past any double's, else None, which is + the number it is. + """ + number = Decimal(text) + if number == NOT_A_NUMBER: + return math.nan + if number >= INFINITE: + return math.inf + if number <= -INFINITE: + return -math.inf + return None + + +def _float(text: str) -> float: + """A number with a fraction or an exponent: as `json.loads` reads it, unless it stands for NaN or an infinity.""" + number = float(text) + # A number float reads as finite is under any double's greatest, 1.8e308, so stands for neither. + if math.isfinite(number): + return number + try: + found = non_finite_of(text) + except InvalidOperation: + # Its exponent is past even Decimal's (`1e99999999999999999999`, as a `json` value may hold it): far past any + # double's, so the infinity float reads. + return number + return number if found is None else found + + +def _int(text: str) -> int | float: + """A whole number: as `json.loads` reads it, unless it is past any double's, an infinity.""" + if len(text.lstrip("-")) <= _LONGEST_FINITE_WHOLE: + return int(text) + found = non_finite_of(text) + return int(text) if found is None else found + + +def read_non_finite_json(data: str | bytes) -> Any: + """ + The JSON `data`, as Postgres writes a `json` or `jsonb` value, read as `json.loads` reads it, with each number that + stands for NaN or an infinity the double it is. A string is never read as a number, whatever it holds. + """ + return json.loads(data, parse_float=_float, parse_int=_int) + + +def read_non_finite(cursor: Any) -> None: + """ + Has `cursor` read `json` and `jsonb`, as text or binary, with read_non_finite_json. Set on that cursor alone, + never for the connection or the process: it would reach every statement, a user's own SQL and other libraries' too. + """ + from psycopg.types.json import set_json_loads + + set_json_loads(read_non_finite_json, cursor) diff --git a/drivine/cypher/sql/compiled_cypher.py b/drivine/cypher/sql/compiled_cypher.py index d091168b..ecf653b3 100644 --- a/drivine/cypher/sql/compiled_cypher.py +++ b/drivine/cypher/sql/compiled_cypher.py @@ -16,6 +16,10 @@ class CompiledCypher: `_nulls_equal`, which a store made by an earlier version, or by a user who may not make a function, does not have: see `OracleCypherCompiler.nulls_equal`. `calls_stored_era` says the same of the functions it renumbers a year that may be BC by, `_era` and `_era_date`: see `OracleTemporalSql.era`. + + `non_finite` says whether the statement was given NaN or an infinity, so its rows may hold the numbers that stand + for them, as drivine4j's compiler writes them on Postgres: see postgres_non_finite. Drivine's own compilers say + False of every statement, as they write NaN and the infinities as text (liberation-data/drivine#673). """ sql: str @@ -23,3 +27,4 @@ class CompiledCypher: prelude: str | None = None calls_stored_nulls_equal: bool = False calls_stored_era: bool = False + non_finite: bool = False diff --git a/drivine/sync/connection/postgres/postgres_connection.py b/drivine/sync/connection/postgres/postgres_connection.py index 9f7283c2..08618597 100644 --- a/drivine/sync/connection/postgres/postgres_connection.py +++ b/drivine/sync/connection/postgres/postgres_connection.py @@ -10,6 +10,7 @@ from drivine.backend.sql.sql_graph import DEFAULT_GRAPH from drivine.backend.sql.sql_json import to_json +from drivine.connection.postgres.postgres_non_finite import read_non_finite from drivine.connection.postgres.postgres_result_mapper import PostgresResultMapper from drivine.connection.postgres.postgres_spec_compiler import PostgresSpecCompiler from drivine.cursor.cursor_specification import CursorSpecification @@ -32,7 +33,9 @@ class PostgresConnection(Connection): A connection to Postgres, through psycopg. Outside a transaction each statement commits itself, as a statement does on the graph engines. A statement without parameters is sent as it is, so it may hold several statements, or a block. A Cypher statement is compiled to SQL over the graph's tables (see SqlCypherCompiler), and gives a row as a - graph engine gives it: its one item, or a list of its items. + graph engine gives it: its one item, or a list of its items. A compiled statement that says it was given NaN or an + infinity has its cursor read the numbers that stand for them (see postgres_non_finite); every other is read as + psycopg reads it. """ def __init__(self, connection: Any, release: Any, session_id: str, graph: str = DEFAULT_GRAPH) -> None: @@ -50,6 +53,8 @@ def query(self, spec: QuerySpecification[T]) -> list[T]: compiled = PostgresSpecCompiler(sql).compile() start = time.perf_counter() with self.connection.cursor() as cursor: + if cypher is not None and cypher.non_finite: + read_non_finite(cursor) if prelude is not None: cursor.execute(prelude) cursor.execute(compiled.statement, _bindable(compiled.parameters)) @@ -76,6 +81,8 @@ def open_cursor(self, spec: CursorSpecification[T]) -> PostgresCursor[T]: cursor = self.connection.cursor() else: cursor = self.connection.cursor(name=f"drivine_{next(_CURSORS)}", withhold=not self._in_transaction) + if cypher is not None and cypher.non_finite: + read_non_finite(cursor) cursor.execute(compiled.statement, _bindable(compiled.parameters)) return PostgresCursor(finalized, cursor, self.result_mapper, compiled=sql is not finalized) diff --git a/tests/integration/postgres/test_postgres_non_finite.py b/tests/integration/postgres/test_postgres_non_finite.py new file mode 100644 index 00000000..039253c2 --- /dev/null +++ b/tests/integration/postgres/test_postgres_non_finite.py @@ -0,0 +1,122 @@ +""" +The numbers that stand for NaN and the infinities in the rows of a statement given one, as drivine4j's compiler writes +them on Postgres, read back as the doubles they are where the statement says so, and as today where it doesn't +(liberation-data/drivine#673). +""" + +import copy +import math +import sys +from collections.abc import AsyncIterator +from typing import Any + +import pytest + +from drivine import CursorSpecification, DatabaseRegistry, QuerySpecification +from drivine.connection.connection import Connection +from drivine.connection.postgres.postgres_connection import PostgresConnection +from drivine.cypher.sql.compiled_cypher import CompiledCypher +from drivine.query.statement import Statement +from tests.integration.postgres.conftest import ON_POSTGRES + +pytestmark = ON_POSTGRES + +NUMBERS = ( + "jsonb_build_array(to_jsonb(2e309 + 0.5), to_jsonb(1e309 + 0.5), to_jsonb(-1e309 - 0.5), to_jsonb(1e308), " + "to_jsonb(2.5), to_jsonb('2' || repeat('0', 309) || '.5'))" +) +"""The numbers drivine4j's compiler writes for NaN and the infinities, two that stand for neither, and a string.""" + +STATEMENT = f"SELECT {NUMBERS} AS result" + +AS_TODAY = [math.inf, math.inf, -math.inf, 10**308, 2.5, "2" + "0" * 309 + ".5"] +"""As json.loads reads them: Postgres writes `numeric` 1e308 as its digits, a whole number.""" + + +def compiled_as(monkeypatch: pytest.MonkeyPatch, non_finite: bool, prelude: str | None = None) -> None: + """ + Has the Postgres connection run each Cypher statement's text as SQL, compiled with `non_finite` as given, as + drivine4j's compiler would hand it over. + """ + + def compiled( + spec: QuerySpecification[Any], *_: Any, **__: Any + ) -> tuple[QuerySpecification[Any], CompiledCypher | None]: + if spec.statement is None or spec.statement.language != "CYPHER": + return spec, None + sql = copy.copy(spec) + sql.statement = Statement(spec.statement.text, "SQL") + sql.parameters = {} + return sql, CompiledCypher(spec.statement.text, {}, prelude=prelude, non_finite=non_finite) + + monkeypatch.setattr(sys.modules[PostgresConnection.__module__], "cypher_compiled", compiled) + + +@pytest.fixture +async def connection(registry: DatabaseRegistry) -> AsyncIterator[Connection]: + provider = registry.connection_provider("postgres") + assert provider is not None + connection = await provider.connect() + assert isinstance(connection, PostgresConnection) + try: + yield connection + finally: + await connection.release() + + +def same(read: list[Any], expected: list[Any]) -> bool: + """Whether the values are the same, NaN the same as NaN.""" + return len(read) == len(expected) and all( + (isinstance(a, float) and isinstance(b, float) and math.isnan(a) and math.isnan(b)) or a == b + for a, b in zip(read, expected, strict=True) + ) + + +@pytest.mark.parametrize("prelude", [None, "SELECT 1"]) +async def test_reads_the_stand_ins_as_nan_and_the_infinities_where_the_statement_says_so( + connection: Connection, monkeypatch: pytest.MonkeyPatch, prelude: str | None +): + compiled_as(monkeypatch, True, prelude) + [row] = await connection.query(QuerySpecification[Any](Statement(STATEMENT, "CYPHER"))) + assert same(row, [math.nan, math.inf, -math.inf, 10**308, 2.5, "2" + "0" * 309 + ".5"]) + # The cursor's alone: the next statement, not given NaN, reads as today. + compiled_as(monkeypatch, False, prelude) + [row] = await connection.query(QuerySpecification[Any](Statement(STATEMENT, "CYPHER"))) + assert same(row, AS_TODAY) + + +@pytest.mark.parametrize("prelude", [None, "SELECT 1"]) +async def test_reads_the_stand_ins_as_today_where_the_statement_doesnt_say_so( + connection: Connection, monkeypatch: pytest.MonkeyPatch, prelude: str | None +): + compiled_as(monkeypatch, False, prelude) + [row] = await connection.query(QuerySpecification[Any](Statement(STATEMENT, "CYPHER"))) + assert same(row, AS_TODAY) + + +@pytest.mark.parametrize(("non_finite", "expected"), [(True, [math.nan, math.inf, -math.inf]), (False, AS_TODAY[:3])]) +@pytest.mark.parametrize("prelude", [None, "SELECT 1"]) +async def test_reads_a_cursors_rows_as_the_statement_says( + connection: Connection, + monkeypatch: pytest.MonkeyPatch, + non_finite: bool, + expected: list[float], + prelude: str | None, +): + compiled_as(monkeypatch, non_finite, prelude) + statement = ( + "SELECT to_jsonb(it) AS result FROM unnest(ARRAY[2e309 + 0.5, 1e309 + 0.5, -1e309 - 0.5]) WITH ORDINALITY " + "AS t(it, i) ORDER BY i" + ) + cursor = await connection.open_cursor(CursorSpecification[Any](Statement(statement, "CYPHER")).batch_size(2)) + try: + read = [it async for it in cursor] + finally: + await cursor.close() + assert same(read, expected) + + +async def test_reads_plain_sql_as_today(connection: Connection, monkeypatch: pytest.MonkeyPatch): + compiled_as(monkeypatch, True) + [row] = await connection.query(QuerySpecification[Any](STATEMENT)) + assert same(row["result"], AS_TODAY) diff --git a/tests/sync_integration/postgres/test_postgres_non_finite.py b/tests/sync_integration/postgres/test_postgres_non_finite.py new file mode 100644 index 00000000..5a1176f5 --- /dev/null +++ b/tests/sync_integration/postgres/test_postgres_non_finite.py @@ -0,0 +1,124 @@ +# Written by scripts/unasync.py from tests/integration/postgres/test_postgres_non_finite.py: +# change that, and run the script again. +""" +The numbers that stand for NaN and the infinities in the rows of a statement given one, as drivine4j's compiler writes +them on Postgres, read back as the doubles they are where the statement says so, and as today where it doesn't +(liberation-data/drivine#673). +""" + +import copy +import math +import sys +from collections.abc import Iterator +from typing import Any + +import pytest + +from drivine.cypher.sql.compiled_cypher import CompiledCypher +from drivine.query.statement import Statement +from drivine.sync import CursorSpecification, DatabaseRegistry, QuerySpecification +from drivine.sync.connection.connection import Connection +from drivine.sync.connection.postgres.postgres_connection import PostgresConnection +from tests.sync_integration.postgres.conftest import ON_POSTGRES + +pytestmark = ON_POSTGRES + +NUMBERS = ( + "jsonb_build_array(to_jsonb(2e309 + 0.5), to_jsonb(1e309 + 0.5), to_jsonb(-1e309 - 0.5), to_jsonb(1e308), " + "to_jsonb(2.5), to_jsonb('2' || repeat('0', 309) || '.5'))" +) +"""The numbers drivine4j's compiler writes for NaN and the infinities, two that stand for neither, and a string.""" + +STATEMENT = f"SELECT {NUMBERS} AS result" + +AS_TODAY = [math.inf, math.inf, -math.inf, 10**308, 2.5, "2" + "0" * 309 + ".5"] +"""As json.loads reads them: Postgres writes `numeric` 1e308 as its digits, a whole number.""" + + +def compiled_as(monkeypatch: pytest.MonkeyPatch, non_finite: bool, prelude: str | None = None) -> None: + """ + Has the Postgres connection run each Cypher statement's text as SQL, compiled with `non_finite` as given, as + drivine4j's compiler would hand it over. + """ + + def compiled( + spec: QuerySpecification[Any], *_: Any, **__: Any + ) -> tuple[QuerySpecification[Any], CompiledCypher | None]: + if spec.statement is None or spec.statement.language != "CYPHER": + return spec, None + sql = copy.copy(spec) + sql.statement = Statement(spec.statement.text, "SQL") + sql.parameters = {} + return sql, CompiledCypher(spec.statement.text, {}, prelude=prelude, non_finite=non_finite) + + monkeypatch.setattr(sys.modules[PostgresConnection.__module__], "cypher_compiled", compiled) + + +@pytest.fixture +def connection(registry: DatabaseRegistry) -> Iterator[Connection]: + provider = registry.connection_provider("postgres") + assert provider is not None + connection = provider.connect() + assert isinstance(connection, PostgresConnection) + try: + yield connection + finally: + connection.release() + + +def same(read: list[Any], expected: list[Any]) -> bool: + """Whether the values are the same, NaN the same as NaN.""" + return len(read) == len(expected) and all( + (isinstance(a, float) and isinstance(b, float) and math.isnan(a) and math.isnan(b)) or a == b + for a, b in zip(read, expected, strict=True) + ) + + +@pytest.mark.parametrize("prelude", [None, "SELECT 1"]) +def test_reads_the_stand_ins_as_nan_and_the_infinities_where_the_statement_says_so( + connection: Connection, monkeypatch: pytest.MonkeyPatch, prelude: str | None +): + compiled_as(monkeypatch, True, prelude) + [row] = connection.query(QuerySpecification[Any](Statement(STATEMENT, "CYPHER"))) + assert same(row, [math.nan, math.inf, -math.inf, 10**308, 2.5, "2" + "0" * 309 + ".5"]) + # The cursor's alone: the next statement, not given NaN, reads as today. + compiled_as(monkeypatch, False, prelude) + [row] = connection.query(QuerySpecification[Any](Statement(STATEMENT, "CYPHER"))) + assert same(row, AS_TODAY) + + +@pytest.mark.parametrize("prelude", [None, "SELECT 1"]) +def test_reads_the_stand_ins_as_today_where_the_statement_doesnt_say_so( + connection: Connection, monkeypatch: pytest.MonkeyPatch, prelude: str | None +): + compiled_as(monkeypatch, False, prelude) + [row] = connection.query(QuerySpecification[Any](Statement(STATEMENT, "CYPHER"))) + assert same(row, AS_TODAY) + + +@pytest.mark.parametrize(("non_finite", "expected"), [(True, [math.nan, math.inf, -math.inf]), (False, AS_TODAY[:3])]) +@pytest.mark.parametrize("prelude", [None, "SELECT 1"]) +def test_reads_a_cursors_rows_as_the_statement_says( + connection: Connection, + monkeypatch: pytest.MonkeyPatch, + non_finite: bool, + expected: list[float], + prelude: str | None, +): + compiled_as(monkeypatch, non_finite, prelude) + statement = ( + "SELECT to_jsonb(it) AS result FROM unnest(ARRAY[2e309 + 0.5, 1e309 + 0.5, -1e309 - 0.5]) WITH ORDINALITY " + "AS t(it, i) ORDER BY i" + ) + cursor = connection.open_cursor(CursorSpecification[Any](Statement(statement, "CYPHER")).batch_size(2)) + try: + read = [it for it in cursor] + finally: + cursor.close() + assert same(read, expected) + + +def test_reads_plain_sql_as_today(connection: Connection, monkeypatch: pytest.MonkeyPatch): + compiled_as(monkeypatch, True) + [row] = connection.query(QuerySpecification[Any](STATEMENT)) + assert same(row["result"], AS_TODAY) diff --git a/tests/unit/test_postgres_non_finite.py b/tests/unit/test_postgres_non_finite.py new file mode 100644 index 00000000..eda58407 --- /dev/null +++ b/tests/unit/test_postgres_non_finite.py @@ -0,0 +1,306 @@ +""" +How the rows of a statement given NaN or an infinity are read on Postgres, as drivine4j's `PostgresCypherCompiler.read` +reads them, and which statements are read so: those that say they were given one, alone (liberation-data/drivine#673). +""" + +import json +import math +from typing import Any + +import psycopg +import pytest +from psycopg.adapt import AdaptersMap +from psycopg.pq import Format + +from drivine.connection.postgres import postgres_connection +from drivine.connection.postgres.postgres_non_finite import non_finite_of, read_non_finite, read_non_finite_json +from drivine.cursor.cursor_specification import CursorSpecification +from drivine.cypher.cypher_specification import cypher_compiled +from drivine.cypher.sql.compiled_cypher import CompiledCypher +from drivine.query.query_specification import QuerySpecification +from drivine.query.statement import Statement +from drivine.sync.connection.postgres import postgres_connection as sync_postgres_connection + +NAN = "2" + "0" * 309 + ".5" +"""The number drivine4j's compiler writes on Postgres for NaN, as `jsonb` writes it.""" + +INFINITY = "1" + "0" * 309 + ".5" +NEGATIVE_INFINITY = f"-{INFINITY}" + +JSON_OID = 114 +JSONB_OID = 3802 + + +def is_nan(value: Any) -> bool: + return isinstance(value, float) and math.isnan(value) + + +@pytest.mark.parametrize( + ("text", "number"), + [ + (NAN, math.nan), + (f"{NAN}0", math.nan), + (f"{NAN}e0", math.nan), + (f"{NAN}E+0", math.nan), + (f"0.2{'0' * 309}5e310", math.nan), + (f"2{'0' * 309}5e-1", math.nan), + (INFINITY, math.inf), + (NEGATIVE_INFINITY, -math.inf), + ("1e309", math.inf), + ("1" + "0" * 309, math.inf), + ("2e309", math.inf), + (f"2{'0' * 309}.4", math.inf), + (f"2{'0' * 309}.6", math.inf), + ("-2e309", -math.inf), + (f"-{NAN}", -math.inf), + ("-1" + "0" * 309, -math.inf), + ("1e400", math.inf), + ], +) +def test_reads_a_stand_in_as_the_double_it_stands_for(text: str, number: float): + found = non_finite_of(text) + assert is_nan(found) if math.isnan(number) else found == number + [read] = read_non_finite_json(f"[{text}]") + assert is_nan(read) if math.isnan(number) else read == number + + +@pytest.mark.parametrize( + "text", + [ + "1e308", + "-1e308", + "1.7976931348623157e308", + "1.8e308", + "9" * 309, + "-" + "9" * 309, + f"{'9' * 309}.9", + "0", + "-0", + "-0.0", + "2.5", + "12345678901234567890123", + "0.0000000000000000000000000000000000000000000000000001", + ], +) +def test_reads_any_other_number_as_json_loads_does(text: str): + read = read_non_finite_json(text) + today = json.loads(text) + assert read == today + assert type(read) is type(today) + if isinstance(today, float): + assert math.copysign(1, read) == math.copysign(1, today) + + +@pytest.mark.parametrize( + ("text", "number"), [("1e99999999999999999999", math.inf), ("-1e99999999999999999999", -math.inf)] +) +def test_reads_a_number_past_decimals_exponent_as_json_loads_does(text: str, number: float): + assert read_non_finite_json(f"[{text}]") == json.loads(f"[{text}]") == [number] + + +def test_reads_the_stand_ins_nested_in_lists_and_maps_and_leaves_the_rest(): + text = ( + f'{{"a": [1, {NAN}, {{"b": {INFINITY}}}], "c": {NEGATIVE_INFINITY}, "d": "{NAN}", "e": 2e309, "f": 1e308, ' + f'"g": true, "h": null}}' + ) + read = read_non_finite_json(text.encode()) + assert is_nan(read["a"][1]) + read["a"][1] = "NaN" + assert read == { + "a": [1, "NaN", {"b": math.inf}], + "c": -math.inf, + "d": NAN, + "e": math.inf, + "f": 1e308, + "g": True, + "h": None, + } + + +def test_reads_the_stand_in_for_nan_as_json_loads_does_an_infinity_where_not_asked_to(): + assert json.loads(NAN) == math.inf + assert json.loads(f'{{"a": [{NAN}]}}') == {"a": [math.inf]} + + +class Cursor: + """A cursor of psycopg's that reads its JSON as `adapters` say, as psycopg's own reads it.""" + + def __init__(self) -> None: + self.adapters = AdaptersMap(psycopg.adapters) + + +def loaded(adapters: AdaptersMap, oid: int, format: Format, text: str) -> Any: + loader = adapters.get_loader(oid, format) + assert loader is not None + data = text.encode() + return loader(oid).load(b"\x01" + data if (oid, format) == (JSONB_OID, Format.BINARY) else data) + + +@pytest.mark.parametrize("oid", [JSON_OID, JSONB_OID]) +@pytest.mark.parametrize("format", [Format.TEXT, Format.BINARY]) +def test_has_a_cursor_read_json_and_jsonb_for_nan_and_nothing_else(oid: int, format: Format): + cursor = Cursor() + read_non_finite(cursor) + assert is_nan(loaded(cursor.adapters, oid, format, f"[{NAN}]")[0]) + # Nothing is set for the process, nor for another cursor. + assert loaded(psycopg.adapters, oid, format, f"[{NAN}]") == [math.inf] + assert loaded(Cursor().adapters, oid, format, f"[{NAN}]") == [math.inf] + + +CYPHER = [ + "MATCH (n) RETURN n.x", + "MATCH (n:A {x: 1.5}) RETURN n.x", + "CREATE (n:A {x: $x}) RETURN n.x", + "MATCH (n:A) SET n.x = 1 RETURN n", + "RETURN 1 AS x UNION RETURN 2 AS x", + "UNWIND [1, 2] AS x RETURN x ORDER BY x", +] + + +@pytest.mark.parametrize("cypher", CYPHER) +@pytest.mark.parametrize("store", ["POSTGRES", "ORACLE", "SQLITE"]) +def test_drivines_own_compiler_says_no_statement_was_given_nan(cypher: str, store: Any): + _, compiled = cypher_compiled(QuerySpecification[Any](Statement(cypher, "CYPHER")).bind({"x": 2}), store, "g") + assert compiled is not None + assert compiled.non_finite is False + + +def test_a_compiled_statement_says_it_was_given_no_nan_unless_told(): + assert CompiledCypher("SELECT 1", {}).non_finite is False + assert CompiledCypher("SELECT 1", {}, non_finite=True).non_finite is True + + +ROW = {"result": 1} + + +class AsyncCursor: + """A psycopg cursor that gives the row ROW, and records whether its JSON was read for NaN.""" + + def __init__(self, record: list[str]) -> None: + self.record = record + self.description = [("result",)] + self._adapters: AdaptersMap | None = None + + @property + def adapters(self) -> AdaptersMap: + self.record.append("adapters") + if self._adapters is None: + self._adapters = AdaptersMap(psycopg.adapters) + return self._adapters + + async def __aenter__(self) -> "AsyncCursor": + return self + + async def __aexit__(self, *_: Any) -> None: + pass + + async def execute(self, statement: str, parameters: Any = None) -> None: + self.record.append(statement) + + async def fetchall(self) -> list[dict[str, Any]]: + return [ROW] + + async def fetchmany(self, _: int) -> list[dict[str, Any]]: + return [ROW] + + async def close(self) -> None: + pass + + +class SyncCursor(AsyncCursor): + def __enter__(self) -> "SyncCursor": + return self + + def __exit__(self, *_: Any) -> None: + pass + + def execute(self, statement: str, parameters: Any = None) -> None: # type: ignore[override] + self.record.append(statement) + + def fetchall(self) -> list[dict[str, Any]]: # type: ignore[override] + return [ROW] + + def fetchmany(self, _: int) -> list[dict[str, Any]]: # type: ignore[override] + return [ROW] + + def close(self) -> None: # type: ignore[override] + pass + + +class Client: + """A psycopg connection whose cursors are `cursor_type`'s.""" + + def __init__(self, cursor_type: type[AsyncCursor]) -> None: + self.record: list[str] = [] + self.cursors: list[AsyncCursor] = [] + self.cursor_type = cursor_type + + def cursor(self, **_: Any) -> AsyncCursor: + cursor = self.cursor_type(self.record) + self.cursors.append(cursor) + return cursor + + async def execute(self, statement: str) -> None: + self.record.append(statement) + + +class SyncClient(Client): + def execute(self, statement: str) -> None: # type: ignore[override] + self.record.append(statement) + + +def compiled_as(monkeypatch: pytest.MonkeyPatch, module: Any, non_finite: bool, prelude: str | None) -> None: + """Has the connection of `module` compile each Cypher statement to `SELECT 2`, with `non_finite` as given.""" + + def compiled(spec: QuerySpecification[Any], *_: Any) -> tuple[QuerySpecification[Any], CompiledCypher | None]: + if spec.statement is None or spec.statement.language != "CYPHER": + return spec, None + sql = QuerySpecification[Any](Statement("SELECT 2", "SQL")) + return sql, CompiledCypher("SELECT 2", {}, prelude=prelude, non_finite=non_finite) + + monkeypatch.setattr(module, "cypher_compiled", compiled) + + +async def run(sync: bool, how: str, spec: Any) -> tuple[list[Any], Client]: + """The rows `spec` gives through a Postgres connection, async or `sync`, by `query` or a cursor, and its client.""" + module = sync_postgres_connection if sync else postgres_connection + client = (SyncClient if sync else Client)(SyncCursor if sync else AsyncCursor) + connection: Any = module.PostgresConnection(client, lambda _: None, "pg.1") + if how == "query": + rows = connection.query(QuerySpecification[Any](spec)) + return (rows if sync else await rows), client + cursor = connection.open_cursor(CursorSpecification[Any](spec).batch_size(5)) + cursor = cursor if sync else await cursor + return (cursor.read(5) if sync else await cursor.read(5))[0], client + + +@pytest.mark.parametrize("sync", [False, True]) +@pytest.mark.parametrize("how", ["query", "cursor"]) +@pytest.mark.parametrize("prelude", [None, "PRELUDE"]) +async def test_a_connection_reads_for_nan_only_a_statement_that_says_it_was_given_one( + monkeypatch: pytest.MonkeyPatch, sync: bool, how: str, prelude: str | None +): + module = sync_postgres_connection if sync else postgres_connection + compiled_as(monkeypatch, module, False, prelude) + rows, client = await run(sync, how, Statement("RETURN 1", "CYPHER")) + assert rows == [1] + # Not given NaN, its cursor is left as psycopg made it: never even asked for its adapters. + assert "adapters" not in client.record + assert client.record == [it for it in [prelude, "SELECT 2"] if it is not None] + + compiled_as(monkeypatch, module, True, prelude) + rows, client = await run(sync, how, Statement("RETURN 1", "CYPHER")) + assert rows == [1] + [cursor] = client.cursors + assert cursor._adapters is not None + assert is_nan(loaded(cursor._adapters, JSONB_OID, Format.TEXT, f"[{NAN}]")[0]) + # Its JSON is read for NaN before the statement runs. + assert client.record.index("adapters") < client.record.index("SELECT 2") + + +@pytest.mark.parametrize("sync", [False, True]) +@pytest.mark.parametrize("how", ["query", "cursor"]) +async def test_a_connection_reads_plain_sql_and_drivines_own_compiled_statements_as_psycopg_does(sync: bool, how: str): + for spec in [Statement("SELECT 1", "SQL"), Statement("RETURN 1 AS x", "CYPHER")]: + _, client = await run(sync, how, spec) + assert "adapters" not in client.record