Skip to content

Hide session management for table client methods #505

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 14 commits into from
Oct 16, 2024
91 changes: 91 additions & 0 deletions tests/aio/test_table_client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
import pytest
import ydb


class TestTableClient:
@pytest.mark.asyncio
async def test_create_table(self, driver: ydb.aio.Driver):
client = driver.table_client
table_name = "/local/testtableclient"
try:
await client.drop_table(table_name)
except ydb.SchemeError:
pass

with pytest.raises(ydb.SchemeError):
await client.describe_table(table_name)

description = (
ydb.TableDescription()
.with_primary_keys("key1", "key2")
.with_columns(
ydb.Column("key1", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("key2", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("value", ydb.OptionalType(ydb.PrimitiveType.Utf8)),
)
)

await client.create_table(table_name, description)

actual_description = await client.describe_table(table_name)

assert actual_description.columns == description.columns

@pytest.mark.asyncio
async def test_alter_table(self, driver: ydb.aio.Driver):
client = driver.table_client

table_name = "/local/testtableclient"
try:
await client.drop_table(table_name)
except ydb.SchemeError:
pass

description = (
ydb.TableDescription()
.with_primary_keys("key1", "key2")
.with_columns(
ydb.Column("key1", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("key2", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("value", ydb.OptionalType(ydb.PrimitiveType.Utf8)),
)
)

await client.create_table(table_name, description)

await client.alter_table(
table_name,
add_columns=[
ydb.Column("value2", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
],
)

description = await client.describe_table(table_name)
assert len(description.columns) == 4

@pytest.mark.asyncio
async def test_copy_table(self, driver: ydb.aio.Driver):
client = driver.table_client
table_name = "/local/testtableclient"
try:
await client.drop_table(table_name)
except ydb.SchemeError:
pass

description = (
ydb.TableDescription()
.with_primary_keys("key1", "key2")
.with_columns(
ydb.Column("key1", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("key2", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("value", ydb.OptionalType(ydb.PrimitiveType.Utf8)),
)
)

await client.create_table(table_name, description)

await client.copy_table(table_name, table_name + "_copy")

copied_description = await client.describe_table(table_name + "_copy")

assert description.columns == copied_description.columns
88 changes: 88 additions & 0 deletions tests/table/test_table_client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
import pytest
import ydb


class TestTableClient:
def test_create_table(self, driver_sync: ydb.Driver):
client = driver_sync.table_client
table_name = "/local/testtableclient"
try:
client.drop_table(table_name)
except ydb.SchemeError:
pass

with pytest.raises(ydb.SchemeError):
client.describe_table(table_name)

description = (
ydb.TableDescription()
.with_primary_keys("key1", "key2")
.with_columns(
ydb.Column("key1", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("key2", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("value", ydb.OptionalType(ydb.PrimitiveType.Utf8)),
)
)

client.create_table(table_name, description)

actual_description = client.describe_table(table_name)

assert actual_description.columns == description.columns

def test_alter_table(self, driver_sync: ydb.Driver):
client = driver_sync.table_client

table_name = "/local/testtableclient"
try:
client.drop_table(table_name)
except ydb.SchemeError:
pass

description = (
ydb.TableDescription()
.with_primary_keys("key1", "key2")
.with_columns(
ydb.Column("key1", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("key2", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("value", ydb.OptionalType(ydb.PrimitiveType.Utf8)),
)
)

client.create_table(table_name, description)

client.alter_table(
table_name,
add_columns=[
ydb.Column("value2", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
],
)

description = client.describe_table(table_name)
assert len(description.columns) == 4

def test_copy_table(self, driver_sync: ydb.Driver):
client = driver_sync.table_client
table_name = "/local/testtableclient"
try:
client.drop_table(table_name)
except ydb.SchemeError:
pass

description = (
ydb.TableDescription()
.with_primary_keys("key1", "key2")
.with_columns(
ydb.Column("key1", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("key2", ydb.OptionalType(ydb.PrimitiveType.Uint64)),
ydb.Column("value", ydb.OptionalType(ydb.PrimitiveType.Utf8)),
)
)

client.create_table(table_name, description)

client.copy_table(table_name, table_name + "_copy")

copied_description = client.describe_table(table_name + "_copy")

assert description.columns == copied_description.columns
4 changes: 4 additions & 0 deletions ydb/aio/driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,3 +59,7 @@ def __init__(
self.scheme_client = scheme.SchemeClient(self)
self.table_client = table.TableClient(self, config.table_client_settings)
self.topic_client = topic.TopicClientAsyncIO(self, config.topic_client_settings)

async def stop(self, timeout=10):
await self.table_client._stop_pool_if_needed(timeout=timeout)
await super().stop(timeout=timeout)
Loading
Loading