Display catalog information on the schema pane when connecting to Trino (#6578)

This commit is contained in:
Masayuki Takahashi
2023-11-07 10:58:12 +09:00
committed by GitHub
parent 63cef6632e
commit 198b422eaf
2 changed files with 7 additions and 18 deletions

View File

@@ -103,9 +103,7 @@ class Trino(BaseQueryRunner):
results = json_loads(results) results = json_loads(results)
for row in results["rows"]: for row in results["rows"]:
table_name = f'{row["table_schema"]}.{row["table_name"]}' table_name = f'{catalog}.{row["table_schema"]}.{row["table_name"]}'
if not self.configuration.get("catalog"):
table_name = f"{catalog}." + table_name
if table_name not in schema: if table_name not in schema:
schema[table_name] = {"name": table_name, "columns": []} schema[table_name] = {"name": table_name, "columns": []}

View File

@@ -17,34 +17,25 @@ class TestTrino(TestCase):
@patch.object(Trino, "_get_catalogs") @patch.object(Trino, "_get_catalogs")
@patch.object(Trino, "run_query") @patch.object(Trino, "run_query")
def test_get_schema_no_catalog_set(self, mock_run_query, mock__get_catalogs): def test_get_schema_no_catalog_set(self, mock_run_query, mock__get_catalogs):
mock_run_query.return_value = (
f'{{"rows": [{{"table_schema": "{TestTrino.schema_name}", "table_name": "{TestTrino.table_name}", "column_name": "{TestTrino.column_name}", "data_type": "{TestTrino.column_type}"}}]}}',
None,
)
mock__get_catalogs.return_value = [TestTrino.catalog_name]
runner = Trino({}) runner = Trino({})
schema = runner.get_schema() self._assert_schema_catalog(mock_run_query, mock__get_catalogs, runner)
expected_schema = [
{
"name": f"{TestTrino.catalog_name}.{TestTrino.schema_name}.{TestTrino.table_name}",
"columns": [{"name": f"{TestTrino.column_name}", "type": f"{TestTrino.column_type}"}],
}
]
self.assertEqual(schema, expected_schema)
@patch.object(Trino, "_get_catalogs") @patch.object(Trino, "_get_catalogs")
@patch.object(Trino, "run_query") @patch.object(Trino, "run_query")
def test_get_schema_catalog_set(self, mock_run_query, mock__get_catalogs): def test_get_schema_catalog_set(self, mock_run_query, mock__get_catalogs):
runner = Trino({"catalog": TestTrino.catalog_name})
self._assert_schema_catalog(mock_run_query, mock__get_catalogs, runner)
def _assert_schema_catalog(self, mock_run_query, mock__get_catalogs, runner):
mock_run_query.return_value = ( mock_run_query.return_value = (
f'{{"rows": [{{"table_schema": "{TestTrino.schema_name}", "table_name": "{TestTrino.table_name}", "column_name": "{TestTrino.column_name}", "data_type": "{TestTrino.column_type}"}}]}}', f'{{"rows": [{{"table_schema": "{TestTrino.schema_name}", "table_name": "{TestTrino.table_name}", "column_name": "{TestTrino.column_name}", "data_type": "{TestTrino.column_type}"}}]}}',
None, None,
) )
mock__get_catalogs.return_value = [TestTrino.catalog_name] mock__get_catalogs.return_value = [TestTrino.catalog_name]
runner = Trino({"catalog": TestTrino.catalog_name})
schema = runner.get_schema() schema = runner.get_schema()
expected_schema = [ expected_schema = [
{ {
"name": f"{TestTrino.schema_name}.{TestTrino.table_name}", "name": f"{TestTrino.catalog_name}.{TestTrino.schema_name}.{TestTrino.table_name}",
"columns": [{"name": f"{TestTrino.column_name}", "type": f"{TestTrino.column_type}"}], "columns": [{"name": f"{TestTrino.column_name}", "type": f"{TestTrino.column_type}"}],
} }
] ]