Adding upstream version 1.1.1.

Signed-off-by: Daniel Baumann <daniel@debian.org>

tests/conftest.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import sys
+from typing import Generator
+
+import psycopg
+import pytest
+from harlequin_postgres.adapter import (
+    HarlequinPostgresAdapter,
+    HarlequinPostgresConnection,
+)
+
+if sys.version_info < (3, 10):
+    pass
+else:
+    pass
+
+TEST_DB_CONN = "postgresql://postgres:for-testing@localhost:5432"
+
+
+@pytest.fixture
+def connection() -> Generator[HarlequinPostgresConnection, None, None]:
+    pgconn = psycopg.connect(conninfo=TEST_DB_CONN, dbname="postgres")
+    pgconn.autocommit = True
+    cur = pgconn.cursor()
+    cur.execute("drop database if exists test;")
+    cur.execute("create database test;")
+    cur.close()
+    pgconn.close()
+    conn = HarlequinPostgresAdapter(
+        conn_str=(f"{TEST_DB_CONN}",), dbname="test"
+    ).connect()
+    yield conn
+    conn.close()
+    pgconn = psycopg.connect(conninfo=TEST_DB_CONN, dbname="postgres")
+    pgconn.autocommit = True
+    cur = pgconn.cursor()
+    cur.execute("drop database if exists test;")
+    cur.close()
+    pgconn.close()

tests/test_adapter.py

@@ -0,0 +1,158 @@
+from __future__ import annotations
+
+import sys
+from datetime import date, datetime
+
+import pytest
+from harlequin.adapter import HarlequinAdapter, HarlequinConnection, HarlequinCursor
+from harlequin.catalog import Catalog, CatalogItem
+from harlequin.exception import HarlequinConnectionError, HarlequinQueryError
+from harlequin_postgres.adapter import (
+    HarlequinPostgresAdapter,
+    HarlequinPostgresConnection,
+)
+from textual_fastdatatable.backend import create_backend
+
+if sys.version_info < (3, 10):
+    from importlib_metadata import entry_points
+else:
+    from importlib.metadata import entry_points
+
+TEST_DB_CONN = "postgresql://postgres:for-testing@localhost:5432"
+
+
+def test_plugin_discovery() -> None:
+    PLUGIN_NAME = "postgres"
+    eps = entry_points(group="harlequin.adapter")
+    assert eps[PLUGIN_NAME]
+    adapter_cls = eps[PLUGIN_NAME].load()
+    assert issubclass(adapter_cls, HarlequinAdapter)
+    assert adapter_cls == HarlequinPostgresAdapter
+
+
+def test_connect() -> None:
+    conn = HarlequinPostgresAdapter(conn_str=(TEST_DB_CONN,)).connect()
+    assert isinstance(conn, HarlequinConnection)
+
+
+def test_init_extra_kwargs() -> None:
+    assert HarlequinPostgresAdapter(
+        conn_str=(TEST_DB_CONN,), foo=1, bar="baz"
+    ).connect()
+
+
+@pytest.mark.parametrize(
+    "conn_str",
+    [
+        ("foo",),
+        ("host=foo",),
+        ("postgresql://admin:pass@foo:5432/db",),
+    ],
+)
+def test_connect_raises_connection_error(conn_str: tuple[str]) -> None:
+    with pytest.raises(HarlequinConnectionError):
+        _ = HarlequinPostgresAdapter(conn_str=conn_str, connect_timeout=0.1).connect()
+
+
+@pytest.mark.parametrize(
+    "conn_str,options,expected",
+    [
+        (("",), {}, "localhost:5432/postgres"),
+        (("host=foo",), {}, "foo:5432/postgres"),
+        (("postgresql://foo",), {}, "foo:5432/postgres"),
+        (("postgresql://foo",), {"port": 5431}, "foo:5431/postgres"),
+        (("postgresql://foo/mydb",), {"port": 5431}, "foo:5431/mydb"),
+        (("postgresql://admin:pass@foo/mydb",), {"port": 5431}, "foo:5431/mydb"),
+        (("postgresql://admin:pass@foo:5431/mydb",), {}, "foo:5431/mydb"),
+    ],
+)
+def test_connection_id(
+    conn_str: tuple[str], options: dict[str, int | float | str | None], expected: str
+) -> None:
+    adapter = HarlequinPostgresAdapter(
+        conn_str=conn_str,
+        **options,  # type: ignore[arg-type]
+    )
+    assert adapter.connection_id == expected
+
+
+def test_get_catalog(connection: HarlequinPostgresConnection) -> None:
+    catalog = connection.get_catalog()
+    assert isinstance(catalog, Catalog)
+    assert catalog.items
+    assert isinstance(catalog.items[0], CatalogItem)
+
+
+def test_get_completions(connection: HarlequinPostgresConnection) -> None:
+    completions = connection.get_completions()
+    test_labels = ["atomic", "greatest", "point_right", "autovacuum"]
+    filtered = list(filter(lambda x: x.label in test_labels, completions))
+    assert len(filtered) == 4
+    value_filtered = list(filter(lambda x: x.value in test_labels, completions))
+    assert len(value_filtered) == 4
+
+
+def test_execute_ddl(connection: HarlequinPostgresConnection) -> None:
+    cur = connection.execute("create table foo (a int)")
+    assert cur is None
+
+
+def test_execute_select(connection: HarlequinPostgresConnection) -> None:
+    cur = connection.execute("select 1 as a")
+    assert isinstance(cur, HarlequinCursor)
+    assert cur.columns() == [("a", "#")]
+    data = cur.fetchall()
+    backend = create_backend(data)
+    assert backend.column_count == 1
+    assert backend.row_count == 1
+
+
+def test_execute_select_dupe_cols(connection: HarlequinPostgresConnection) -> None:
+    cur = connection.execute("select 1 as a, 2 as a, 3 as a")
+    assert isinstance(cur, HarlequinCursor)
+    assert len(cur.columns()) == 3
+    data = cur.fetchall()
+    backend = create_backend(data)
+    assert backend.column_count == 3
+    assert backend.row_count == 1
+
+
+def test_set_limit(connection: HarlequinPostgresConnection) -> None:
+    cur = connection.execute("select 1 as a union all select 2 union all select 3")
+    assert isinstance(cur, HarlequinCursor)
+    cur = cur.set_limit(2)
+    assert isinstance(cur, HarlequinCursor)
+    data = cur.fetchall()
+    backend = create_backend(data)
+    assert backend.column_count == 1
+    assert backend.row_count == 2
+
+
+def test_execute_raises_query_error(connection: HarlequinPostgresConnection) -> None:
+    with pytest.raises(HarlequinQueryError):
+        _ = connection.execute("sel;")
+
+
+def test_inf_timestamps(connection: HarlequinPostgresConnection) -> None:
+    cur = connection.execute(
+        """select
+            'infinity'::date,
+            'infinity'::timestamp,
+            'infinity'::timestamptz,
+            '-infinity'::date,
+            '-infinity'::timestamp,
+            '-infinity'::timestamptz
+        """
+    )
+    assert cur is not None
+    data = cur.fetchall()
+    assert data == [
+        (
+            date.max,
+            datetime.max,
+            datetime.max,
+            date.min,
+            datetime.min,
+            datetime.min,
+        )
+    ]

tests/test_catalog.py

@@ -0,0 +1,92 @@
+import pytest
+from harlequin.catalog import InteractiveCatalogItem
+from harlequin_postgres.adapter import HarlequinPostgresConnection
+from harlequin_postgres.catalog import (
+    ColumnCatalogItem,
+    DatabaseCatalogItem,
+    RelationCatalogItem,
+    SchemaCatalogItem,
+    TableCatalogItem,
+    ViewCatalogItem,
+)
+
+
+@pytest.fixture
+def connection_with_objects(
+    connection: HarlequinPostgresConnection,
+) -> HarlequinPostgresConnection:
+    connection.execute("create schema one")
+    connection.execute("create table one.foo as select 1 as a, '2' as b")
+    connection.execute("create table one.bar as select 1 as a, '2' as b")
+    connection.execute("create table one.baz as select 1 as a, '2' as b")
+    connection.execute("create schema two")
+    connection.execute("create view two.qux as select * from one.foo")
+    connection.execute("create schema three")
+    # the original connection fixture will clean this up.
+    return connection
+
+
+def test_catalog(connection_with_objects: HarlequinPostgresConnection) -> None:
+    conn = connection_with_objects
+
+    catalog = conn.get_catalog()
+
+    # at least two databases, postgres and test
+    assert len(catalog.items) >= 2
+
+    [test_db_item] = filter(lambda item: item.label == "test", catalog.items)
+    assert isinstance(test_db_item, InteractiveCatalogItem)
+    assert isinstance(test_db_item, DatabaseCatalogItem)
+    assert not test_db_item.children
+    assert not test_db_item.loaded
+
+    schema_items = test_db_item.fetch_children()
+    assert all(isinstance(item, SchemaCatalogItem) for item in schema_items)
+
+    [schema_one_item] = filter(lambda item: item.label == "one", schema_items)
+    assert isinstance(schema_one_item, SchemaCatalogItem)
+    assert not schema_one_item.children
+    assert not schema_one_item.loaded
+
+    table_items = schema_one_item.fetch_children()
+    assert all(isinstance(item, RelationCatalogItem) for item in table_items)
+
+    [foo_item] = filter(lambda item: item.label == "foo", table_items)
+    assert isinstance(foo_item, TableCatalogItem)
+    assert not foo_item.children
+    assert not foo_item.loaded
+
+    foo_column_items = foo_item.fetch_children()
+    assert all(isinstance(item, ColumnCatalogItem) for item in foo_column_items)
+
+    [schema_two_item] = filter(lambda item: item.label == "two", schema_items)
+    assert isinstance(schema_two_item, SchemaCatalogItem)
+    assert not schema_two_item.children
+    assert not schema_two_item.loaded
+
+    view_items = schema_two_item.fetch_children()
+    assert all(isinstance(item, ViewCatalogItem) for item in view_items)
+
+    [qux_item] = filter(lambda item: item.label == "qux", view_items)
+    assert isinstance(qux_item, ViewCatalogItem)
+    assert not qux_item.children
+    assert not qux_item.loaded
+
+    qux_column_items = qux_item.fetch_children()
+    assert all(isinstance(item, ColumnCatalogItem) for item in qux_column_items)
+
+    assert [item.label for item in foo_column_items] == [
+        item.label for item in qux_column_items
+    ]
+
+    # ensure calling fetch_children on cols doesn't raise
+    children_items = foo_column_items[0].fetch_children()
+    assert not children_items
+
+    [schema_three_item] = filter(lambda item: item.label == "three", schema_items)
+    assert isinstance(schema_two_item, SchemaCatalogItem)
+    assert not schema_two_item.children
+    assert not schema_two_item.loaded
+
+    three_children = schema_three_item.fetch_children()
+    assert not three_children