ddl and dml queries

Author: vgvoleg

ydb_dbapi/connection.py

@@ -142,7 +142,7 @@ async def describe(self, table_path: str) -> ydb.TableSchemeEntry:
         return self.driver.table_client.describe_table(abs_table_path)
 
     @handle_ydb_errors
-    async def check_exists(self, table_path: str) -> ydb.SchemeEntry:
+    async def check_exists(self, table_path: str) -> bool:
         abs_table_path = posixpath.join(
             self.database, self.table_path_prefix, table_path
         )
@@ -154,7 +154,7 @@ async def get_table_names(self) -> List[str]:
         names = await self._get_table_names(abs_dir_path)
         return [posixpath.relpath(path, abs_dir_path) for path in names]
 
-    async def _check_path_exists(self, table_path: str) -> ydb.SchemeEntry:
+    async def _check_path_exists(self, table_path: str) -> bool:
         try:
             await retry_operation_async(
                 self.driver.scheme_client.describe_path, table_path

ydb_dbapi/cursors.py

@@ -1,4 +1,5 @@
 import dataclasses
+import typing
 from typing import (
     Any,
     Dict,
@@ -10,6 +11,9 @@
 
 import ydb
 
+if typing.TYPE_CHECKING:
+    from ydb.aio.query.base import AsyncResponseContextIterator
+
 
 ParametersType = Dict[
     str,
@@ -32,9 +36,10 @@ def __init__(
         self,
         session_pool: ydb.aio.QuerySessionPool,
         session: ydb.aio.QuerySession,
-        tx_mode: ydb.BaseQueryTxMode,
-        tx_context: Optional[ydb.QueryTxContext] = None,
+        tx_mode: ydb.aio.QueryTxContext,
+        tx_context: ydb.aio.QueryTxContext,
         table_path_prefix: str = "",
+        autocommit: bool = True,
     ):
         self.arraysize = 1
         self._description = None
@@ -44,18 +49,37 @@ def __init__(
         self._tx_mode = tx_mode
         self._tx_context = tx_context
         self._table_path_prefix = table_path_prefix
+        self._autocommit = autocommit
 
     async def _execute_ddl_query(
-        self, query: YdbQuery, parameters: Optional[ParametersType] = None
+        self, query: str, parameters: Optional[ParametersType] = None
     ) -> List[ydb.convert.ResultSet]:
         return await self._pool.execute_with_retries(
             query=query, parameters=parameters
         )
 
+    async def _execute_dml_query(
+        self, query: str, parameters: Optional[ParametersType] = None
+    ) -> AsyncResponseContextIterator:
+        return await self._tx_context.execute(
+            query=query,
+            parameters=parameters,
+            commit_tx=self._autocommit,
+        )
+
     async def execute(
         self, operation: YdbQuery, parameters: Optional[ParametersType] = None
     ):
-        pass
+        if operation.is_ddl:
+            result_sets = await self._execute_ddl_query(
+                query=operation.yql_text, parameters=parameters
+            )
+        else:
+            result_sets_stream = await self._execute_dml_query(
+                query=operation.yql_text, parameters=parameters
+            )
+
+        return result_sets or result_sets_stream
 
     async def executemany(self):
         pass