wherobots-python-dbapi 0.4.0__tar.gz

wherobots_python_dbapi-0.4.0/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

wherobots_python_dbapi-0.4.0/PKG-INFO

@@ -0,0 +1,75 @@
+Metadata-Version: 2.1
+Name: wherobots-python-dbapi
+Version: 0.4.0
+Summary: Python DB-API driver for Wherobots DB
+License: Apache 2.0
+Author: Maxime Petazzoni
+Author-email: max@wherobots.com
+Requires-Python: >=3.9,<4.0
+Classifier: License :: Other/Proprietary License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: StrEnum (>=0.4.15,<0.5.0)
+Requires-Dist: cbor2 (>=5.6.3,<6.0.0)
+Requires-Dist: pandas (>=2.2.2,<3.0.0)
+Requires-Dist: pyarrow (>=15.0.2,<16.0.0)
+Requires-Dist: requests (>=2.31.0,<3.0.0)
+Requires-Dist: tenacity (>=8.2.3,<9.0.0)
+Requires-Dist: websockets (>=12.0,<13.0)
+Description-Content-Type: text/markdown
+
+# wherobots-python-dbapi-driver
+
+Python DB-API implementation for Wherobots DB. This package implements a
+PEP-0249 compatible driver to programmatically connect to a Wherobots DB
+runtime and execute Spatial SQL queries.
+
+## Installation
+
+If you use [Poetry](https://python-poetry.org) in your project, add the
+dependency with `poetry add`:
+
+```
+$ poetry add git+https://github.com/wherobots/wherobots-python-dbapi-driver
+```
+
+Otherwise, just `pip install` it:
+
+```
+$ pip install git+https://github.com/wherobots/wherobots-python-dbapi-driver
+```
+
+## Usage
+
+```python
+from wherobots.db import connect
+from wherobots.db.region import Region
+from wherobots.db.runtime import Runtime
+
+with connect(
+        api_key='...',
+        runtime=Runtime.SEDONA,
+        region=Region.AWS_US_WEST_2) as conn:
+    curr = conn.cursor()
+    curr.execute("SHOW SCHEMAS IN wherobots_open_data")
+    results = curr.fetchall()
+    print(results)
+```
+
+The `Cursor` also supports the context manager protocol, so you can use
+it within a `with` statement when needed:
+
+```python
+with connect(...) as conn:
+    with conn.cursor() as curr:
+        curr.execute(...)
+        results = curr.fetchall()
+```
+
+It also implements the `close()` method, as suggested by the PEP-2049
+specification, to support situations where the cursor is wrapped in a
+`contextmanager.closing()`.
+

wherobots_python_dbapi-0.4.0/README.md

@@ -0,0 +1,51 @@
+# wherobots-python-dbapi-driver
+
+Python DB-API implementation for Wherobots DB. This package implements a
+PEP-0249 compatible driver to programmatically connect to a Wherobots DB
+runtime and execute Spatial SQL queries.
+
+## Installation
+
+If you use [Poetry](https://python-poetry.org) in your project, add the
+dependency with `poetry add`:
+
+```
+$ poetry add git+https://github.com/wherobots/wherobots-python-dbapi-driver
+```
+
+Otherwise, just `pip install` it:
+
+```
+$ pip install git+https://github.com/wherobots/wherobots-python-dbapi-driver
+```
+
+## Usage
+
+```python
+from wherobots.db import connect
+from wherobots.db.region import Region
+from wherobots.db.runtime import Runtime
+
+with connect(
+        api_key='...',
+        runtime=Runtime.SEDONA,
+        region=Region.AWS_US_WEST_2) as conn:
+    curr = conn.cursor()
+    curr.execute("SHOW SCHEMAS IN wherobots_open_data")
+    results = curr.fetchall()
+    print(results)
+```
+
+The `Cursor` also supports the context manager protocol, so you can use
+it within a `with` statement when needed:
+
+```python
+with connect(...) as conn:
+    with conn.cursor() as curr:
+        curr.execute(...)
+        results = curr.fetchall()
+```
+
+It also implements the `close()` method, as suggested by the PEP-2049
+specification, to support situations where the cursor is wrapped in a
+`contextmanager.closing()`.

wherobots_python_dbapi-0.4.0/pyproject.toml

@@ -0,0 +1,41 @@
+[tool.poetry]
+name = "wherobots-python-dbapi"
+version = "0.4.0"
+description = "Python DB-API driver for Wherobots DB"
+authors = ["Maxime Petazzoni <max@wherobots.com>"]
+license = "Apache 2.0"
+readme = "README.md"
+packages = [{ include = "wherobots" }]
+
+[project.urls]
+Homepage = "https://github.com/wherobots/wherobots-python-dbapi-driver"
+Tracker = "https://github.com/wherobots/wherobots-python-dbapi-driver/issues"
+
+[tool.poetry.dependencies]
+python = "^3.9"
+requests = "^2.31.0"
+websockets = "^12.0"
+tenacity = "^8.2.3"
+pyarrow = "^15.0.2"
+cbor2 = "^5.6.3"
+pandas = "^2.2.2"
+StrEnum = "^0.4.15"
+
+[tool.poetry.group.dev.dependencies]
+mypy = "^1.8.0"
+pytest = "^8.0.2"
+black = "^24.2.0"
+pre-commit = "^3.6.2"
+conventional-pre-commit = "^3.1.0"
+rich = "^13.7.1"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.pylint.format]
+max-line-length = "88"
+
+[tool.mypy]
+strict = true
+show_error_codes = true

wherobots_python_dbapi-0.4.0/wherobots/db/__init__.py

@@ -0,0 +1,6 @@
+from .connection import Connection
+from .cursor import Cursor
+from .driver import connect, connect_direct
+from .errors import *
+from .region import Region
+from .runtime import Runtime

wherobots_python_dbapi-0.4.0/wherobots/db/connection.py

@@ -0,0 +1,219 @@
+import json
+import logging
+import threading
+import uuid
+from dataclasses import dataclass
+from typing import Callable, Any
+
+import cbor2
+import pyarrow
+import websockets.exceptions
+import websockets.protocol
+import websockets.sync.client
+
+from wherobots.db.constants import (
+    DEFAULT_READ_TIMEOUT_SECONDS,
+    RequestKind,
+    EventKind,
+    ExecutionState,
+    ResultsFormat,
+    DataCompression,
+)
+from wherobots.db.cursor import Cursor
+from wherobots.db.errors import NotSupportedError, OperationalError
+
+_DEFAULT_RESULTS_FORMAT = ResultsFormat.ARROW
+_DEFAULT_DATA_COMPRESSION = DataCompression.BROTLI
+
+
+@dataclass
+class Query:
+    sql: str
+    execution_id: str
+    state: ExecutionState
+    handler: Callable[[Any], None]
+
+
+class Connection:
+    """
+    A PEP-0249 compatible Connection object for Wherobots DB.
+
+    The connection is backed by the WebSocket connected to the Wherobots SQL session instance.
+    Transactions are not supported, so commit() and rollback() raise NotSupportedError.
+
+    This class handles all the interactions with the remote SQL session, and the details of the
+    Wherobots Spatial SQL API protocol. It supports multiple concurrent cursors, each one executing
+    a single query at a time.
+
+    A background thread listens for events from the SQL session, and handles update to the
+    corresponding query state. Queries are tracked by their unique execution ID.
+    """
+
+    def __init__(
+        self,
+        ws: websockets.sync.client.ClientConnection,
+        read_timeout: float = DEFAULT_READ_TIMEOUT_SECONDS,
+    ):
+        self.__ws = ws
+        self.__read_timeout = read_timeout
+        self.__queries: dict[str, Query] = {}
+        self.__thread = threading.Thread(
+            target=self.__main_loop, daemon=True, name="wherobots-connection"
+        )
+        self.__thread.start()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def close(self):
+        self.__ws.close()
+
+    def commit(self):
+        raise NotSupportedError
+
+    def rollback(self):
+        raise NotSupportedError
+
+    def cursor(self) -> Cursor:
+        return Cursor(self.__execute_sql, self.__cancel_query)
+
+    def __main_loop(self):
+        """Main background loop listening for messages from the SQL session."""
+        logging.info("Starting background connection handling loop...")
+        while self.__ws.protocol.state < websockets.protocol.State.CLOSING:
+            try:
+                self.__listen()
+            except TimeoutError:
+                # Expected, retry next time
+                continue
+            except websockets.exceptions.ConnectionClosedOK:
+                logging.info("Connection closed; stopping main loop.")
+                return
+            except Exception as e:
+                logging.exception("Error handling message from SQL session", exc_info=e)
+
+    def __listen(self):
+        """Waits for the next message from the SQL session and processes it.
+
+        The code in this method is purposefully defensive to avoid unexpected situations killing the thread.
+        """
+        message = self.__recv()
+        kind = message.get("kind")
+        execution_id = message.get("execution_id")
+        if not kind or not execution_id:
+            # Invalid event.
+            return
+
+        query = self.__queries.get(execution_id)
+        if not query:
+            logging.warning(
+                "Received %s event for unknown execution ID %s", kind, execution_id
+            )
+            return
+
+        if kind == EventKind.STATE_UPDATED:
+            try:
+                query.state = ExecutionState[message["state"].upper()]
+                logging.info("Query %s is now %s.", execution_id, query.state)
+            except KeyError:
+                logging.warning("Invalid state update message for %s", execution_id)
+                return
+
+            # Incoming state transitions are handled here.
+            if query.state == ExecutionState.SUCCEEDED:
+                self.__request_results(execution_id)
+            elif query.state == ExecutionState.FAILED:
+                query.handler(OperationalError("Query execution failed"))
+
+        elif kind == EventKind.EXECUTION_RESULT:
+            results = message.get("results")
+            if not results or not isinstance(results, dict):
+                logging.warning("Got no results back from %s.", execution_id)
+                return
+
+            result_bytes = results.get("result_bytes")
+            result_format = results.get("format")
+            result_compression = results.get("compression")
+            logging.info(
+                "Received %d bytes of %s-compressed %s results from %s.",
+                len(result_bytes),
+                result_compression,
+                result_format,
+                execution_id,
+            )
+
+            query.state = ExecutionState.COMPLETED
+            if result_format == ResultsFormat.JSON:
+                query.handler(json.loads(result_bytes.decode("utf-8")))
+            elif result_format == ResultsFormat.ARROW:
+                buffer = pyarrow.py_buffer(result_bytes)
+                stream = pyarrow.input_stream(buffer, result_compression)
+                with pyarrow.ipc.open_stream(stream) as reader:
+                    query.handler(reader.read_pandas())
+            else:
+                query.handler(
+                    OperationalError(
+                        f"Unsupported results format {result_format}"
+                    )
+                )
+        elif kind == EventKind.ERROR:
+            query.state = ExecutionState.FAILED
+            error = message.get("message")
+            query.handler(OperationalError(error))
+        else:
+            logging.warning("Received unknown %s event!", kind)
+
+    def __send(self, message: dict[str, Any]) -> None:
+        self.__ws.send(json.dumps(message))
+
+    def __recv(self) -> dict[str, Any]:
+        frame = self.__ws.recv(timeout=self.__read_timeout)
+        if isinstance(frame, str):
+            message = json.loads(frame)
+        elif isinstance(frame, bytes):
+            message = cbor2.loads(frame)
+        else:
+            raise ValueError("Unexpected frame type received")
+        return message
+
+    def __execute_sql(self, sql: str, handler: Callable[[Any], None]) -> str:
+        """Triggers the execution of the given SQL query."""
+        execution_id = str(uuid.uuid4())
+        request = {
+            "kind": RequestKind.EXECUTE_SQL.value,
+            "execution_id": execution_id,
+            "statement": sql,
+        }
+
+        self.__queries[execution_id] = Query(
+            sql=sql,
+            execution_id=execution_id,
+            state=ExecutionState.EXECUTION_REQUESTED,
+            handler=handler,
+        )
+        self.__send(request)
+        return execution_id
+
+    def __request_results(self, execution_id: str) -> None:
+        query = self.__queries.get(execution_id)
+        if not query:
+            return
+
+        request = {
+            "kind": RequestKind.RETRIEVE_RESULTS.value,
+            "execution_id": execution_id,
+            "format": _DEFAULT_RESULTS_FORMAT.value,
+            "compression": _DEFAULT_DATA_COMPRESSION.value,
+        }
+        query.state = ExecutionState.RESULTS_REQUESTED
+        logging.info("Requesting results from %s ...", execution_id)
+        self.__send(request)
+
+    def __cancel_query(self, execution_id: str) -> None:
+        query = self.__queries.pop(execution_id)
+        if query:
+            logging.info("Cancelled query %s.", execution_id)
+            # TODO: when protocol supports it, send cancellation request.

wherobots_python_dbapi-0.4.0/wherobots/db/constants.py

@@ -0,0 +1,60 @@
+from enum import auto
+from strenum import LowercaseStrEnum
+
+from .region import Region
+from .runtime import Runtime
+
+
+DEFAULT_ENDPOINT: str = "api.wherobots.services"  # "api.cloud.wherobots.com"
+STAGING_ENDPOINT: str = "api.staging.wherobots.services"  # "api.staging.wherobots.com"
+DEFAULT_RUNTIME: Runtime = Runtime.SEDONA
+DEFAULT_REGION: Region = Region.AWS_US_WEST_2
+DEFAULT_READ_TIMEOUT_SECONDS: float = 0.25
+DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS: float = 300
+MAX_MESSAGE_SIZE: int = 100 * 2**20  # 100MiB
+
+
+class ExecutionState(LowercaseStrEnum):
+    IDLE = auto()
+    "Not executing any operation."
+
+    EXECUTION_REQUESTED = auto()
+    "Execution of a query has been requested by the driver."
+
+    RUNNING = auto()
+    "The SQL session has reported the query is running."
+
+    SUCCEEDED = auto()
+    "The SQL session has reported the query has completed successfully."
+
+    FAILED = auto()
+    "The SQL session has reported the query has failed."
+
+    RESULTS_REQUESTED = auto()
+    "The driver has requested the query results from the SQL session."
+
+    COMPLETED = auto()
+    "The driver has completed processing the query results."
+
+    def is_terminal_state(self):
+        return self in (ExecutionState.COMPLETED, ExecutionState.FAILED)
+
+
+class RequestKind(LowercaseStrEnum):
+    EXECUTE_SQL = auto()
+    RETRIEVE_RESULTS = auto()
+
+
+class EventKind(LowercaseStrEnum):
+    STATE_UPDATED = auto()
+    EXECUTION_RESULT = auto()
+    ERROR = auto()
+
+
+class ResultsFormat(LowercaseStrEnum):
+    JSON = auto()
+    ARROW = auto()
+
+
+class DataCompression(LowercaseStrEnum):
+    BROTLI = auto()

wherobots_python_dbapi-0.4.0/wherobots/db/cursor.py

@@ -0,0 +1,94 @@
+import queue
+from typing import Any, Optional
+
+from .errors import ProgrammingError, DatabaseError
+
+
+class Cursor:
+
+    def __init__(self, exec_fn, cancel_fn):
+        self.__exec_fn = exec_fn
+        self.__cancel_fn = cancel_fn
+
+        self.__queue: queue.Queue = queue.Queue()
+        self.__results: Optional[list[Any]] = None
+        self.__current_execution_id: Optional[str] = None
+        self.__current_row: int = 0
+
+        # Description and row count are set by the last executed operation.
+        # Their default values are defined by PEP-0249.
+        self.__description: Optional[str] = None
+        self.__rowcount: int = -1
+
+        # Array-size is also defined by PEP-0249 and is expected to be read/writable.
+        self.arraysize: int = 1
+
+    @property
+    def description(self) -> Optional[str]:
+        return self.__description
+
+    @property
+    def rowcount(self) -> int:
+        return self.__rowcount
+
+    def __on_execution_result(self, result) -> None:
+        self.__queue.put(result)
+
+    def __get_results(self) -> Optional[list[Any]]:
+        if not self.__current_execution_id:
+            raise ProgrammingError("No query has been executed yet")
+        if self.__results is not None:
+            return self.__results
+
+        result = self.__queue.get()
+        if isinstance(result, DatabaseError):
+            raise result
+        self.__rowcount = len(result)
+        self.__results = result
+        return self.__results
+
+    def execute(self, operation: str, parameters: dict[str, Any] = None):
+        if self.__current_execution_id:
+            self.__cancel_fn(self.__current_execution_id)
+
+        self.__results = None
+        self.__current_row = 0
+        self.__rowcount = -1
+
+        sql = operation.format(**(parameters or {}))
+        self.__current_execution_id = self.__exec_fn(sql, self.__on_execution_result)
+
+    def executemany(self, operation: str, seq_of_parameters: list[dict[str, Any]]):
+        raise NotImplementedError
+
+    def fetchone(self):
+        results = self.__get_results()[self.__current_row :]
+        if not results:
+            return None
+        self.__current_row += 1
+        return results[0]
+
+    def fetchmany(self, size: int = None):
+        size = size or self.arraysize
+        results = self.__get_results()[self.__current_row : self.__current_row + size]
+        self.__current_row += size
+        return results
+
+    def fetchall(self):
+        return self.__get_results()[self.__current_row :]
+
+    def close(self):
+        """Close the cursor."""
+        pass
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        raise StopIteration
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()

wherobots_python_dbapi-0.4.0/wherobots/db/driver.py

@@ -0,0 +1,159 @@
+"""Wherobots DB driver.
+
+A PEP-0249 compatible driver for interfacing with Wherobots DB.
+"""
+
+import logging
+import urllib.parse
+import queue
+import requests
+import tenacity
+import threading
+import websockets.sync.client
+
+from .constants import (
+    DEFAULT_ENDPOINT,
+    DEFAULT_REGION,
+    DEFAULT_RUNTIME,
+    DEFAULT_READ_TIMEOUT_SECONDS,
+    DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
+    MAX_MESSAGE_SIZE,
+)
+from .errors import (
+    InterfaceError,
+    OperationalError,
+)
+from .region import Region
+from .runtime import Runtime
+from .connection import Connection
+
+apilevel = "2.0"
+threadsafety = 1
+paramstyle = "pyformat"
+
+
+def connect(
+    host: str = DEFAULT_ENDPOINT,
+    token: str = None,
+    api_key: str = None,
+    runtime: Runtime = None,
+    region: Region = None,
+    wait_timeout: float = DEFAULT_SESSION_WAIT_TIMEOUT_SECONDS,
+    read_timeout: float = DEFAULT_READ_TIMEOUT_SECONDS,
+) -> Connection:
+    if not token and not api_key:
+        raise ValueError("At least one of `token` or `api_key` is required")
+    if token and api_key:
+        raise ValueError("`token` and `api_key` can't be both provided")
+
+    headers = {}
+    if token:
+        headers["Authorization"] = f"Bearer {token}"
+    elif api_key:
+        headers["X-API-Key"] = api_key
+
+    runtime = runtime or DEFAULT_RUNTIME
+    region = region or DEFAULT_REGION
+
+    logging.info(
+        "Requesting %s/%s runtime in %s from %s ...",
+        runtime.name,
+        runtime.value,
+        region.value,
+        host,
+    )
+
+    # Default to HTTPS if the hostname doesn't explicitly specify a scheme.
+    if not host.startswith("http:"):
+        host = f"https://{host}"
+
+    try:
+        resp = requests.post(
+            url=f"{host}/sql/session",
+            params={"region": region.value},
+            json={"runtimeId": runtime.value},
+            headers=headers,
+        )
+        resp.raise_for_status()
+    except requests.HTTPError as e:
+        raise InterfaceError("Failed to create SQL session!", e)
+
+    # At this point we've been redirected to /sql/session/{session_id}, which we'll need to keep polling until the
+    # session is in READY state.
+    session_id_url = resp.url
+
+    @tenacity.retry(
+        stop=tenacity.stop_after_delay(wait_timeout),
+        wait=tenacity.wait_exponential(multiplier=1, min=1, max=5),
+        retry=tenacity.retry_if_not_exception_type(
+            (requests.HTTPError, OperationalError)
+        ),
+    )
+    def get_session_uri() -> str:
+        r = requests.get(session_id_url, headers=headers)
+        r.raise_for_status()
+        payload = r.json()
+        status = payload.get("status")
+        logging.info(" ... %s", status)
+        if status in ("REQUESTED", "DEPLOYING", "DEPLOYED", "INITIALIZING"):
+            raise tenacity.TryAgain("SQL Session is not ready yet")
+        elif status == "READY":
+            return payload["appMeta"]["url"]
+        else:
+            logging.error("SQL session creation failed: %s; should not retry.", status)
+            raise OperationalError(f"Failed to create SQL session: {status}")
+
+    try:
+        logging.info("Getting SQL session status from %s ...", session_id_url)
+        session_uri = get_session_uri()
+    except Exception as e:
+        raise InterfaceError("Could not acquire SQL session!", e)
+
+    return connect_direct(
+        uri=http_to_ws(session_uri),
+        headers=headers,
+        read_timeout=read_timeout,
+    )
+
+
+def http_to_ws(uri: str) -> str:
+    """Converts an HTTP URI to a WebSocket URI."""
+    parsed = urllib.parse.urlparse(uri)
+    for from_scheme, to_scheme in [("http", "ws"), ("https", "wss")]:
+        if parsed.scheme == from_scheme:
+            parsed = parsed._replace(scheme=to_scheme)
+    return str(urllib.parse.urlunparse(parsed))
+
+
+def connect_direct(
+    uri: str,
+    headers: dict[str, str] = None,
+    read_timeout: float = DEFAULT_READ_TIMEOUT_SECONDS,
+) -> Connection:
+    q = queue.SimpleQueue()
+
+    def create_ws_connection():
+        try:
+            logging.info("Connecting to SQL session at %s ...", uri)
+            ws = websockets.sync.client.connect(
+                uri=uri,
+                additional_headers=headers,
+                max_size=MAX_MESSAGE_SIZE,
+            )
+            q.put(ws)
+        except Exception as e:
+            q.put(e)
+
+    dt = threading.Thread(
+        name="wherobots-ws-connector",
+        target=create_ws_connection,
+        daemon=True,
+    )
+    dt.start()
+    dt.join()
+
+    result = q.get()
+    if isinstance(result, Exception):
+        raise InterfaceError("Failed to connect to SQL session!") from result
+
+    return Connection(result, read_timeout)

wherobots_python_dbapi-0.4.0/wherobots/db/errors.py

@@ -0,0 +1,26 @@
+class Error(Exception):
+    pass
+
+
+class InterfaceError(Error):
+    pass
+
+
+class DatabaseError(Error):
+    pass
+
+
+class InternalError(DatabaseError):
+    pass
+
+
+class OperationalError(DatabaseError):
+    pass
+
+
+class ProgrammingError(DatabaseError):
+    pass
+
+
+class NotSupportedError(DatabaseError):
+    pass

wherobots_python_dbapi-0.4.0/wherobots/db/region.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+
+class Region(Enum):
+    AWS_US_WEST_2 = "aws-us-west-2"

wherobots_python_dbapi-0.4.0/wherobots/db/runtime.py

@@ -0,0 +1,21 @@
+from enum import Enum
+
+
+class Runtime(Enum):
+    SEDONA = "TINY"
+    SAN_FRANCISCO = "SMALL"
+    NEW_YORK = "MEDIUM"
+    CAIRO = "LARGE"
+    DELHI = "XLARGE"
+    TOKYO = "XXLARGE"
+    ATLANTIS = "4x-large"
+
+    NEW_YORK_HIMEM = "medium-himem"
+    CAIRO_HIMEM = "large-himem"
+    DELHI_HIMEM = "x-large-himem"
+    TOKYO_HIMEM = "2x-large-himem"
+    ATLANTIS_HIMEM = "4x-large-himem"
+
+    SEDONA_GPU = "tiny-a10-gpu"
+    SAN_FRANCISCO_GPU = "small-a10-gpu"
+    NEW_YORK_GPU = "medium-a10-gpu"