Skip to content

Commit

Permalink
rename interfaces
Browse files Browse the repository at this point in the history
  • Loading branch information
vgvoleg committed Sep 4, 2024
1 parent e1bafec commit 9489f7a
Show file tree
Hide file tree
Showing 21 changed files with 141 additions and 141 deletions.
2 changes: 1 addition & 1 deletion examples/basic_example_v2/basic_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ def select_with_parameters(pool: ydb.QuerySessionPool, path: str, series_id, sea
# calls instead to avoid additional hops to YDB cluster and allow more efficient
# execution of queries.
def explicit_transaction_control(pool: ydb.QuerySessionPool, path: str, series_id, season_id, episode_id):
def callee(session: ydb.QuerySessionSync):
def callee(session: ydb.QuerySession):
query = f"""
PRAGMA TablePathPrefix("{path}");
UPDATE episodes
Expand Down
18 changes: 9 additions & 9 deletions examples/basic_example_v2/basic_example_async.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@
"""


async def fill_tables_with_data(pool: ydb.aio.QuerySessionPoolAsync, path: str):
async def fill_tables_with_data(pool: ydb.aio.QuerySessionPool, path: str):
print("\nFilling tables with data...")

query = FillDataQuery.format(path)
Expand All @@ -84,7 +84,7 @@ async def fill_tables_with_data(pool: ydb.aio.QuerySessionPoolAsync, path: str):
)


async def select_simple(pool: ydb.aio.QuerySessionPoolAsync, path: str):
async def select_simple(pool: ydb.aio.QuerySessionPool, path: str):
print("\nCheck series table...")
result_sets = await pool.execute_with_retries(
f"""
Expand All @@ -111,7 +111,7 @@ async def select_simple(pool: ydb.aio.QuerySessionPoolAsync, path: str):
return first_set


async def upsert_simple(pool: ydb.aio.QuerySessionPoolAsync, path: str):
async def upsert_simple(pool: ydb.aio.QuerySessionPool, path: str):
print("\nPerforming UPSERT into episodes...")

await pool.execute_with_retries(
Expand All @@ -122,7 +122,7 @@ async def upsert_simple(pool: ydb.aio.QuerySessionPoolAsync, path: str):
)


async def select_with_parameters(pool: ydb.aio.QuerySessionPoolAsync, path: str, series_id, season_id, episode_id):
async def select_with_parameters(pool: ydb.aio.QuerySessionPool, path: str, series_id, season_id, episode_id):
result_sets = await pool.execute_with_retries(
f"""
PRAGMA TablePathPrefix("{path}");
Expand Down Expand Up @@ -152,9 +152,9 @@ async def select_with_parameters(pool: ydb.aio.QuerySessionPoolAsync, path: str,
# calls instead to avoid additional hops to YDB cluster and allow more efficient
# execution of queries.
async def explicit_transaction_control(
pool: ydb.aio.QuerySessionPoolAsync, path: str, series_id, season_id, episode_id
pool: ydb.aio.QuerySessionPool, path: str, series_id, season_id, episode_id
):
async def callee(session: ydb.aio.QuerySessionAsync):
async def callee(session: ydb.aio.QuerySession):
query = f"""
PRAGMA TablePathPrefix("{path}");
UPDATE episodes
Expand Down Expand Up @@ -185,12 +185,12 @@ async def callee(session: ydb.aio.QuerySessionAsync):
return await pool.retry_operation_async(callee)


async def drop_tables(pool: ydb.aio.QuerySessionPoolAsync, path: str):
async def drop_tables(pool: ydb.aio.QuerySessionPool, path: str):
print("\nCleaning up existing tables...")
await pool.execute_with_retries(DropTablesQuery.format(path))


async def create_tables(pool: ydb.aio.QuerySessionPoolAsync, path: str):
async def create_tables(pool: ydb.aio.QuerySessionPool, path: str):
print("\nCreating table series...")
await pool.execute_with_retries(
f"""
Expand Down Expand Up @@ -266,7 +266,7 @@ async def run(endpoint, database, path):
) as driver:
await driver.wait(timeout=5, fail_fast=True)

async with ydb.aio.QuerySessionPoolAsync(driver) as pool:
async with ydb.aio.QuerySessionPool(driver) as pool:

await ensure_path_exists(driver, database, path)

Expand Down
2 changes: 1 addition & 1 deletion examples/query-service/basic_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ def callee(session):

pool.retry_operation_sync(callee)

def callee(session: ydb.QuerySessionSync):
def callee(session: ydb.QuerySession):
query_print = """select $a"""

print("=" * 50)
Expand Down
4 changes: 2 additions & 2 deletions examples/query-service/basic_example_asyncio.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ async def main():
except TimeoutError:
raise RuntimeError("Connect failed to YDB")

pool = ydb.aio.QuerySessionPoolAsync(driver)
pool = ydb.aio.QuerySessionPool(driver)

print("=" * 50)
print("DELETE TABLE IF EXISTS")
Expand Down Expand Up @@ -83,7 +83,7 @@ async def callee(session):

await pool.retry_operation_async(callee)

async def callee(session: ydb.aio.QuerySessionAsync):
async def callee(session: ydb.aio.QuerySession):
query_print = """select $a"""

print("=" * 50)
Expand Down
8 changes: 4 additions & 4 deletions tests/aio/query/conftest.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
import pytest
from ydb.aio.query.session import QuerySessionAsync
from ydb.aio.query.pool import QuerySessionPoolAsync
from ydb.aio.query.session import QuerySession
from ydb.aio.query.pool import QuerySessionPool


@pytest.fixture
async def session(driver):
session = QuerySessionAsync(driver)
session = QuerySession(driver)

yield session

Expand All @@ -30,5 +30,5 @@ async def tx(session):

@pytest.fixture
async def pool(driver):
async with QuerySessionPoolAsync(driver) as pool:
async with QuerySessionPool(driver) as pool:
yield pool
30 changes: 15 additions & 15 deletions tests/aio/query/test_query_session.py
Original file line number Diff line number Diff line change
@@ -1,22 +1,22 @@
import pytest
from ydb.aio.query.session import QuerySessionAsync
from ydb.aio.query.session import QuerySession


def _check_session_state_empty(session: QuerySessionAsync):
def _check_session_state_empty(session: QuerySession):
assert session._state.session_id is None
assert session._state.node_id is None
assert not session._state.attached


def _check_session_state_full(session: QuerySessionAsync):
def _check_session_state_full(session: QuerySession):
assert session._state.session_id is not None
assert session._state.node_id is not None
assert session._state.attached


class TestAsyncQuerySession:
@pytest.mark.asyncio
async def test_session_normal_lifecycle(self, session: QuerySessionAsync):
async def test_session_normal_lifecycle(self, session: QuerySession):
_check_session_state_empty(session)

await session.create()
Expand All @@ -26,7 +26,7 @@ async def test_session_normal_lifecycle(self, session: QuerySessionAsync):
_check_session_state_empty(session)

@pytest.mark.asyncio
async def test_second_create_do_nothing(self, session: QuerySessionAsync):
async def test_second_create_do_nothing(self, session: QuerySession):
await session.create()
_check_session_state_full(session)

Expand All @@ -40,30 +40,30 @@ async def test_second_create_do_nothing(self, session: QuerySessionAsync):
assert session._state.node_id == node_id_before

@pytest.mark.asyncio
async def test_second_delete_do_nothing(self, session: QuerySessionAsync):
async def test_second_delete_do_nothing(self, session: QuerySession):
await session.create()

await session.delete()
await session.delete()

@pytest.mark.asyncio
async def test_delete_before_create_not_possible(self, session: QuerySessionAsync):
async def test_delete_before_create_not_possible(self, session: QuerySession):
with pytest.raises(RuntimeError):
await session.delete()

@pytest.mark.asyncio
async def test_create_after_delete_not_possible(self, session: QuerySessionAsync):
async def test_create_after_delete_not_possible(self, session: QuerySession):
await session.create()
await session.delete()
with pytest.raises(RuntimeError):
await session.create()

def test_transaction_before_create_raises(self, session: QuerySessionAsync):
def test_transaction_before_create_raises(self, session: QuerySession):
with pytest.raises(RuntimeError):
session.transaction()

@pytest.mark.asyncio
async def test_transaction_after_delete_raises(self, session: QuerySessionAsync):
async def test_transaction_after_delete_raises(self, session: QuerySession):
await session.create()

await session.delete()
Expand All @@ -72,24 +72,24 @@ async def test_transaction_after_delete_raises(self, session: QuerySessionAsync)
session.transaction()

@pytest.mark.asyncio
async def test_transaction_after_create_not_raises(self, session: QuerySessionAsync):
async def test_transaction_after_create_not_raises(self, session: QuerySession):
await session.create()
session.transaction()

@pytest.mark.asyncio
async def test_execute_before_create_raises(self, session: QuerySessionAsync):
async def test_execute_before_create_raises(self, session: QuerySession):
with pytest.raises(RuntimeError):
await session.execute("select 1;")

@pytest.mark.asyncio
async def test_execute_after_delete_raises(self, session: QuerySessionAsync):
async def test_execute_after_delete_raises(self, session: QuerySession):
await session.create()
await session.delete()
with pytest.raises(RuntimeError):
await session.execute("select 1;")

@pytest.mark.asyncio
async def test_basic_execute(self, session: QuerySessionAsync):
async def test_basic_execute(self, session: QuerySession):
await session.create()
it = await session.execute("select 1;")
result_sets = [result_set async for result_set in it]
Expand All @@ -100,7 +100,7 @@ async def test_basic_execute(self, session: QuerySessionAsync):
assert list(result_sets[0].rows[0].values()) == [1]

@pytest.mark.asyncio
async def test_two_results(self, session: QuerySessionAsync):
async def test_two_results(self, session: QuerySession):
await session.create()
res = []

Expand Down
38 changes: 19 additions & 19 deletions tests/aio/query/test_query_session_pool.py
Original file line number Diff line number Diff line change
@@ -1,42 +1,42 @@
import asyncio
import pytest
import ydb
from ydb.aio.query.pool import QuerySessionPoolAsync
from ydb.aio.query.session import QuerySessionAsync, QuerySessionStateEnum
from ydb.aio.query.pool import QuerySessionPool
from ydb.aio.query.session import QuerySession, QuerySessionStateEnum


class TestQuerySessionPoolAsync:
class TestQuerySessionPool:
@pytest.mark.asyncio
async def test_checkout_provides_created_session(self, pool: QuerySessionPoolAsync):
async def test_checkout_provides_created_session(self, pool: QuerySessionPool):
async with pool.checkout() as session:
assert session._state._state == QuerySessionStateEnum.CREATED

@pytest.mark.asyncio
async def test_oneshot_query_normal(self, pool: QuerySessionPoolAsync):
async def test_oneshot_query_normal(self, pool: QuerySessionPool):
res = await pool.execute_with_retries("select 1;")
assert len(res) == 1

@pytest.mark.asyncio
async def test_oneshot_ddl_query(self, pool: QuerySessionPoolAsync):
async def test_oneshot_ddl_query(self, pool: QuerySessionPool):
await pool.execute_with_retries("drop table if exists Queen;")
await pool.execute_with_retries("create table Queen(key UInt64, PRIMARY KEY (key));")
await pool.execute_with_retries("drop table Queen;")

@pytest.mark.asyncio
async def test_oneshot_query_raises(self, pool: QuerySessionPoolAsync):
async def test_oneshot_query_raises(self, pool: QuerySessionPool):
with pytest.raises(ydb.GenericError):
await pool.execute_with_retries("Is this the real life? Is this just fantasy?")

@pytest.mark.asyncio
async def test_retry_op_uses_created_session(self, pool: QuerySessionPoolAsync):
async def callee(session: QuerySessionAsync):
async def test_retry_op_uses_created_session(self, pool: QuerySessionPool):
async def callee(session: QuerySession):
assert session._state._state == QuerySessionStateEnum.CREATED

await pool.retry_operation_async(callee)

@pytest.mark.asyncio
async def test_retry_op_normal(self, pool: QuerySessionPoolAsync):
async def callee(session: QuerySessionAsync):
async def test_retry_op_normal(self, pool: QuerySessionPool):
async def callee(session: QuerySession):
async with session.transaction() as tx:
iterator = await tx.execute("select 1;", commit_tx=True)
return [result_set async for result_set in iterator]
Expand All @@ -45,18 +45,18 @@ async def callee(session: QuerySessionAsync):
assert len(res) == 1

@pytest.mark.asyncio
async def test_retry_op_raises(self, pool: QuerySessionPoolAsync):
async def test_retry_op_raises(self, pool: QuerySessionPool):
class CustomException(Exception):
pass

async def callee(session: QuerySessionAsync):
async def callee(session: QuerySession):
raise CustomException()

with pytest.raises(CustomException):
await pool.retry_operation_async(callee)

@pytest.mark.asyncio
async def test_pool_size_limit_logic(self, pool: QuerySessionPoolAsync):
async def test_pool_size_limit_logic(self, pool: QuerySessionPool):
target_size = 5
pool._size = target_size
ids = set()
Expand All @@ -78,7 +78,7 @@ async def test_pool_size_limit_logic(self, pool: QuerySessionPoolAsync):
assert pool._current_size == target_size

@pytest.mark.asyncio
async def test_checkout_do_not_increase_size(self, pool: QuerySessionPoolAsync):
async def test_checkout_do_not_increase_size(self, pool: QuerySessionPool):
session_id = None
for _ in range(10):
async with pool.checkout() as session:
Expand All @@ -88,7 +88,7 @@ async def test_checkout_do_not_increase_size(self, pool: QuerySessionPoolAsync):
assert session_id == session._state.session_id

@pytest.mark.asyncio
async def test_pool_recreates_bad_sessions(self, pool: QuerySessionPoolAsync):
async def test_pool_recreates_bad_sessions(self, pool: QuerySessionPool):
async with pool.checkout() as session:
session_id = session._state.session_id
await session.delete()
Expand All @@ -98,20 +98,20 @@ async def test_pool_recreates_bad_sessions(self, pool: QuerySessionPoolAsync):
assert pool._current_size == 1

@pytest.mark.asyncio
async def test_acquire_from_closed_pool_raises(self, pool: QuerySessionPoolAsync):
async def test_acquire_from_closed_pool_raises(self, pool: QuerySessionPool):
await pool.stop()
with pytest.raises(RuntimeError):
await pool.acquire()

@pytest.mark.asyncio
async def test_acquire_with_timeout_from_closed_pool_raises(self, pool: QuerySessionPoolAsync):
async def test_acquire_with_timeout_from_closed_pool_raises(self, pool: QuerySessionPool):
await pool.stop()
with pytest.raises(RuntimeError):
await asyncio.wait_for(pool.acquire(), timeout=0.1)

@pytest.mark.asyncio
async def test_no_session_leak(self, driver, docker_project):
pool = ydb.aio.QuerySessionPoolAsync(driver, 1)
pool = ydb.aio.QuerySessionPool(driver, 1)
docker_project.stop()
try:
await asyncio.wait_for(pool.acquire(), timeout=0.1)
Expand Down
Loading

0 comments on commit 9489f7a

Please sign in to comment.