Skip to content

Commit 8e78525

Browse files
fix(starrocks): reuse shared connection for table comments (#27030)
1 parent c827b6a commit 8e78525

2 files changed

Lines changed: 34 additions & 12 deletions

File tree

ingestion/src/metadata/ingestion/source/database/starrocks/metadata.py

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -282,24 +282,20 @@ def query_view_names_and_types(
282282
return tables
283283

284284
def get_table_description(
285-
self, schema_name: str, table_name: str, inspector: Inspector
285+
self, schema_name: str, table_name: str, _inspector: Inspector
286286
) -> Optional[str]:
287-
description = None
288287
try:
289-
with self.engine.connect() as conn:
290-
rows = conn.execute(
291-
sql.text(STARROCKS_TABLE_COMMENTS),
292-
{"table_name": table_name, "schema": schema_name},
293-
)
294-
for row in rows:
295-
description = row[0]
296-
break
288+
row = self.connection.execute(
289+
sql.text(STARROCKS_TABLE_COMMENTS),
290+
{"table_name": table_name, "schema": schema_name},
291+
).fetchone()
292+
return row[0] if row else None
297293
except Exception as exc: # pylint: disable=broad-except
298294
logger.debug(traceback.format_exc())
299295
logger.warning(
300296
f"Table description error for table [{schema_name}.{table_name}]: {exc}"
301297
)
302-
return description or None
298+
return None
303299

304300
def _get_columns(self, table_name, schema=None):
305301
"""Get column information and primary key columns of the specified table"""

ingestion/tests/unit/topology/database/test_starrocks.py

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
"""
1515

1616
from unittest import TestCase
17-
from unittest.mock import patch
17+
from unittest.mock import Mock, patch
1818

1919
import pytest
2020
from sqlalchemy import types as sqltypes
@@ -114,6 +114,32 @@ def test_close_connection(self, engine, connection):
114114
connection.return_value = True
115115
self.starrocks_source.close()
116116

117+
def test_get_table_description_returns_comment(self):
118+
mock_result = Mock()
119+
mock_result.fetchone.return_value = ("table comment",)
120+
mock_connection = Mock()
121+
mock_connection.execute.return_value = mock_result
122+
self.starrocks_source.connection = mock_connection
123+
124+
result = self.starrocks_source.get_table_description(
125+
schema_name="public", table_name="my_table", _inspector=Mock()
126+
)
127+
128+
self.assertEqual(result, "table comment")
129+
130+
def test_get_table_description_returns_none_when_missing(self):
131+
mock_result = Mock()
132+
mock_result.fetchone.return_value = None
133+
mock_connection = Mock()
134+
mock_connection.execute.return_value = mock_result
135+
self.starrocks_source.connection = mock_connection
136+
137+
result = self.starrocks_source.get_table_description(
138+
schema_name="public", table_name="my_table", _inspector=Mock()
139+
)
140+
141+
self.assertIsNone(result)
142+
117143

118144
class StarRocksSSLUnitTest(TestCase):
119145
@patch(

0 commit comments

Comments
 (0)