Updates to clean up baseline project

Changes
=======
* Static type checking in place (via mypy)
* Baseline unit tests available and passing
* Baseline integration query tests available and passing

Author: thejcfactor

MANIFEST.in

@@ -0,0 +1,8 @@
+include *.txt LICENSE CONTRIBUTING.md pyproject.toml couchbase_analytics_version.py
+include couchbase-sdk-analytics-python-black-duck-manifest.yaml
+include couchbase_analytics/common/core/_nonprod_certificates/*.pem
+include couchbase_analytics/common/core/_capella_certificates/*.pem
+recursive-include couchbase_analytics *.py
+recursive-include acouchbase_analytics *.py
+global-exclude *.py[cod] *.DS_Store
+exclude .git .gitignore .gitmodules gocaves* *.jar .clang* .cmake* .pre* .flake* MANIFEST.in

acouchbase_analytics/options.py

@@ -13,7 +13,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from couchbase_analytics.common.enums import IpProtocol as IpProtocol  # noqa: F401
 from couchbase_analytics.common.options import ClusterOptions as ClusterOptions  # noqa: F401
 from couchbase_analytics.common.options import ClusterOptionsKwargs as ClusterOptionsKwargs  # noqa: F401
 from couchbase_analytics.common.options import QueryOptions as QueryOptions  # noqa: F401

acouchbase_analytics/protocol/cluster.pyi

@@ -21,7 +21,7 @@ if sys.version_info < (3, 11):
 else:
     from typing import Unpack
 
-from acouchbase_analytics.protocol.core.client_adapter import _ClientAdapter
+from acouchbase_analytics.protocol.core.client_adapter import _AsyncClientAdapter
 from acouchbase_analytics.protocol.database import AsyncDatabase
 from couchbase_analytics.common.credential import Credential
 from couchbase_analytics.common.result import AsyncQueryResult
@@ -54,7 +54,7 @@ class AsyncCluster:
                  **kwargs: Unpack[ClusterOptionsKwargs]) -> None: ...
 
     @property
-    def client_adapter(self) -> _ClientAdapter: ...
+    def client_adapter(self) -> _AsyncClientAdapter: ...
 
     @property
     def connected(self) -> bool: ...

acouchbase_analytics/protocol/core/_request_context.py

@@ -5,6 +5,7 @@
 from typing import (Any,
                     Awaitable,
                     Callable,
+                    Coroutine,
                     Dict,
                     List,
                     Optional,
@@ -20,7 +21,7 @@
                                                              get_time)
 from couchbase_analytics.common.core.net_utils import get_request_ip_async
 from couchbase_analytics.common.deserializer import Deserializer
-from couchbase_analytics.common.errors import AnalyticsError
+from couchbase_analytics.common.errors import AnalyticsError, InvalidCredentialError
 from couchbase_analytics.common.streaming import StreamingState
 from couchbase_analytics.protocol.connection import DEFAULT_TIMEOUTS
 from couchbase_analytics.protocol.errors import ErrorMapper
@@ -41,7 +42,7 @@ def __init__(self,
         self._client_adapter = client_adapter
         self._request = request
         self._backend = backend or current_async_library()
-        self._response_task: Optional[Task] = None
+        # self._response_task: Optional[Task] = None
         self._request_state = StreamingState.NotStarted
         self._stage_completed: Optional[anyio.Event] = None
         self._request_error: Optional[Exception] = None
@@ -80,9 +81,9 @@ def request_state(self, state: StreamingState) -> None:
             raise TypeError('request_state must be an instance of StreamingState')
         self._request_state = state
 
-    @property
-    def stage_completed(self) -> anyio.Event:
-        return self._stage_completed
+    # @property
+    # def stage_completed(self) -> Optional[anyio.Event]:
+    #     return self._stage_completed
 
     @property
     def timed_out(self) -> bool:
@@ -94,9 +95,10 @@ def cancelled(self) -> bool:
 
     async def _execute(self, fn: Callable[..., Awaitable[Any]], *args: object) -> None:
         await fn(*args)
-        self._stage_completed.set()
+        if self._stage_completed is not None:
+            self._stage_completed.set()
 
-    async def _trace_handler(self, event_name, _) -> None:
+    async def _trace_handler(self, event_name: str, _: str) -> None:
         if event_name == 'connection.connect_tcp.complete':
             # after connection is established, we need to update the cancel_scope deadline to match the query_timeout
             self._update_cancel_scope_deadline(self._request_deadline, is_absolute=True)
@@ -115,24 +117,31 @@ async def initialize(self) -> None:
         self._request_state = StreamingState.Started
         # we set the request timeout once the context is initialized in order to create the deadline 
         # closer to when the upstream logic will begin to use the request context
-        timeouts = self._request.get_request_timeouts()
-        self._request_deadline = get_time() + timeouts.get('read', DEFAULT_TIMEOUTS['query_timeout'])
+        timeouts = self._request.get_request_timeouts() or {}
+        self._request_deadline = get_time() + (timeouts.get('read', None) or DEFAULT_TIMEOUTS['query_timeout'])
         self._update_cancel_scope_deadline(self._connect_timeout)
 
     async def send_request(self, enable_trace_handling: Optional[bool]=False) -> HttpCoreResponse:
         ip = await get_request_ip_async(self._request.host, self._request.port, self._request.previous_ips)
         if ip is None:
             attempted_ips = ', '.join(self._request.previous_ips or [])
-            raise AnalyticsError(f'Connect failure.  Attempted to connect to resolved IPs: {attempted_ips}.')
+            raise AnalyticsError(message=f'Connect failure.  Attempted to connect to resolved IPs: {attempted_ips}.')
 
         if enable_trace_handling is True:
             (self._request.update_url(ip, self._client_adapter.analytics_path)
-                          .update_extensions({'trace': self._trace_handler})
+                          .add_trace_to_extensions(self._trace_handler)
                           .update_previous_ips(ip))
         else:
             self._request.update_url(ip, self._client_adapter.analytics_path).update_previous_ips(ip)
         response = await self._client_adapter.send_request(self._request)
         self._request.set_client_server_addrs(response)
+        if response.status_code == 401:
+            context = {
+                'client_addr': self._request.client_addr,
+                'server_addr': self._request.server_addr,
+                'http_status': response.status_code,
+            }
+            raise InvalidCredentialError(str(context))
         return response
 
     async def shutdown(self,
@@ -149,14 +158,16 @@ async def shutdown(self,
         if StreamingState.is_okay(self._request_state):
             self._request_state = StreamingState.Completed
 
-    def create_response_task(self, fn: Callable[..., Awaitable[Any]], *args: object) -> Task:
+    def create_response_task(self, fn: Callable[..., Coroutine[Any, Any, Any]], *args: object) -> Task[Any]:
         if self._backend is None or self._backend.backend_lib != 'asyncio':
             raise RuntimeError('Must use the asyncio backend to create a response task.')
+        if self._backend.loop is None:
+            raise RuntimeError('Async backend loop is not initialized.')
         task_name = f'{self._id}-response-task'
         print(f'Creating response task: {task_name}')
-        task = self._backend.loop.create_task(fn(*args), name=task_name)
+        task: Task[Any] = self._backend.loop.create_task(fn(*args), name=task_name)
         # TODO: I don't think this callback is necessary...need to add more tests to confirm
-        def task_done(t: Task) -> None:
+        def task_done(t: Task[Any]) -> None:
             print(f'Task ({t.get_name()}) done: {t.done()}, cancelled: {t.cancelled()}')
 
         task.add_done_callback(task_done)
@@ -170,15 +181,23 @@ def start_next_stage(self,
                          fn: Callable[..., Awaitable[Any]],
                          *args: object,
                          reset_previous_stage: Optional[bool]=False) -> None:
-        if reset_previous_stage is True:
-            if self._stage_completed is not None:
+        # if reset_previous_stage is True:
+        #     if self._stage_completed is not None:
+        #         self._stage_completed = None
+        if self._stage_completed is not None:
+            if reset_previous_stage is True:
                 self._stage_completed = None
-        elif self._stage_completed is not None:
-            raise RuntimeError('Task already running in this context.')
+            else:
+                raise RuntimeError('Task already running in this context.')
 
         self._stage_completed = anyio.Event()
         self._taskgroup.start_soon(self._execute, fn, *args)
 
+    async def wait_for_stage_to_complete(self) -> None:
+        if self._stage_completed is None:
+            return
+        await self._stage_completed.wait()
+
     async def process_error(self, json_data: List[Dict[str, Any]]) -> None:
         self._request_state = StreamingState.Error
         if not isinstance(json_data, list):
@@ -209,4 +228,6 @@ async def __aexit__(self,
                 self._request_state = StreamingState.Cancelled
             elif exc_val is not None:
                 self._request_state = StreamingState.Error
-            del self._taskgroup
+            del self._taskgroup
+            # TODO:  should we suppress here (e.g., return True)
+            return None

acouchbase_analytics/protocol/core/client_adapter.py

@@ -115,6 +115,8 @@ async def create_client(self) -> None:
         """
         if not hasattr(self, '_client'):
             if self._conn_details.is_secure():
+                if self._conn_details.ssl_context is None:
+                    raise ValueError('SSL context is required for secure connections.')
                 transport = None
                 if self._http_transport_cls is not None:
                     transport = self._http_transport_cls(verify=self._conn_details.ssl_context)
@@ -137,6 +139,9 @@ async def send_request(self, request: QueryRequest) -> Response:
         if not hasattr(self, '_client'):
             raise RuntimeError('Client not created yet')
 
+        if request.url is None:
+            raise ValueError('Request URL cannot be None')
+
         req = self._client.build_request(request.method,
                                          request.url,
                                          json=request.body,

acouchbase_analytics/protocol/scope.pyi

@@ -21,7 +21,7 @@ if sys.version_info < (3, 11):
 else:
     from typing import Unpack
 
-from acouchbase_analytics.protocol.core.client_adapter import _ClientAdapter
+from acouchbase_analytics.protocol.core.client_adapter import _AsyncClientAdapter
 from acouchbase_analytics.protocol.database import AsyncDatabase as AsyncDatabase
 from couchbase_analytics.options import QueryOptions, QueryOptionsKwargs
 from couchbase_analytics.result import AsyncQueryResult
@@ -30,7 +30,7 @@ class AsyncScope:
     def __init__(self, database: AsyncDatabase, scope_name: str) -> None: ...
 
     @property
-    def client_adapter(self) -> _ClientAdapter: ...
+    def client_adapter(self) -> _AsyncClientAdapter: ...
 
     @property
     def name(self) -> str: ...

acouchbase_analytics/protocol/streaming.py

@@ -16,13 +16,20 @@
 from __future__ import annotations
 
 import json
+import sys
+
 from asyncio import CancelledError
 from functools import wraps
 from typing import (Any,
                     Callable,
                     Coroutine,
                     Optional)
 
+if sys.version_info < (3, 10):
+    from typing_extensions import TypeAlias
+else:
+    from typing import TypeAlias
+
 from httpx import Response as HttpCoreResponse
 
 # TODO: errors?
@@ -39,20 +46,20 @@
 from couchbase_analytics.common.streaming import StreamingState
 
 
+
+
 class RequestWrapper:
     """
         **INTERNAL**
     """
 
     @classmethod
-    def handle_retries(cls,  # noqa: C901
-                       ) -> Callable[[Callable[[], None]], Callable[[AsyncHttpStreamingResponse], Coroutine[Any, Any, None]]]:
+    def handle_retries(cls) -> Callable[[SendRequestFunc], WrappedSendRequestFunc]:  # noqa: C901
         """
             **INTERNAL**
         """
 
-        def decorator(fn: Callable[[], None]  # noqa: C901
-                      ) -> Callable[[AsyncHttpStreamingResponse], Coroutine[Any, Any, None]]:
+        def decorator(fn: SendRequestFunc) -> WrappedSendRequestFunc:  # noqa: C901
             @wraps(fn)
             async def wrapped_fn(self: AsyncHttpStreamingResponse) -> None:
                 try:
@@ -93,11 +100,11 @@ def __init__(self,
 
     async def _finish_processing_stream(self) -> None:
         if not self._request_context.has_stage_completed:
-            await self._request_context.stage_completed.wait()
+            await self._request_context.wait_for_stage_to_complete()
 
         while not self._json_stream.token_stream_exhausted:
             self._request_context.start_next_stage(self._json_stream.continue_parsing, reset_previous_stage=True)
-            await self._request_context.stage_completed.wait()
+            await self._request_context.wait_for_stage_to_complete()
 
     def _maybe_continue_to_process_stream(self) -> None:
         if not self._request_context.has_stage_completed:
@@ -112,10 +119,11 @@ async def _process_response(self, raw_response: Optional[ParsedResult]=None) ->
         if raw_response is None:
             raw_response = await self._json_stream.get_result()
             if raw_response is None:
-                # TODO: logging??
-                # TODO: exception??
-                raise RuntimeError('No result from JsonStream')
-        
+                raise AnalyticsError(message='Received unexpected empty result from JsonStream.')
+                
+        if raw_response.value is None:
+            raise AnalyticsError(message='Received unexpected empty result from JsonStream.')
+
         json_response = json.loads(raw_response.value)
         if 'errors' in json_response:
             await self._request_context.process_error(json_response['errors'])
@@ -173,6 +181,8 @@ async def get_next_row(self) -> Any:
         self._maybe_continue_to_process_stream()
         raw_response = await self._json_stream.get_result()
         if raw_response.result_type == ParsedResultType.ROW:
+            if raw_response.value is None:
+                raise AnalyticsError(message='Unexpected empty row response while streaming.')
             return self._request_context.deserializer.deserialize(raw_response.value)
         elif raw_response.result_type in [ParsedResultType.ERROR, ParsedResultType.UNKNOWN]:
             await self._process_response(raw_response=raw_response)
@@ -200,3 +210,8 @@ async def send_request(self) -> None:
         else:
             await self._finish_processing_stream()
             await self._process_response()
+
+SendRequestFunc: TypeAlias = Callable[[AsyncHttpStreamingResponse], Coroutine[Any, Any, None]]
+# Although, SendRequestFunc is the same type as WrappedSendRequestFunc, keep separate for clarity and indicate
+# WrappedSendRequestFunc is a decorator
+WrappedSendRequestFunc: TypeAlias = Callable[[AsyncHttpStreamingResponse], Coroutine[Any, Any, None]]