From 55e510eef0293ed0886819d2ce76eb24241455de Mon Sep 17 00:00:00 2001 From: Phillip Cloud <417981+cpcloud@users.noreply.github.com> Date: Thu, 16 Jan 2025 05:00:16 -0500 Subject: [PATCH] chore: move pandas import inside use site --- .github/renovate.json | 2 +- ibis/backends/risingwave/__init__.py | 35 +++---------------- ibis/backends/risingwave/tests/conftest.py | 2 +- ibis/backends/risingwave/tests/test_client.py | 2 +- .../risingwave/tests/test_functions.py | 2 +- ibis/backends/tests/errors.py | 19 ++++++++++ 6 files changed, 27 insertions(+), 35 deletions(-) diff --git a/.github/renovate.json b/.github/renovate.json index 9f1d4036dfea..6fbdba515037 100644 --- a/.github/renovate.json +++ b/.github/renovate.json @@ -89,7 +89,7 @@ }, { "addLabels": ["risingwave"], - "matchPackageNames": ["/psycopg/", "/risingwave/"] + "matchPackageNames": ["/psycopg2/", "/risingwave/"] }, { "addLabels": ["snowflake"], diff --git a/ibis/backends/risingwave/__init__.py b/ibis/backends/risingwave/__init__.py index 7dd8ce7ec861..3d5f72dbbb8e 100644 --- a/ibis/backends/risingwave/__init__.py +++ b/ibis/backends/risingwave/__init__.py @@ -9,7 +9,6 @@ import sqlglot as sg import sqlglot.expressions as sge -from pandas.api.types import is_float_dtype import ibis import ibis.backends.sql.compilers as sc @@ -29,6 +28,7 @@ import pandas as pd import polars as pl + import psycopg2.extensions import pyarrow as pa @@ -105,6 +105,7 @@ def _from_url(self, url: ParseResult, **kwargs): return self.connect(**kwargs) def _register_in_memory_table(self, op: ops.InMemoryTable) -> None: + from pandas.api.types import is_float_dtype from psycopg2.extras import execute_batch schema = op.schema @@ -528,13 +529,14 @@ def do_connect( year int32 month int32 """ + import psycopg2 self.con = psycopg2.connect( host=host, port=port, user=user, password=password, - dbname=database, + database=database, options=(f"-csearch_path={schema}" * (schema is not None)) or None, ) @@ -681,35 +683,6 @@ def create_table( name, schema=schema, source=self, namespace=ops.Namespace(database=database) ).to_expr() - def _register_in_memory_table(self, op: ops.InMemoryTable) -> None: - schema = op.schema - if null_columns := [col for col, dtype in schema.items() if dtype.is_null()]: - raise com.IbisTypeError( - f"{self.name} cannot yet reliably handle `null` typed columns; " - f"got null typed columns: {null_columns}" - ) - - name = op.name - quoted = self.compiler.quoted - - create_stmt = sg.exp.Create( - kind="TABLE", - this=sg.exp.Schema( - this=sg.to_identifier(name, quoted=quoted), - expressions=schema.to_sqlglot(self.dialect), - ), - ) - create_stmt_sql = create_stmt.sql(self.dialect) - - df = op.data.to_frame() - data = df.itertuples(index=False) - sql = self._build_insert_template( - name, schema=schema, columns=True, placeholder="%s" - ) - with self.begin() as cur: - cur.execute(create_stmt_sql) - cur.executemany(sql, data) - def list_databases( self, *, like: str | None = None, catalog: str | None = None ) -> list[str]: diff --git a/ibis/backends/risingwave/tests/conftest.py b/ibis/backends/risingwave/tests/conftest.py index 2832cddedab2..06940a6573b6 100644 --- a/ibis/backends/risingwave/tests/conftest.py +++ b/ibis/backends/risingwave/tests/conftest.py @@ -33,7 +33,7 @@ class TestConf(ServiceBackendTest): supports_structs = False rounding_method = "half_to_even" service_name = "risingwave" - deps = ("psycopg",) + deps = ("psycopg2",) @property def test_files(self) -> Iterable[Path]: diff --git a/ibis/backends/risingwave/tests/test_client.py b/ibis/backends/risingwave/tests/test_client.py index c0533a768e17..1d2ce761242d 100644 --- a/ibis/backends/risingwave/tests/test_client.py +++ b/ibis/backends/risingwave/tests/test_client.py @@ -12,7 +12,7 @@ import ibis.expr.types as ir from ibis.util import gen_name -pytest.importorskip("psycopg") +pytest.importorskip("psycopg2") RISINGWAVE_TEST_DB = os.environ.get("IBIS_TEST_RISINGWAVE_DATABASE", "dev") IBIS_RISINGWAVE_HOST = os.environ.get("IBIS_TEST_RISINGWAVE_HOST", "localhost") diff --git a/ibis/backends/risingwave/tests/test_functions.py b/ibis/backends/risingwave/tests/test_functions.py index dafbd315d98d..38be06b1281d 100644 --- a/ibis/backends/risingwave/tests/test_functions.py +++ b/ibis/backends/risingwave/tests/test_functions.py @@ -14,7 +14,7 @@ import ibis.expr.datatypes as dt from ibis import literal as L -pytest.importorskip("psycopg") +pytest.importorskip("psycopg2") @pytest.mark.parametrize(("value", "expected"), [(0, None), (5.5, 5.5)]) diff --git a/ibis/backends/tests/errors.py b/ibis/backends/tests/errors.py index 5824deec489f..7c601bf075a8 100644 --- a/ibis/backends/tests/errors.py +++ b/ibis/backends/tests/errors.py @@ -112,6 +112,25 @@ except ImportError: TrinoUserError = None +try: + from psycopg2.errors import ArraySubscriptError as PsycoPg2ArraySubscriptError + from psycopg2.errors import DivisionByZero as PsycoPg2DivisionByZero + from psycopg2.errors import IndeterminateDatatype as PsycoPg2IndeterminateDatatype + from psycopg2.errors import InternalError_ as PsycoPg2InternalError + from psycopg2.errors import ( + InvalidTextRepresentation as PsycoPg2InvalidTextRepresentation, + ) + from psycopg2.errors import OperationalError as PsycoPg2OperationalError + from psycopg2.errors import ProgrammingError as PsycoPg2ProgrammingError + from psycopg2.errors import SyntaxError as PsycoPg2SyntaxError + from psycopg2.errors import UndefinedObject as PsycoPg2UndefinedObject +except ImportError: + PsycoPg2SyntaxError = PsycoPg2IndeterminateDatatype = ( + PsycoPg2InvalidTextRepresentation + ) = PsycoPg2DivisionByZero = PsycoPg2InternalError = PsycoPg2ProgrammingError = ( + PsycoPg2OperationalError + ) = PsycoPg2UndefinedObject = PsycoPg2ArraySubscriptError = None + try: from psycopg.errors import ArraySubscriptError as PsycoPgArraySubscriptError from psycopg.errors import DivisionByZero as PsycoPgDivisionByZero