Skip to content

Commit d05ffc1

Browse files
committed
Ability to get view names
1 parent ac0cb5d commit d05ffc1

File tree

5 files changed

+100
-12
lines changed

5 files changed

+100
-12
lines changed

examples/basic_example.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
import ydb_dbapi as dbapi
2+
3+
4+
def main() -> None:
5+
connection = dbapi.connect(
6+
host="localhost",
7+
port=2136,
8+
database="/local",
9+
)
10+
11+
print(f"Existing tables: {connection.get_table_names()}")
12+
print(f"Existing views: {connection.get_view_names()}")
13+
# TODO: fill example
14+
15+
connection.close()
16+
17+
18+
if __name__ == "__main__":
19+
main()

poetry.lock

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ force-single-line = true
7878

7979
[tool.ruff.lint.per-file-ignores]
8080
"**/test_*.py" = ["S", "SLF", "ANN201", "ARG", "PLR2004", "PT012"]
81+
"examples/*.py" = ["T201", "INP001"]
8182
"conftest.py" = ["S", "ARG001"]
8283
"__init__.py" = ["F401", "F403"]
8384

tests/test_connections.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -239,6 +239,41 @@ def _test_error_with_interactive_tx(
239239
maybe_await(cur.close())
240240
maybe_await(connection.rollback())
241241

242+
def _test_get_view_names(
243+
self,
244+
connection: dbapi.Connection,
245+
) -> None:
246+
cur = connection.cursor()
247+
248+
maybe_await(
249+
cur.execute_scheme(
250+
"""
251+
DROP VIEW if exists test_view;
252+
"""
253+
)
254+
)
255+
256+
res = maybe_await(connection.get_view_names())
257+
258+
assert len(res) == 0
259+
260+
maybe_await(
261+
cur.execute_scheme(
262+
"""
263+
CREATE VIEW test_view WITH (security_invoker = TRUE) as (
264+
select 1 as res
265+
);
266+
"""
267+
)
268+
)
269+
270+
res = maybe_await(connection.get_view_names())
271+
272+
assert len(res) == 1
273+
assert res[0] == "test_view"
274+
275+
maybe_await(cur.close())
276+
242277

243278
class TestConnection(BaseDBApiTestSuit):
244279
@pytest.fixture
@@ -289,6 +324,11 @@ def test_errors_with_interactive_tx(
289324
) -> None:
290325
self._test_error_with_interactive_tx(connection)
291326

327+
def test_get_view_names(
328+
self, connection: dbapi.Connection
329+
) -> None:
330+
self._test_get_view_names(connection)
331+
292332

293333
class TestAsyncConnection(BaseDBApiTestSuit):
294334
@pytest_asyncio.fixture
@@ -358,3 +398,9 @@ async def test_errors_with_interactive_tx(
358398
self, connection: dbapi.AsyncConnection
359399
) -> None:
360400
await greenlet_spawn(self._test_error_with_interactive_tx, connection)
401+
402+
@pytest.mark.asyncio
403+
async def test_get_view_names(
404+
self, connection: dbapi.AsyncConnection
405+
) -> None:
406+
await greenlet_spawn(self._test_get_view_names, connection)

ydb_dbapi/connections.py

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -277,7 +277,13 @@ def check_exists(self, table_path: str) -> bool:
277277
@handle_ydb_errors
278278
def get_table_names(self) -> list[str]:
279279
abs_dir_path = posixpath.join(self.database, self.table_path_prefix)
280-
names = self._get_table_names(abs_dir_path)
280+
names = self._get_entity_names(abs_dir_path, ydb.SchemeEntryType.TABLE)
281+
return [posixpath.relpath(path, abs_dir_path) for path in names]
282+
283+
@handle_ydb_errors
284+
def get_view_names(self) -> list[str]:
285+
abs_dir_path = posixpath.join(self.database, self.table_path_prefix)
286+
names = self._get_entity_names(abs_dir_path, ydb.SchemeEntryType.VIEW)
281287
return [posixpath.relpath(path, abs_dir_path) for path in names]
282288

283289
def _check_path_exists(self, table_path: str) -> bool:
@@ -295,7 +301,7 @@ def callee() -> None:
295301
else:
296302
return True
297303

298-
def _get_table_names(self, abs_dir_path: str) -> list[str]:
304+
def _get_entity_names(self, abs_dir_path: str, etype: int) -> list[str]:
299305
settings = self._get_request_settings()
300306

301307
def callee() -> ydb.Directory:
@@ -308,10 +314,10 @@ def callee() -> ydb.Directory:
308314
result = []
309315
for child in directory.children:
310316
child_abs_path = posixpath.join(abs_dir_path, child.name)
311-
if child.is_table():
317+
if child.type == etype:
312318
result.append(child_abs_path)
313319
elif child.is_directory() and not child.name.startswith("."):
314-
result.extend(self._get_table_names(child_abs_path))
320+
result.extend(self._get_entity_names(child_abs_path, etype))
315321
return result
316322

317323
@handle_ydb_errors
@@ -452,7 +458,19 @@ async def check_exists(self, table_path: str) -> bool:
452458
@handle_ydb_errors
453459
async def get_table_names(self) -> list[str]:
454460
abs_dir_path = posixpath.join(self.database, self.table_path_prefix)
455-
names = await self._get_table_names(abs_dir_path)
461+
names = await self._get_entity_names(
462+
abs_dir_path,
463+
ydb.SchemeEntryType.TABLE,
464+
)
465+
return [posixpath.relpath(path, abs_dir_path) for path in names]
466+
467+
@handle_ydb_errors
468+
async def get_view_names(self) -> list[str]:
469+
abs_dir_path = posixpath.join(self.database, self.table_path_prefix)
470+
names = await self._get_entity_names(
471+
abs_dir_path,
472+
ydb.SchemeEntryType.VIEW,
473+
)
456474
return [posixpath.relpath(path, abs_dir_path) for path in names]
457475

458476
async def _check_path_exists(self, table_path: str) -> bool:
@@ -471,7 +489,9 @@ async def callee() -> None:
471489
else:
472490
return True
473491

474-
async def _get_table_names(self, abs_dir_path: str) -> list[str]:
492+
async def _get_entity_names(
493+
self, abs_dir_path: str, etype: int
494+
) -> list[str]:
475495
settings = self._get_request_settings()
476496

477497
async def callee() -> ydb.Directory:
@@ -484,10 +504,12 @@ async def callee() -> ydb.Directory:
484504
result = []
485505
for child in directory.children:
486506
child_abs_path = posixpath.join(abs_dir_path, child.name)
487-
if child.is_table():
507+
if child.type == etype:
488508
result.append(child_abs_path)
489509
elif child.is_directory() and not child.name.startswith("."):
490-
result.extend(await self._get_table_names(child_abs_path))
510+
result.extend(
511+
await self._get_entity_names(child_abs_path, etype)
512+
)
491513
return result
492514

493515
@handle_ydb_errors

0 commit comments

Comments
 (0)