e6data-python-connector 3.0.0__py3-none-any.whl → 3.0.1rc2__py3-none-any.whl

e6data_python_connector/async_connection.py

@@ -50,7 +50,11 @@ class AsyncConnection:
                  token_url=None, oauth_scope=None, access_token=None,
                  client_auth_method='basic', *, operation_timeout=600.0,
                  oauth_timeout=10.0, cleanup_timeout=10.0, auto_resume_timeout=300.0,
-                 max_receive_message_bytes=64 * 1024 * 1024):
+                 max_receive_message_bytes=64 * 1024 * 1024,
+                 enable_result_batch_v2=False):
+        from .result_decode_worker import is_decode_worker
+        if is_decode_worker():
+            raise RuntimeError('A result decode worker cannot create a connection.')
         if sys.version_info < (3, 11):
             raise RuntimeError('The async API requires Python 3.11 or newer.')
         try:
@@ -78,13 +82,19 @@ class AsyncConnection:
             raise ValueError('Async OAuth requires verified TLS (secure=True).')
         if ssl_cert is not None and not isinstance(ssl_cert, (str, bytes)):
             raise ValueError('ssl_cert must be a certificate path or PEM bytes.')
+        if not isinstance(enable_result_batch_v2, bool):
+            raise ValueError('enable_result_batch_v2 must be a boolean.')
+        if enable_result_batch_v2:
+            from .result_decode import validate_decode_runtime
+            validate_decode_runtime()
         settings = dict(host=host, port=port, username=username, password=password,
                         catalog=catalog, database=database, cluster_name=cluster_name,
                         secure=bool(secure), ssl_cert=ssl_cert, auto_resume=auto_resume,
                         scheme=scheme, debug=debug, require_fastbinary=require_fastbinary,
                         client_id=client_id, client_secret=client_secret, token_url=token_url,
                         oauth_scope=oauth_scope, access_token=access_token,
-                        client_auth_method=client_auth_method)
+                        client_auth_method=client_auth_method,
+                        enable_result_batch_v2=enable_result_batch_v2)
         for name, value in [('operation_timeout', operation_timeout), ('oauth_timeout', oauth_timeout),
                             ('cleanup_timeout', cleanup_timeout), ('auto_resume_timeout', auto_resume_timeout)]:
             settings[name] = validate_positive_timeout(value, name)
@@ -103,6 +113,9 @@ class AsyncConnection:
         settings['grpc_options'] = MappingProxyType(options)
         settings['max_receive_message_bytes'] = max_receive_message_bytes
         self._config = MappingProxyType(settings)
+        if debug:
+            from .e6data_grpc import _configure_debug_logging
+            _configure_debug_logging()
         self._oauth_enabled = not modes[0]
         self._channel = self._client = self._token_provider = None
         self._owns_token_provider = True
@@ -125,6 +138,11 @@ class AsyncConnection:
         self._cleanup_error = None
         self._cluster_manager = None
         self._owns_cluster_manager = True
+        self._decoder_started = False
+        self._decoder_lease = None
+        if enable_result_batch_v2:
+            from .result_decode import DecoderLease
+            self._decoder_lease = DecoderLease()
 
     def __getattr__(self, name):
         config = self.__dict__.get('_config', {})
@@ -300,6 +318,24 @@ class AsyncConnection:
         self._check_owner(_cleanup=_cleanup, _internal=True)
         return metadata
 
+    async def _ensure_decoder(self, deadline):
+        self._check_owner()
+        if self._decoder_lease is not None and not self._decoder_started:
+            from .async_work import run_blocking
+            await run_blocking(self._decoder_lease.start, deadline, deadline=deadline)
+            self._check_owner()
+            self._decoder_started = True
+
+    async def _raw_result_rpc(self, method_name, request, metadata, deadline):
+        """Transport only: the owning task already froze identity and metadata."""
+        call = getattr(self._client, method_name)(
+            request, metadata=metadata, timeout=self._remaining(deadline))
+        self._calls.add(call)
+        try:
+            return await call
+        finally:
+            self._calls.discard(call)
+
     async def _rpc(self, method_name, request, *, deadline, route=None, safe_retry=False, _cleanup=False, _on_dispatch=None, _response_metadata=None):
         self._check_owner(_cleanup=_cleanup, _internal=True)
         refreshed = switched = False
@@ -500,6 +536,15 @@ class AsyncConnection:
             # Native channel close with no grace cancels transports without waiting.
             if self._channel is not None:
                 await self._channel.close()
+            if self._decoder_lease is not None:
+                # Retirement must survive loop shutdown and saturated work slots.
+                self._decoder_lease.retire(deadline)
+                try:
+                    async with asyncio.timeout_at(deadline):
+                        while self._decoder_lease.cleanup_pending:
+                            await asyncio.sleep(min(0.005, self._remaining(deadline)))
+                except (Exception, asyncio.CancelledError):
+                    self._cleanup_error = OperationalError('Result decoder cleanup remains pending.')
             for resource, owned in ((self._cluster_manager, self._owns_cluster_manager),
                                     (self._token_provider, self._owns_token_provider)):
                 if owned and resource is not None:
@@ -538,6 +583,12 @@ class AsyncConnection:
         if self._owner is not None and self._owner != (os.getpid(), threading.get_ident(), asyncio.get_running_loop()):
             raise ProgrammingError('Connection belongs to another owner.')
         await self.close()
+        if self._decoder_lease is not None and self._decoder_lease.cleanup_pending:
+            raise ProgrammingError('Cannot reopen while result decoder cleanup remains pending.')
+        if self.enable_result_batch_v2:
+            from .result_decode import DecoderLease
+            self._decoder_lease = DecoderLease()
+            self._decoder_started = False
         self._state = 'new'
         self._session_id = self._session_task = self._close_task = self._open_task = None
         if self._owns_token_provider:

e6data_python_connector/async_connection_pool.py

@@ -316,6 +316,8 @@ class AsyncConnectionPool:
                     healthy = healthy and cursor._state != 'SUBMISSION_UNKNOWN'
                     await cursor._close_owned(deadline)
                     healthy = healthy and cursor.cleanup_error is None
+                    healthy = (healthy and cursor._pending_result is None
+                               and not cursor._retired_results)
                 healthy = (healthy and not connection._routes and not connection._ambiguous_submissions
                            and connection._state == 'open')
         except (Exception, asyncio.CancelledError):

e6data_python_connector/async_cursor.py

@@ -1,5 +1,7 @@
 """Native asynchronous query lifecycle, with no replay of consumed results."""
 import asyncio
+import logging
+import time
 from collections import deque
 from contextlib import contextmanager
 from io import BytesIO
@@ -11,6 +13,28 @@ from .e6data_grpc import _escaper
 from .exceptions import (AmbiguousSubmissionError, IncompleteResultError,
                          NotSupportedError, OperationalError, ProgrammingError, OAuthError)
 from .server import e6x_engine_pb2 as pb
+from .result_batch import ResultBatchBuffer, decode_result_batches
+from .result_prefetch import ContinuousResultStream, PendingFetch
+
+
+logger = logging.getLogger(__name__)
+
+
+def _log_result_batch(protocol, rpc_seconds, status, response=None, decode_seconds=0.0):
+    chunk_count = ((len(response.resultBatches) if protocol == 'v2'
+                    else int(bool(response.resultBatch))) if response else 0)
+    serialized_bytes = response.ByteSize() if response else 0
+    logger.debug(
+        'Result batch fetch protocol=%s rpc_seconds=%.6f decode_seconds=%.6f '
+        'chunk_count=%d serialized_bytes=%d status=%s',
+        protocol, rpc_seconds, decode_seconds, chunk_count, serialized_bytes, status, extra={
+        'result_batch_protocol': protocol,
+        'result_batch_rpc_seconds': rpc_seconds,
+        'result_batch_chunk_count': chunk_count,
+        'result_batch_serialized_bytes': serialized_bytes,
+        'result_batch_decode_seconds': decode_seconds,
+        'result_batch_status': status,
+    })
 
 
 def _metadata_columns(payload):
@@ -46,8 +70,20 @@ class AsyncCursor:
         self._active_call = None
         self._close_task = None
         self._cleanup_error = None
+        self._result_batches = ResultBatchBuffer()
+        self._pending_result = None
+        self._result_available = None
+        self._retired_results = []
+        self._decode_owner = None
+        self._reset_result_protocol()
         connection._cursors.add(self)
 
+    def _reset_result_protocol(self):
+        self._result_batch_v2 = self._connection.enable_result_batch_v2
+        self._result_protocol = 'v2' if self._result_batch_v2 else 'v1'
+        self._result_session_id = None
+        self._result_batches.clear()
+
     @property
     def connection(self):
         return self._public_connection
@@ -133,10 +169,14 @@ class AsyncCursor:
                 self._active_call = None
 
     def _fail_result(self, reason):
+        self._retire_prefetch()
+        self._cancel_decode()
         if self._failure is None:
             self._failure = IncompleteResultError(reason, query_id=self.query_id)
         if self._state != 'CLOSED':
             self._state = 'RESULT_FAILED'
+        self._rows.clear()
+        self._result_batches.clear()
         return self._failure
 
     def _accept_rows(self, rows):
@@ -147,10 +187,18 @@ class AsyncCursor:
         self._rownumber += len(rows)
         return rows
 
-    async def _query_request(self, cls, deadline):
+    async def _query_request(self, cls, deadline, *, owned=False):
         if self._route is None:
             raise ProgrammingError('No query handle is known.')
-        session = await self._connection.get_session_id(deadline=deadline)
+        if self._result_batch_v2 and not self._connection._oauth_enabled and self._result_session_id:
+            session = self._result_session_id
+        elif owned:
+            # Owned cleanup cannot mint a replacement authentication session.
+            session = '' if self._connection._oauth_enabled else self._connection._session_id
+            if session is None:
+                raise OperationalError('Cleanup has no authenticated legacy session.')
+        else:
+            session = await self._connection.get_session_id(deadline=deadline)
         return cls(sessionId=session, queryId=self.query_id, engineIP=self._route.engine_ip)
 
     async def _refresh_metadata(self, deadline):
@@ -168,11 +216,13 @@ class AsyncCursor:
         from .async_connection import QueryRoute
         if self._state != 'EMPTY':
             raise ProgrammingError('Clear the previous query before executing another.')
+        self._reset_result_protocol()
         if not isinstance(operation, str) or not operation.strip():
             raise ProgrammingError('operation must be a nonempty SQL string.')
         sql = operation.strip().removesuffix(';')
         if parameters is not None:
             sql = sql % _escaper.escape_args(parameters)
+        await self._connection._ensure_decoder(deadline)
         session = await self._connection.get_session_id(deadline=deadline)
         v2 = bool(self._catalog)
         request_type = pb.PrepareStatementV2Request if v2 else pb.PrepareStatementRequest
@@ -268,6 +318,8 @@ class AsyncCursor:
 
     async def _next_batch(self, deadline):
         from .async_work import reserve_work
+        if self._result_batch_v2:
+            return await self._next_v2_batch(deadline)
         if self._state == 'EXHAUSTED':
             return None
         if self._columns is None:
@@ -279,17 +331,31 @@ class AsyncCursor:
         def mark_dispatch():
             nonlocal consumed
             consumed = True
+        started = time.monotonic()
+        response = None
+        decode_seconds = 0.0
         try:
-            response = await self._call('getNextResultBatch', request, deadline, on_dispatch=mark_dispatch)
+            try:
+                response = await self._call('getNextResultBatch', request, deadline, on_dispatch=mark_dispatch)
+            finally:
+                rpc_seconds = time.monotonic() - started
             payload = response.resultBatch
             if not payload:
                 self._state = 'EXHAUSTED'
                 return None
-            return await self._decode_batch(payload, reservation, revision)
+            decode_started = time.monotonic()
+            try:
+                return await self._decode_batch(payload, reservation, revision)
+            finally:
+                decode_seconds = time.monotonic() - decode_started
         except asyncio.CancelledError:
             if consumed:
                 self._fail_result('ambiguous_result')
             raise
+        except grpc.RpcError:
+            if consumed:
+                self._fail_result('ambiguous_result')
+            raise
         except IncompleteResultError:
             raise
         except Exception as error:
@@ -298,6 +364,281 @@ class AsyncCursor:
             raise
         finally:
             reservation.release()
+            if consumed:
+                status = 'ok' if response is not None and self._failure is None else 'error'
+                _log_result_batch('v1', rpc_seconds,
+                                  status, response if status == 'ok' else None, decode_seconds if status == 'ok' else 0.0)
+
+    def _retain_result_session(self, response, revision):
+        self._publish(revision)
+        if self._result_batch_v2 and not self._connection._oauth_enabled and response.sessionId:
+            pending = self._pending_result[0] if self._pending_result is not None else None
+            self._result_session_id = (pending.session_id if isinstance(pending, ContinuousResultStream)
+                                       and pending.session_id else response.sessionId)
+
+    async def _decode_v2_response(self, response, reservation, revision, deadline=None):
+        self._retain_result_session(response, revision)
+        owner = None
+        try:
+            if self._connection._decoder_started:
+                owner = self._connection._decoder_lease.new_owner()
+                self._decode_owner = owner
+                chunks = await reservation.run(
+                    self._connection._decoder_lease.decode, tuple(self._columns),
+                    tuple(response.resultBatches), deadline, owner)
+            else:
+                chunks = await reservation.run(decode_result_batches, self._columns, response.resultBatches)
+        except asyncio.CancelledError:
+            self._cancel_decode()
+            raise
+        except Exception as error:
+            raise self._fail_result('decode_failed') from error
+        finally:
+            if self._decode_owner is owner:
+                self._decode_owner = None
+        self._publish(revision)
+        self._result_batches.accept(chunks, response.endOfStream)
+
+    def _cancel_decode(self):
+        if self._decode_owner is not None:
+            self._connection._decoder_lease.cancel(self._decode_owner)
+
+    def _retire_prefetch(self):
+        if self._pending_result is not None:
+            pending, _, _ = self._pending_result
+            self._pending_result = None
+            self._result_available = None
+            pending.retire()
+            self._retired_results.append(pending)
+
+    async def _dispatch_result(self, deadline, permit=None):
+        """Freeze identity in the owner; the background task only runs transport."""
+        revision, route = self._revision, self._route
+        request = await self._query_request(pb.GetNextResultBatchRequest, deadline)
+        metadata = await self._connection._metadata(deadline, route)
+        self._publish(revision)
+        self._connection._remaining(deadline)
+        method = 'getNextResultBatchV2' if self._result_protocol == 'v2' else 'getNextResultBatch'
+        task = asyncio.create_task(self._connection._raw_result_rpc(method, request, tuple(metadata), deadline))
+        pending = PendingFetch(task, permit, session_id=request.sessionId, deadline=deadline)
+        self._pending_result = (pending, revision, route)
+
+    async def _take_pending_result(self, deadline):
+        pending, revision, route = self._pending_result
+        if isinstance(pending, ContinuousResultStream):
+            event = self._result_available
+            while not pending.ready:
+                event.clear()
+                if pending.ready:
+                    break
+                async with asyncio.timeout_at(deadline):
+                    await event.wait()
+        else:
+            handle = pending.handle
+            if handle is not None and not pending.settled:
+                done, _ = await asyncio.wait({handle}, timeout=self._connection._remaining(deadline))
+                if not done:
+                    raise TimeoutError('Result transport exceeded the operation deadline.')
+        self._publish(revision)
+        if route != self._route:
+            raise asyncio.CancelledError()
+        if not self._connection._oauth_enabled and pending.session_id:
+            self._result_session_id = pending.session_id
+        if not isinstance(pending, ContinuousResultStream):
+            self._pending_result = None
+        response = pending.take()
+        new_strategy = getattr(response, 'new_strategy', None)
+        if new_strategy in ('blue', 'green'):
+            self._connection._pending_strategy = new_strategy
+            self._connection._apply_pending_strategy()
+        return response, pending.elapsed
+
+    async def _start_prefetch(self, response, deadline):
+        if (self._result_protocol != 'v2' or response.endOfStream
+                or self._pending_result is not None):
+            return
+        revision, route = self._revision, self._route
+        try:
+            self._publish(revision)
+            self._connection._remaining(deadline)
+        except BaseException:
+            logger.debug('Continuous result download preparation failed.', extra={
+                'result_batch_download_status': 'preparation_failed'})
+            raise
+        event = asyncio.Event()
+        stream = ContinuousResultStream(session_id=self._result_session_id, notify=event.set)
+        producer = asyncio.create_task(self._download_results(stream, revision, route))
+        stream.set_handle(producer)
+        self._result_available = event
+        self._pending_result = (stream, revision, route)
+        logger.debug('Continuous result download started.', extra={'result_batch_download_status': 'started'})
+
+    async def _download_results(self, stream, revision, route):
+        """Download ordered envelopes without using any decode/work reservation."""
+        budget = self._connection.operation_timeout
+        no_progress_deadline = time.monotonic() + budget
+        backoff = .01
+        rpc_started = None
+        try:
+            while not stream.retired:
+                deadline = min(time.monotonic() + budget, no_progress_deadline)
+                rpc_started = None
+                metadata = await self._connection._metadata(deadline, route)
+                if stream.retired or revision != self._revision or route != self._route:
+                    stream.retire()
+                    return
+                if self._lease_guard is not None:
+                    self._lease_guard(internal=True)
+                self._connection._check_owner(_internal=True)
+                self._connection._remaining(deadline)
+                session = '' if self._connection._oauth_enabled else stream.session_id
+                if not self._connection._oauth_enabled and not session:
+                    raise ProgrammingError('An active query session is required to download results.')
+                request = pb.GetNextResultBatchRequest(
+                    sessionId=session, queryId=route.query_id, engineIP=route.engine_ip)
+                rpc_started = time.monotonic()
+                transport = asyncio.create_task(self._connection._raw_result_rpc(
+                    'getNextResultBatchV2', request, tuple(metadata), deadline))
+                stream.set_transport(transport)
+                try:
+                    response = await asyncio.shield(transport)
+                finally:
+                    if transport.done():
+                        stream.clear_transport(transport)
+                stream.push(response, time.monotonic() - rpc_started)
+                if stream.retired or response.endOfStream:
+                    return
+                if any(response.resultBatches):
+                    no_progress_deadline = time.monotonic() + budget
+                    backoff = .01
+                else:
+                    await asyncio.sleep(min(backoff, self._connection._remaining(no_progress_deadline)))
+                    backoff = min(backoff * 2, .1)
+        except asyncio.CancelledError:
+            stream.retire()
+            raise
+        except Exception as error:
+            stream.fail(error, 0.0 if rpc_started is None else time.monotonic() - rpc_started)
+        finally:
+            stream.finish()
+
+    async def _next_v2_batch(self, deadline):
+        from .async_work import reserve_work
+        batch = self._result_batches.pop()
+        if batch is not None:
+            return batch
+        if self._result_batches.finished or self._state == 'EXHAUSTED':
+            self._state = 'EXHAUSTED'
+            return None
+        consumed = self._pending_result is not None
+        backoff = .01
+        try:
+            await self._connection._ensure_decoder(deadline)
+            if self._columns is None:
+                await self._refresh_metadata(deadline)
+            while True:
+                consumed = consumed or self._pending_result is not None
+                reservation = None
+                revision = self._revision
+                protocol = self._result_protocol
+                previously_consumed = consumed
+                dispatched = False
+                response = None
+                status = 'error'
+                rpc_seconds = decode_seconds = 0.0
+                started = time.monotonic()
+                prefetched = self._pending_result is not None
+                transport_record = None
+                transport_status = 'error'
+
+                def mark_dispatch():
+                    nonlocal consumed, dispatched
+                    consumed = dispatched = True
+
+                try:
+                    try:
+                        if self._pending_result is None:
+                            await self._dispatch_result(deadline)
+                        transport_record = self._pending_result[0]
+                        mark_dispatch()
+                        response, rpc_seconds = await self._take_pending_result(deadline)
+                        transport_status = 'ok'
+                    except grpc.RpcError as error:
+                        if protocol == 'v2' and error.code() == grpc.StatusCode.UNIMPLEMENTED:
+                            consumed = previously_consumed
+                            self._retire_prefetch()
+                            # Settle V2 before V1 can advance the same query session.
+                            async with asyncio.timeout_at(deadline):
+                                await self._settle_retired_results()
+                            self._publish(revision)
+                            self._result_protocol = 'v1'
+                            status = 'unimplemented'
+                            logger.debug('Result batch compatibility fallback.',
+                                         extra={'result_batch_fallback': True})
+                            continue
+                        raise
+                    finally:
+                        wait_seconds = time.monotonic() - started
+                        rpc_seconds = (transport_record.elapsed
+                                       if transport_record is not None and (
+                                           isinstance(transport_record, ContinuousResultStream)
+                                           or transport_record.settled)
+                                       else wait_seconds)
+                        if prefetched:
+                            logger.debug('Result prefetched RPC seconds=%.6f wait_seconds=%.6f status=%s',
+                                         rpc_seconds, wait_seconds, transport_status, extra={
+                                             'result_batch_prefetched_rpc_seconds': rpc_seconds,
+                                             'result_batch_prefetch_wait_seconds': wait_seconds,
+                                             'result_batch_prefetch_status': transport_status})
+                    decode_started = time.monotonic()
+                    try:
+                        if protocol == 'v2':
+                            await self._start_prefetch(response, deadline)
+                            reservation = await reserve_work(deadline=deadline)
+                            await self._decode_v2_response(response, reservation, revision, deadline)
+                        else:
+                            self._retain_result_session(response, revision)
+                            if response.resultBatch:
+                                reservation = await reserve_work(deadline=deadline)
+                                batch = await self._decode_batch(response.resultBatch, reservation, revision)
+                            else:
+                                self._state = 'EXHAUSTED'
+                                batch = None
+                    finally:
+                        decode_seconds = time.monotonic() - decode_started
+                    status = 'ok'
+                    if protocol == 'v1':
+                        return batch
+                    batch = self._result_batches.pop()
+                    if batch is not None:
+                        return batch
+                    if self._result_batches.finished:
+                        self._state = 'EXHAUSTED'
+                        return None
+                finally:
+                    if reservation is not None:
+                        reservation.release()
+                    if dispatched:
+                        _log_result_batch(protocol, rpc_seconds, status,
+                                          response if status == 'ok' else None,
+                                          decode_seconds if status == 'ok' else 0.0)
+                if not isinstance(transport_record, ContinuousResultStream):
+                    await asyncio.sleep(min(backoff, self._connection._remaining(deadline)))
+                    backoff = min(backoff * 2, .1)
+        except asyncio.CancelledError:
+            if consumed:
+                self._fail_result('ambiguous_result')
+            raise
+        except grpc.RpcError:
+            if consumed:
+                self._fail_result('ambiguous_result')
+            raise
+        except IncompleteResultError:
+            raise
+        except Exception as error:
+            if consumed:
+                raise self._fail_result('ambiguous_result') from error
+            raise
 
     async def _fetchmany(self, size, deadline):
         while len(self._rows) < size and self._state != 'EXHAUSTED':
@@ -317,7 +658,7 @@ class AsyncCursor:
                 raise ProgrammingError('Fetch size must be a nonnegative integer.')
             try:
                 return await self._fetchmany(size, self._connection._deadline(timeout))
-            except (asyncio.CancelledError, OperationalError, ProgrammingError, OAuthError):
+            except (asyncio.CancelledError, grpc.RpcError, OperationalError, ProgrammingError, OAuthError):
                 raise
             except Exception as error:
                 raise _operational(error) from error
@@ -335,7 +676,7 @@ class AsyncCursor:
                 if rows is not None:
                     self._rownumber += len(rows)
                 return rows
-            except (asyncio.CancelledError, OperationalError, ProgrammingError, OAuthError):
+            except (asyncio.CancelledError, grpc.RpcError, OperationalError, ProgrammingError, OAuthError):
                 raise
             except Exception as error:
                 raise _operational(error) from error
@@ -353,6 +694,8 @@ class AsyncCursor:
                         return rows
                     rows.extend(batch)
                     self._rownumber += len(batch)
+            except grpc.RpcError:
+                raise
             except asyncio.CancelledError:
                 if rows:
                     self._fail_result('aggregation_failed')
@@ -440,38 +783,50 @@ class AsyncCursor:
             return await self._connection.get_schema_names(self._catalog, timeout=timeout)
 
     async def _stop_active(self):
+        self._cancel_decode()
+        self._retire_prefetch()
         self._revision += 1
         task = self._operation_task or self._active_call
         if task is not None and task is not asyncio.current_task():
-            task.cancel()
-            await asyncio.gather(task, return_exceptions=True)
+                task.cancel()
+                await asyncio.gather(task, return_exceptions=True)
+        await self._settle_retired_results()
+
+    async def _settle_retired_results(self):
+        for pending in tuple(self._retired_results):
+            handle = pending.handle
+            if handle is not None and not pending.settled:
+                await asyncio.gather(asyncio.shield(handle), return_exceptions=True)
+            while not pending.settled:
+                await asyncio.sleep(.005)
+            if not self._connection._oauth_enabled and pending.session_id:
+                self._result_session_id = pending.session_id
+            self._retired_results.remove(pending)
 
     async def _clear(self, deadline, owned=False):
         await self._stop_active()
         if self._route is None:
             if self._state == 'SUBMISSION_UNKNOWN' or self in self._connection._ambiguous_submissions:
                 raise AmbiguousSubmissionError()
+            self._reset_query_state()
             return None
-        if owned:
-            # Owned cleanup has no authority to mint a new legacy session silently.
-            session = '' if self._connection._oauth_enabled else self._connection._session_id
-            if session is None:
-                raise OperationalError('Cleanup has no authenticated legacy session.')
-        else:
-            session = await self._connection.get_session_id(deadline=deadline)
-        request = pb.ClearOrCancelQueryRequest(sessionId=session, queryId=self.query_id, engineIP=self._route.engine_ip)
+        request = await self._query_request(pb.ClearOrCancelQueryRequest, deadline, owned=owned)
         response = await self._connection._rpc('clearOrCancelQuery', request, deadline=deadline,
                                               route=self._route, _cleanup=owned)
         self._connection._release_route(self.query_id)
         self._connection._ambiguous_submissions.discard(self)
+        self._reset_query_state()
+        return response
+
+    def _reset_query_state(self):
         self._route = None
         self._rows.clear()
+        self._reset_result_protocol()
         self._columns = self._description = None
         self._rowcount, self._rownumber = -1, 0
         self._failure = self._cleanup_error = None
         if self._state != 'CLOSED':
             self._state = 'EMPTY'
-        return response
 
     async def clear(self, query_id=None, *, timeout=None):
         self._check()
@@ -496,6 +851,8 @@ class AsyncCursor:
         try:
             async with asyncio.timeout_at(deadline):
                 await self._stop_active()
+                if self._result_batch_v2:
+                    self._fail_result('cancelled_result')
                 request = await self._query_request(pb.CancelQueryRequest, deadline)
                 await self._connection._rpc('cancelQuery', request, deadline=deadline, route=self._route)
                 self._fail_result('cancelled_result')
@@ -505,7 +862,7 @@ class AsyncCursor:
             raise _operational(error) from error
 
     async def _close_owned(self, deadline):
-        if self._state == 'CLOSED':
+        if self._state == 'CLOSED' and self._route is None and not self._retired_results:
             return
         unknown_without_handle = self._state == 'SUBMISSION_UNKNOWN' and self._route is None
         if unknown_without_handle:
@@ -524,7 +881,10 @@ class AsyncCursor:
             self._cleanup_error = OperationalError('Query cleanup is unconfirmed.')
         finally:
             self._rows.clear()
-            self._connection._cursors.discard(self)
+            self._result_batches.clear()
+            if self._route is None and not self._retired_results:
+                self._result_session_id = None
+                self._connection._cursors.discard(self)
 
     async def close(self):
         if self._lease_guard is not None:

e6data_python_connector/async_dialect.py

@@ -216,7 +216,8 @@ class E6dataAsyncDialect(E6dataDialect):
         aliases = {'schema': 'database', 'cluster-name': 'cluster_name',
                    'cluster-uuid': 'cluster_name', 'auto-resume': 'auto_resume'}
         strings = {'catalog', 'database', 'cluster_name', 'ssl_cert'}
-        booleans = {'secure', 'auto_resume', 'debug', 'require_fastbinary'}
+        booleans = {'secure', 'auto_resume', 'debug', 'require_fastbinary',
+                    'enable_result_batch_v2'}
         floats = {'operation_timeout', 'cleanup_timeout', 'oauth_timeout', 'auto_resume_timeout'}
         options = {}
         if url.host is not None: