diff --git a/src/harlequin_bigquery/adapter.py b/src/harlequin_bigquery/adapter.py index 9147ccf..b8aa9ec 100644 --- a/src/harlequin_bigquery/adapter.py +++ b/src/harlequin_bigquery/adapter.py @@ -183,8 +183,8 @@ def get_catalog(self) -> Catalog: datasets[dataset_id].children.append(table_catalog_item) if column_name: - # remove anything in <> from the column_type - column_type_cleaned = re.sub(r"\<.*\>", "", row.column_type) + # remove anything in <> or () from the column_type + column_type_cleaned = re.sub(r"[<(].*[>)]", "", row.column_type) column_type_label = COLUMN_TYPE_MAPPING[ StandardSqlTypeNames(column_type_cleaned) ] diff --git a/tests/test_adapter.py b/tests/test_adapter.py index e250828..2a6a231 100644 --- a/tests/test_adapter.py +++ b/tests/test_adapter.py @@ -1,6 +1,8 @@ import sys +from typing import Any, Generator import pytest +from google.cloud.bigquery.client import Client from google.cloud.bigquery.enums import StandardSqlTypeNames from harlequin.adapter import HarlequinAdapter, HarlequinConnection, HarlequinCursor from harlequin.catalog import Catalog, CatalogItem @@ -18,6 +20,19 @@ from importlib.metadata import entry_points +@pytest.fixture(scope="session") +def bigquery_client() -> Client: + return Client() + + +@pytest.fixture(scope="module") +def test_dataset(bigquery_client: Client) -> Generator[str, None, None]: + dataset_name = "tmp__harlequin_bigquery_test" + bigquery_client.create_dataset(dataset_name, exists_ok=True) + yield dataset_name + bigquery_client.delete_dataset(dataset_name, delete_contents=True) + + def test_plugin_discovery() -> None: PLUGIN_NAME = "bigquery-adapter" eps = entry_points(group="harlequin.adapter") @@ -48,6 +63,19 @@ def test_get_catalog(connection: BigQueryConnection) -> None: assert isinstance(catalog.items[0], CatalogItem) +def test_get_catalog_with_parameterized_types( + connection: BigQueryConnection, test_dataset: str +) -> None: + # create a temporary table with a parametrized type + query = f"CREATE TABLE {test_dataset}.test_table (id NUMERIC(9, 8));" + cursor = connection.execute(query) + try: + catalog = connection.get_catalog() + finally: + connection.execute(f"DROP TABLE {test_dataset}.test_table;") + assert True + + def test_execute_select(connection: BigQueryConnection) -> None: cur = connection.execute("select 1 as a") assert isinstance(cur, HarlequinCursor)