Skip to content

Commit 41c5082

Browse files
judahrandjudahrand
committed
Use psycopg rather than psycopg2 for Risingwave
1 parent 4c83730 commit 41c5082

22 files changed

+107
-145
lines changed

‎.github/renovate.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@
8989
},
9090
{
9191
"addLabels": ["risingwave"],
92-
"matchPackageNames": ["/psycopg2/", "/risingwave/"]
92+
"matchPackageNames": ["/psycopg/", "/risingwave/"]
9393
},
9494
{
9595
"addLabels": ["snowflake"],

‎ibis/backends/postgres/__init__.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -713,11 +713,11 @@ def raw_sql(self, query: str | sg.Expression, **kwargs: Any) -> Any:
713713

714714
try:
715715
# try to load hstore
716-
with contextlib.suppress(TypeError):
717-
type_info = psycopg.types.TypeInfo.fetch(con, "hstore")
718-
with contextlib.suppress(psycopg.ProgrammingError, TypeError):
719-
psycopg.types.hstore.register_hstore(type_info, cursor)
720-
except Exception:
716+
psycopg.types.hstore.register_hstore(
717+
psycopg.types.TypeInfo.fetch(con, "hstore"),
718+
cursor,
719+
)
720+
except (psycopg.InternalError, psycopg.ProgrammingError, TypeError):
721721
cursor.close()
722722
raise
723723

‎ibis/backends/risingwave/__init__.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,9 @@
55
from operator import itemgetter
66
from typing import TYPE_CHECKING
77

8-
import psycopg2
8+
import psycopg
99
import sqlglot as sg
1010
import sqlglot.expressions as sge
11-
from psycopg2 import extras
1211

1312
import ibis
1413
import ibis.backends.sql.compilers as sc
@@ -110,12 +109,12 @@ def do_connect(
110109
month int32
111110
"""
112111

113-
self.con = psycopg2.connect(
112+
self.con = psycopg.connect(
114113
host=host,
115114
port=port,
116115
user=user,
117116
password=password,
118-
database=database,
117+
dbname=database,
119118
options=(f"-csearch_path={schema}" * (schema is not None)) or None,
120119
)
121120

@@ -289,7 +288,7 @@ def _register_in_memory_table(self, op: ops.InMemoryTable) -> None:
289288
)
290289
with self.begin() as cur:
291290
cur.execute(create_stmt_sql)
292-
extras.execute_batch(cur, sql, data, 128)
291+
cur.executemany(sql, data)
293292

294293
def list_databases(
295294
self, *, like: str | None = None, catalog: str | None = None

‎ibis/backends/risingwave/tests/conftest.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ class TestConf(ServiceBackendTest):
3333
supports_structs = False
3434
rounding_method = "half_to_even"
3535
service_name = "risingwave"
36-
deps = ("psycopg2",)
36+
deps = ("psycopg",)
3737

3838
@property
3939
def test_files(self) -> Iterable[Path]:

‎ibis/backends/risingwave/tests/test_client.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import ibis.expr.types as ir
1313
from ibis.util import gen_name
1414

15-
pytest.importorskip("psycopg2")
15+
pytest.importorskip("psycopg")
1616

1717
RISINGWAVE_TEST_DB = os.environ.get("IBIS_TEST_RISINGWAVE_DATABASE", "dev")
1818
IBIS_RISINGWAVE_HOST = os.environ.get("IBIS_TEST_RISINGWAVE_HOST", "localhost")

‎ibis/backends/risingwave/tests/test_functions.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
import ibis.expr.datatypes as dt
1515
from ibis import literal as L
1616

17-
pytest.importorskip("psycopg2")
17+
pytest.importorskip("psycopg")
1818

1919

2020
@pytest.mark.parametrize(("value", "expected"), [(0, None), (5.5, 5.5)])

‎ibis/backends/tests/errors.py

Lines changed: 0 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -131,25 +131,6 @@
131131
PsycoPgOperationalError
132132
) = PsycoPgUndefinedObject = PsycoPgArraySubscriptError = None
133133

134-
try:
135-
from psycopg2.errors import ArraySubscriptError as PsycoPg2ArraySubscriptError
136-
from psycopg2.errors import DivisionByZero as PsycoPg2DivisionByZero
137-
from psycopg2.errors import IndeterminateDatatype as PsycoPg2IndeterminateDatatype
138-
from psycopg2.errors import InternalError_ as PsycoPg2InternalError
139-
from psycopg2.errors import (
140-
InvalidTextRepresentation as PsycoPg2InvalidTextRepresentation,
141-
)
142-
from psycopg2.errors import OperationalError as PsycoPg2OperationalError
143-
from psycopg2.errors import ProgrammingError as PsycoPg2ProgrammingError
144-
from psycopg2.errors import SyntaxError as PsycoPg2SyntaxError
145-
from psycopg2.errors import UndefinedObject as PsycoPg2UndefinedObject
146-
except ImportError:
147-
PsycoPg2SyntaxError = PsycoPg2IndeterminateDatatype = (
148-
PsycoPg2InvalidTextRepresentation
149-
) = PsycoPg2DivisionByZero = PsycoPg2InternalError = PsycoPg2ProgrammingError = (
150-
PsycoPg2OperationalError
151-
) = PsycoPg2UndefinedObject = PsycoPg2ArraySubscriptError = None
152-
153134
try:
154135
from MySQLdb import NotSupportedError as MySQLNotSupportedError
155136
from MySQLdb import OperationalError as MySQLOperationalError

‎ibis/backends/tests/test_aggregation.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
MySQLNotSupportedError,
2222
OracleDatabaseError,
2323
PolarsInvalidOperationError,
24-
PsycoPg2InternalError,
24+
PsycoPgInternalError,
2525
Py4JError,
2626
Py4JJavaError,
2727
PyAthenaOperationalError,
@@ -963,7 +963,7 @@ def test_approx_quantile(con, filtered, multi):
963963
),
964964
pytest.mark.notimpl(
965965
["risingwave"],
966-
raises=PsycoPg2InternalError,
966+
raises=PsycoPgInternalError,
967967
reason="function covar_pop(integer, integer) does not exist",
968968
),
969969
],
@@ -983,7 +983,7 @@ def test_approx_quantile(con, filtered, multi):
983983
),
984984
pytest.mark.notimpl(
985985
["risingwave"],
986-
raises=PsycoPg2InternalError,
986+
raises=PsycoPgInternalError,
987987
reason="function covar_pop(integer, integer) does not exist",
988988
),
989989
],
@@ -1005,7 +1005,7 @@ def test_approx_quantile(con, filtered, multi):
10051005
),
10061006
pytest.mark.notimpl(
10071007
["risingwave"],
1008-
raises=PsycoPg2InternalError,
1008+
raises=PsycoPgInternalError,
10091009
reason="function covar_pop(integer, integer) does not exist",
10101010
),
10111011
],
@@ -1062,7 +1062,7 @@ def test_approx_quantile(con, filtered, multi):
10621062
),
10631063
pytest.mark.notimpl(
10641064
["risingwave"],
1065-
raises=PsycoPg2InternalError,
1065+
raises=PsycoPgInternalError,
10661066
reason="function covar_pop(integer, integer) does not exist",
10671067
),
10681068
],
@@ -1088,7 +1088,7 @@ def test_approx_quantile(con, filtered, multi):
10881088
),
10891089
pytest.mark.notimpl(
10901090
["risingwave"],
1091-
raises=PsycoPg2InternalError,
1091+
raises=PsycoPgInternalError,
10921092
reason="function covar_pop(integer, integer) does not exist",
10931093
),
10941094
],

‎ibis/backends/tests/test_array.py

Lines changed: 21 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,10 @@
2222
GoogleBadRequest,
2323
MySQLOperationalError,
2424
PolarsComputeError,
25-
PsycoPg2IndeterminateDatatype,
26-
PsycoPg2InternalError,
27-
PsycoPg2ProgrammingError,
2825
PsycoPgIndeterminateDatatype,
2926
PsycoPgInternalError,
3027
PsycoPgInvalidTextRepresentation,
28+
PsycoPgProgrammingError,
3129
PsycoPgSyntaxError,
3230
Py4JJavaError,
3331
PyAthenaDatabaseError,
@@ -506,7 +504,7 @@ def test_array_slice(backend, start, stop):
506504
)
507505
@pytest.mark.notimpl(
508506
["risingwave"],
509-
raises=PsycoPg2InternalError,
507+
raises=PsycoPgInternalError,
510508
reason="TODO(Kexiang): seems a bug",
511509
)
512510
@pytest.mark.notimpl(["athena"], raises=PyAthenaDatabaseError)
@@ -565,7 +563,7 @@ def test_array_map(con, input, output, func):
565563
)
566564
@pytest.mark.notimpl(
567565
["risingwave"],
568-
raises=PsycoPg2InternalError,
566+
raises=PsycoPgInternalError,
569567
reason="TODO(Kexiang): seems a bug",
570568
)
571569
@pytest.mark.notimpl(["athena"], raises=PyAthenaDatabaseError)
@@ -646,7 +644,7 @@ def test_array_map_with_index(con, input, output, func):
646644
)
647645
@pytest.mark.notyet(
648646
"risingwave",
649-
raises=PsycoPg2InternalError,
647+
raises=PsycoPgInternalError,
650648
reason="no support for not null column constraint",
651649
)
652650
@pytest.mark.parametrize(
@@ -693,7 +691,7 @@ def test_array_filter(con, input, output, predicate):
693691
)
694692
@pytest.mark.notyet(
695693
"risingwave",
696-
raises=PsycoPg2InternalError,
694+
raises=PsycoPgInternalError,
697695
reason="no support for not null column constraint",
698696
)
699697
@pytest.mark.parametrize(
@@ -740,7 +738,7 @@ def test_array_filter_with_index(con, input, output, predicate):
740738
)
741739
@pytest.mark.notyet(
742740
"risingwave",
743-
raises=PsycoPg2InternalError,
741+
raises=PsycoPgInternalError,
744742
reason="no support for not null column constraint",
745743
)
746744
@pytest.mark.parametrize(
@@ -1097,7 +1095,7 @@ def test_array_intersect(con, data):
10971095

10981096
@builtin_array
10991097
@pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError)
1100-
@pytest.mark.notimpl(["risingwave"], raises=PsycoPg2InternalError)
1098+
@pytest.mark.notimpl(["risingwave"], raises=PsycoPgInternalError)
11011099
@pytest.mark.notimpl(
11021100
["trino"], reason="inserting maps into structs doesn't work", raises=TrinoUserError
11031101
)
@@ -1117,7 +1115,7 @@ def test_unnest_struct(con):
11171115

11181116
@builtin_array
11191117
@pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError)
1120-
@pytest.mark.notimpl(["risingwave"], raises=PsycoPg2InternalError)
1118+
@pytest.mark.notimpl(["risingwave"], raises=PsycoPgInternalError)
11211119
@pytest.mark.notimpl(
11221120
["trino"], reason="inserting maps into structs doesn't work", raises=TrinoUserError
11231121
)
@@ -1208,7 +1206,7 @@ def test_zip_null(con, fn):
12081206

12091207
@builtin_array
12101208
@pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError)
1211-
@pytest.mark.notimpl(["risingwave"], raises=PsycoPg2ProgrammingError)
1209+
@pytest.mark.notimpl(["risingwave"], raises=PsycoPgProgrammingError)
12121210
@pytest.mark.notimpl(["datafusion"], raises=Exception, reason="not yet supported")
12131211
@pytest.mark.notimpl(
12141212
["polars"],
@@ -1291,8 +1289,8 @@ def flatten_data():
12911289
reason="Risingwave doesn't truly support arrays of arrays",
12921290
raises=(
12931291
com.OperationNotDefinedError,
1294-
PsycoPg2IndeterminateDatatype,
1295-
PsycoPg2InternalError,
1292+
PsycoPgIndeterminateDatatype,
1293+
PsycoPgInternalError,
12961294
),
12971295
)
12981296
@pytest.mark.parametrize(
@@ -1399,7 +1397,7 @@ def test_range_start_stop_step(con, start, stop, step):
13991397
@pytest.mark.notimpl(["flink"], raises=com.OperationNotDefinedError)
14001398
@pytest.mark.never(
14011399
["risingwave"],
1402-
raises=PsycoPg2InternalError,
1400+
raises=PsycoPgInternalError,
14031401
reason="Invalid parameter step: step size cannot equal zero",
14041402
)
14051403
def test_range_start_stop_step_zero(con, start, stop):
@@ -1432,7 +1430,7 @@ def test_unnest_empty_array(con):
14321430
@pytest.mark.notimpl(["sqlite"], raises=com.UnsupportedBackendType)
14331431
@pytest.mark.notyet(
14341432
"risingwave",
1435-
raises=PsycoPg2InternalError,
1433+
raises=PsycoPgInternalError,
14361434
reason="no support for not null column constraint",
14371435
)
14381436
@pytest.mark.notimpl(["athena"], raises=PyAthenaDatabaseError)
@@ -1515,7 +1513,7 @@ def swap(token):
15151513
id="pos",
15161514
marks=pytest.mark.notimpl(
15171515
["risingwave"],
1518-
raises=PsycoPg2InternalError,
1516+
raises=PsycoPgInternalError,
15191517
reason="function make_interval() does not exist",
15201518
),
15211519
),
@@ -1533,7 +1531,7 @@ def swap(token):
15331531
),
15341532
pytest.mark.notimpl(
15351533
["risingwave"],
1536-
raises=PsycoPg2InternalError,
1534+
raises=PsycoPgInternalError,
15371535
reason="function neg(interval) does not exist",
15381536
),
15391537
],
@@ -1553,7 +1551,7 @@ def swap(token):
15531551
),
15541552
pytest.mark.notimpl(
15551553
["risingwave"],
1556-
raises=PsycoPg2InternalError,
1554+
raises=PsycoPgInternalError,
15571555
reason="function neg(interval) does not exist",
15581556
),
15591557
],
@@ -1585,7 +1583,7 @@ def test_timestamp_range(con, start, stop, step, freq, tzinfo):
15851583
pytest.mark.notyet(["polars"], raises=PolarsComputeError),
15861584
pytest.mark.notyet(
15871585
["risingwave"],
1588-
raises=PsycoPg2InternalError,
1586+
raises=PsycoPgInternalError,
15891587
reason="function make_interval() does not exist",
15901588
),
15911589
],
@@ -1604,7 +1602,7 @@ def test_timestamp_range(con, start, stop, step, freq, tzinfo):
16041602
),
16051603
pytest.mark.notyet(
16061604
["risingwave"],
1607-
raises=PsycoPg2InternalError,
1605+
raises=PsycoPgInternalError,
16081606
reason="function neg(interval) does not exist",
16091607
),
16101608
],
@@ -1760,7 +1758,7 @@ def test_table_unnest_with_keep_empty(con):
17601758
["datafusion", "polars", "flink"], raises=com.OperationNotDefinedError
17611759
)
17621760
@pytest.mark.notyet(
1763-
["risingwave"], raises=PsycoPg2InternalError, reason="not supported in risingwave"
1761+
["risingwave"], raises=PsycoPgInternalError, reason="not supported in risingwave"
17641762
)
17651763
@pytest.mark.notimpl(
17661764
["athena"],
@@ -1781,9 +1779,9 @@ def test_table_unnest_column_expr(backend):
17811779
@pytest.mark.notimpl(["trino"], raises=TrinoUserError)
17821780
@pytest.mark.notimpl(["athena"], raises=PyAthenaOperationalError)
17831781
@pytest.mark.notimpl(["postgres"], raises=PsycoPgSyntaxError)
1784-
@pytest.mark.notimpl(["risingwave"], raises=PsycoPg2ProgrammingError)
1782+
@pytest.mark.notimpl(["risingwave"], raises=PsycoPgProgrammingError)
17851783
@pytest.mark.notyet(
1786-
["risingwave"], raises=PsycoPg2InternalError, reason="not supported in risingwave"
1784+
["risingwave"], raises=PsycoPgInternalError, reason="not supported in risingwave"
17871785
)
17881786
def test_table_unnest_array_of_struct_of_array(con):
17891787
t = ibis.memtable(

0 commit comments

Comments
 (0)