Skip to content
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