Read drivine4j's Postgres stand-ins for NaN and the infinities where a compiled statement says it holds them

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)
 

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)

drivine/cypher/sql/compiled_cypher.py

@@ -16,10 +16,15 @@ class CompiledCypher:
     `<graph>_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, `<graph>_era` and `<graph>_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
     parameters: dict[str, Any]
     prelude: str | None = None
     calls_stored_nulls_equal: bool = False
     calls_stored_era: bool = False
+    non_finite: bool = False

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)
 

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)