CockroachDB Support (#608)

* Cockroach config for tests. * Added feedback from #607 * New CockroachDB engine. * No more extension=[] ! * Running tests as postgres. * Initial build for CockroachDB * Running tests. * Run workflow. * Naming consistency. Making roughly equivalent to Postgres. * Timestamp bug. * DDL for Cockroach. for_engines decorator. * Progress. * Added database prep for cluster settings. * Tests passing. * Passing migration type checks. * Cockroach ALTER COLUMN TYPE moved outside of transaction until #49351 is solved. * New test helpers. * Fixtures. * Cockroach specific stuff. * Array tests. * Column tests. * JSON tests. * M2M tables. as_list is effected by cockroachdb/cockroach#71908 * Skip test for numeric. * Pool test. * Transaction tests. * Related readable tests. * Save tests. * To dict tests. * JOIN tests. * Output tests. * Raw tests. * Repr tests. * Select tests. * "As string" tests. * Cockroach having trouble with one specific operator. * Skipping until CRDB issue is resolved. Should work though! * Minor updates. * Test dump load. Other CRDB things. * Migration manager tests. * ALTER tests. * SQL Generator tests. * Serialization tests. * Test helpers. * Cockroach Specific stuff. * Numerics all the same for CRDB. * Cockroach specific stuff. * Cleanup. * Skipping timestamp tests. * DDL resolution stuff for Cockroach. * WIP. * Cohesion and special test tables for Cockroach. * WIP * Special table for Cockroach. * Removed hack. * Should pass now. * Removed some debug stuff. * LGTM stuff. * Feedback. * Feedback. * Feedback. * More readable "postgres" or "cockroach. * Added Cockroach overrides for COLUMN_TYPE_MAP, COLUMN_DEFAULT_PARSER * Restore SmallInt functionality for Cockroach. * Restored type() comparison for reflection / generators. * Updated to latest Cockroach release. * Cleanup. * Cleanup. * Build stuff. * Cockroach Serial now uses unique_rowid() as default value always instead of DEFAULT. * Removed first_id() test helper. * Time Travel queries using AS OF SYSTEM TIME * Added documentation. * Small docs fix. * Test update. * Refactored out unnecessary code because we use unique_rowid() as default for Cockroach Serial and BigSerial. * BIGINT instead of BIGSERIAL for Cockroach. * LGTM stuff * LGTM... * Made compatible with older Python versions. * Get Cockroach passing without Python 3.9 stuff. * Removed unused variable. * Fixed test for SQLite. * Re-add Python 3.7 for CRDB + new CRDB version. * Consolidated cockroach_skip() into engines_skip() everywhere. * Re-add SQLite DDL. * Moved back to original engine DDL selector. * Reverted mistake in test suite. * Remove migration that should not be available. * Moving back to postgres_only for specific test. Probably cannot detect engine type at this stage. * postgres_only() to engine_only() * Set sane default for JSONB because we override column_type in JSON for some engines. * Ran isort. * Black being obnoxious, part 1. * Black being obnoxious Part 2. * Flake8 * Black * Flake8 * Flake8 * Flake8 * isort * Flake8 * Flake8 * Added alternate test names for duplicates to remove F811 * Added alternative name. * mypy * mypy * mypy * mypy * mypy * mypy * mypy * mypy * mypy * Cockroach: Now testing latest stable and upcoming release. * Cockroach: Testing only v22.2 * Cockroach version documentation. * Engine detection now based on type because CockroachEngine inherits PostgresEngine. * postgres_only is now engines_only("postgres", "cockroach") because CockroachEngine was always tested as a child class. * Updated to Cockroach 22.2 beta. * isort. * assertAlmostEqual for Decimal test. Docs update. * Linter * Cockroach 22.2 Beta 2. * Added notes about special Column types for Cockroach * version pin xpresso, so integration tests pass * add small section about running cockroach in contributing docs * fix typo Co-authored-by: Daniel Townsend <dan@dantownsend.co.uk>
piccolo-orm · Oct 11, 2022 · 721dc2b · 721dc2b
1 parent 45dbb8b
commit 721dc2b
Show file tree

Hide file tree

Showing 78 changed files with 2,179 additions and 445 deletions.
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
@@ -2,7 +2,7 @@ name: Test Suite
 
 on:
     push:
-        branches: ["master"]
+        branches: ["master", "cockroachdb"]
     pull_request:
         branches: ["master"]
 
@@ -132,6 +132,40 @@ jobs:
               uses: codecov/codecov-action@v1
               if: matrix.python-version == '3.7'
 
+    cockroach:
+        runs-on: ubuntu-latest
+        timeout-minutes: 30
+        strategy:
+            matrix:
+                python-version: ["3.7", "3.8", "3.9", "3.10"]
+                cockroachdb-version: ["v22.2.0-beta.2"]
+        steps:
+            - uses: actions/checkout@v3
+            - name: Set up Python ${{ matrix.python-version }}
+              uses: actions/setup-python@v3
+              with:
+                  python-version: ${{ matrix.python-version }}
+            - name: Install dependencies
+              run: |
+                  python -m pip install --upgrade pip
+                  pip install -r requirements/requirements.txt
+                  pip install -r requirements/test-requirements.txt
+                  pip install -r requirements/extras/postgres.txt
+            - name: Setup CockroachDB
+              run: |
+                  wget -qO- https://binaries.cockroachdb.com/cockroach-${{ matrix.cockroachdb-version }}.linux-amd64.tgz | tar xz
+                  ./cockroach-${{ matrix.cockroachdb-version }}.linux-amd64/cockroach start-single-node --insecure --background
+                  ./cockroach-${{ matrix.cockroachdb-version }}.linux-amd64/cockroach sql --insecure -e 'create database piccolo;'
+
+            - name: Test with pytest, CockroachDB
+              run: ./scripts/test-cockroach.sh
+              env:
+                  PG_HOST: localhost
+                  PG_DATABASE: piccolo
+            - name: Upload coverage
+              uses: codecov/codecov-action@v1
+              if: matrix.python-version == '3.7'
+
     sqlite:
         runs-on: ubuntu-latest
         timeout-minutes: 30

diff --git a/docs/src/piccolo/contributing/index.rst b/docs/src/piccolo/contributing/index.rst
@@ -8,6 +8,25 @@ instructions.
 
 -------------------------------------------------------------------------------
 
+Running Cockroach
+-----------------
+
+To get a local Cockroach instance running, you can use:
+
+.. code-block:: console
+
+    cockroach start-single-node --insecure --store=type=mem,size=2GiB
+
+Make sure the test database exists:
+
+.. code-block:: console
+
+    cockroach sql --insecure
+    >>> create database piccolo
+    >>> use piccolo
+
+-------------------------------------------------------------------------------
+
 Get the tests running
 ---------------------
 
@@ -20,6 +39,7 @@ Get the tests running
 * Setup Postgres, and make sure a database called ``piccolo`` exists (see ``tests/postgres_conf.py``).
 * Run the automated code linting/formatting tools: ``./scripts/lint.sh``
 * Run the test suite with Postgres: ``./scripts/test-postgres.sh``
+* Run the test suite with Cockroach: ``./scripts/test-cockroach.sh``
 * Run the test suite with Sqlite: ``./scripts/test-sqlite.sh``
 
 -------------------------------------------------------------------------------

diff --git a/docs/src/piccolo/engines/cockroach_engine.rst b/docs/src/piccolo/engines/cockroach_engine.rst
@@ -0,0 +1,87 @@
+CockroachEngine
+===============
+
+Configuration
+-------------
+
+.. code-block:: python
+
+    # piccolo_conf.py
+    from piccolo.engine.cockroach import CockroachEngine
+
+
+    DB = CockroachEngine(config={
+        'host': 'localhost',
+        'database': 'piccolo',
+        'user': 'root',
+        'password': '',
+        'port': '26257',
+    })
+
+config
+~~~~~~
+
+The config dictionary is passed directly to the underlying database adapter,
+asyncpg. See the `asyncpg docs <https://magicstack.github.io/asyncpg/current/api/index.html#connection>`_
+to learn more.
+
+-------------------------------------------------------------------------------
+
+Connection pool
+---------------
+
+To use a connection pool, you need to first initialise it. The best place to do
+this is in the startup event handler of whichever web framework you are using.
+
+Here's an example using Starlette. Notice that we also close the connection
+pool in the shutdown event handler.
+
+.. code-block:: python
+
+    from piccolo.engine import engine_finder
+    from starlette.applications import Starlette
+
+
+    app = Starlette()
+
+
+    @app.on_event('startup')
+    async def open_database_connection_pool():
+        engine = engine_finder()
+        await engine.start_connection_pool()
+
+
+    @app.on_event('shutdown')
+    async def close_database_connection_pool():
+        engine = engine_finder()
+        await engine.close_connection_pool()
+
+.. hint:: Using a connection pool helps with performance, since connections
+    are reused instead of being created for each query.
+
+Once a connection pool has been started, the engine will use it for making
+queries.
+
+.. hint:: If you're running several instances of an app on the same server,
+    you may prefer an external connection pooler - like pgbouncer.
+
+Configuration
+~~~~~~~~~~~~~
+
+The connection pool uses the same configuration as your engine. You can also
+pass in additional parameters, which are passed to the underlying database
+adapter. Here's an example:
+
+.. code-block:: python
+
+    # To increase the number of connections available:
+    await engine.start_connection_pool(max_size=20)
+
+-------------------------------------------------------------------------------
+
+Source
+------
+
+.. currentmodule:: piccolo.engine.cockroach
+
+.. autoclass:: CockroachEngine
diff --git a/docs/src/piccolo/engines/index.rst b/docs/src/piccolo/engines/index.rst
@@ -126,3 +126,4 @@ Engine types
 
     ./sqlite_engine
     ./postgres_engine
+    ./cockroach_engine
diff --git a/docs/src/piccolo/features/supported_databases.rst b/docs/src/piccolo/features/supported_databases.rst
@@ -8,6 +8,12 @@ will be using in production.
 
 -------------------------------------------------------------------------------
 
+Cockroach DB
+------------
+Cockroach support is in experimental beta.
+
+-------------------------------------------------------------------------------
+
 SQLite
 ------
 SQLite support is not as complete as Postgres, but it is available - mostly

diff --git a/docs/src/piccolo/getting_started/database_support.rst b/docs/src/piccolo/getting_started/database_support.rst
@@ -6,6 +6,8 @@ Database Support
 `Postgres <https://www.postgresql.org/>`_ is the primary database which Piccolo
 was designed for.
 
+`CockroachDB <https://www.cockroachlabs.com/>`_ is in experimental beta.
+
 Limited `SQLite <https://www.sqlite.org/index.html>`_ support is available,
 mostly to enable tooling like the :ref:`playground <Playground>`. Postgres is the only database we
 recommend for use in production with Piccolo.
diff --git a/docs/src/piccolo/getting_started/index.rst b/docs/src/piccolo/getting_started/index.rst
@@ -10,6 +10,7 @@ Getting Started
     ./installing_piccolo
     ./playground
     ./setup_postgres
+    ./setup_cockroach
     ./setup_sqlite
     ./example_schema
     ./sync_and_async
diff --git a/docs/src/piccolo/getting_started/installing_piccolo.rst b/docs/src/piccolo/getting_started/installing_piccolo.rst
@@ -23,7 +23,7 @@ Now install Piccolo, ideally inside a `virtualenv <https://docs.python-guide.org
     # The important bit:
     pip install piccolo
 
-    # Install Piccolo with PostgreSQL driver:
+    # Install Piccolo with PostgreSQL or CockroachDB driver:
     pip install 'piccolo[postgres]'
 
     # Install Piccolo with SQLite driver:

diff --git a/docs/src/piccolo/getting_started/setup_cockroach.rst b/docs/src/piccolo/getting_started/setup_cockroach.rst
@@ -0,0 +1,68 @@
+.. _setting_up_cockroach:
+
+###############
+Setup Cockroach
+###############
+
+Installation
+************
+
+Follow the  `instructions for your OS <https://www.cockroachlabs.com/docs/stable/install-cockroachdb.html>`_.
+
+
+Versions
+--------
+
+We support the latest stable version.
+
+.. note::
+   Features using ``format()`` will be available in v22.2 or higher, but we recommend using the stable version so you can upgrade automatically when it becomes generally available.
+
+   Cockroach is designed to be a "rolling database": Upgrades are as simple as switching out to the next version of a binary (or changing a number in a ``docker-compose.yml``). This has one caveat: You cannot upgrade an "alpha" release. It is best to stay on the latest stable.
+
+
+-------------------------------------------------------------------------------
+
+Creating a database
+*******************
+
+cockroach sql
+-------------
+
+CockroachDB comes with its own management tooling.
+
+.. code-block:: bash
+
+    cd ~/wherever/you/installed/cockroachdb
+    cockroach sql --insecure
+
+Enter the following:
+
+.. code-block:: bash
+
+    create database piccolo;
+    use piccolo;
+
+Management GUI
+--------------
+
+CockroachDB comes with its own web-based management GUI available on localhost: http://127.0.0.1:8080/
+
+Beekeeper Studio
+----------------
+
+If you prefer a GUI, Beekeeper Studio is recommended and has an  `installer available <https://www.beekeeperstudio.io/>`_.
+
+
+-------------------------------------------------------------------------------
+
+
+Column Types
+************
+
+As of this writing, CockroachDB will always convert ``JSON`` to ``JSONB`` and will always report ``INTEGER`` as ``BIGINT``.
+
+Piccolo will automatically handle these special cases for you, but we recommend being explicit about this to prevent complications in future versions of Piccolo.
+
+* Use ``JSONB()`` instead of ``JSON()``
+* Use ``BigInt()`` instead of ``Integer()``
diff --git a/piccolo/apps/asgi/commands/new.py b/piccolo/apps/asgi/commands/new.py
@@ -11,7 +11,7 @@
 TEMPLATE_DIR = os.path.join(os.path.dirname(__file__), "templates/app/")
 SERVERS = ["uvicorn", "Hypercorn"]
 ROUTERS = ["starlette", "fastapi", "blacksheep", "xpresso", "starlite"]
-ROUTER_VERSIONS = {"starlite": "==1.23.0"}
+ROUTER_VERSIONS = {"starlite": "==1.23.0", "xpresso": "==0.43.0"}
 
 
 def print_instruction(message: str):

diff --git a/piccolo/apps/migrations/auto/migration_manager.py b/piccolo/apps/migrations/auto/migration_manager.py
@@ -788,4 +788,11 @@ async def run(self, backwards=False):
             await self._run_drop_columns(backwards=backwards)
             await self._run_drop_tables(backwards=backwards)
             await self._run_rename_columns(backwards=backwards)
+            # We can remove this for cockroach when resolved.
+            # https://github.com/cockroachdb/cockroach/issues/49351
+            # "ALTER COLUMN TYPE is not supported inside a transaction"
+            if engine.engine_type != "cockroach":
+                await self._run_alter_columns(backwards=backwards)
+
+        if engine.engine_type == "cockroach":
             await self._run_alter_columns(backwards=backwards)