Add support for RETURNING in SQLAlchemy dialect

This will enable the usage of `FetchedValue()` as server side default
for primary keys in the declarative model.

```python
import random

from sqlalchemy import create_engine, inspect, func, Column, String, DateTime
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.schema import FetchedValue

engine = create_engine(
    "crate://",
    connect_args={
        "username": "crate",
    },
)
Session = sessionmaker(bind=engine)
Base = declarative_base(bind=engine)

class Log(Base):
    __tablename__ = "logs"
    __table_args__ = {"schema": "doc"}
    __mapper_args__ = {"exclude_properties": ["id"]}

    id = Column("_id", String, server_default=FetchedValue(), primary_key=True)
    ts = Column(DateTime, server_default=func.current_timestamp())
    level = Column(String)
    message = Column(String)

if __name__ == "__main__":
    o = Log(level="info", message="Hello World")
    session = Session()
    session.add(o)
    session.commit()

    print(f"{o.id=}, {o.ts=}, {o.level=}, {o.message=}")
```

Author: chaudum

CHANGES.txt

@@ -5,6 +5,12 @@ Changes for crate
 Unreleased
 ==========
 
+- Added support for the ``RETURNING`` clause to the SQLAlchemy dialect. This
+  requires CrateDB 4.2 or greater. In case you use any server side generated
+  columns in your primary key constraint with earlier CrateDB versions, you can
+  turn this feature off by passing ``implicit_returning=False`` in the
+  ``create_engine()`` call.
+
 - Added support for ``geo_point`` and ``geo_json`` types to the SQLAlchemy
   dialect.
 

docs/sqlalchemy.rst

@@ -198,6 +198,45 @@ In this example, we:
     The SQLAlchemy documentation has more information about `working with
     tables`_.
 
+``_id`` as primary key
+......................
+
+As with version 4.2 CrateDB supports the ``RETURNING`` clause, which makes it
+possible to use the ``_id`` column as fetched value for the ``PRIMARY KEY``
+constraint, since the SQLAlchemy ORM always **requires** a primary key.
+
+A table schema like this
+
+.. code-block:: sql
+
+   CREATE TABLE "doc"."logs" (
+     "ts" TIMESTAMP WITH TIME ZONE,
+     "level" TEXT,
+     "message" TEXT
+   )
+
+would translate into the following declarative model::
+
+    >>> from sqlalchemy.schema import FetchedValue
+
+    >>> class Log(Base):
+    ...
+    ...     __tablename__ = 'logs'
+    ...     __mapper_args__ = {
+    ...         'exclude_properties': ['id']
+    ...     }
+    ...
+    ...     id = sa.Column("_id", sa.String, server_default=FetchedValue(), primary_key=True)
+    ...     ts = sa.Column(sa.DateTime, server_default=sa.func.current_timestamp())
+    ...     level = sa.Column(sa.String)
+    ...     message = sa.Column(sa.String)
+
+    >>> log = Log(level="info", message="Hello World")
+    >>> session.add(log)
+    >>> session.commit()
+    >>> log.id
+    ...
+
 .. _using-extension-types:
 
 Extension types
@@ -336,11 +375,9 @@ You can then set the values of the ``Geopoint`` and ``Geoshape`` columns::
     >>> session.add(tokyo)
     >>> session.commit()
 
-
 Querying
 ========
 
-
 When the ``commit`` method is called, two ``INSERT`` statements are sent to
 CrateDB. However, the newly inserted rows aren't immediately available for
 querying because the table index is only updated periodically (one second, by
@@ -544,7 +581,7 @@ The score is made available via the ``_score`` column, which is a virtual
 column, meaning that it doesn't exist on the source table, and in most cases,
 should not be included in your :ref:`table definition <table-definition>`.
 
-You can select ``_score`` as part of a query, like this:
+You can select ``_score`` as part of a query, like this::
 
     >>> session.query(Character.name, '_score') \
     ...     .filter(match(Character.quote_ft, 'space')) \
@@ -557,6 +594,7 @@ table definition But notice that we select the associated score by passing in
 the virtual column name as a string (``_score``) instead of using a defined
 column on the ``Character`` class.
 
+
 .. _SQLAlchemy: http://www.sqlalchemy.org/
 .. _Object-Relational Mapping: https://en.wikipedia.org/wiki/Object-relational_mapping
 .. _dialect: http://docs.sqlalchemy.org/en/latest/dialects/

src/crate/client/sqlalchemy/compiler.py

@@ -177,6 +177,13 @@ def visit_any(self, element, **kw):
             self.process(element.right, **kw)
         )
 
+    def returning_clause(self, stmt, returning_cols):
+        columns = [
+            self._label_select_column(None, c, True, False, {})
+            for c in sa.sql.expression._select_iterables(returning_cols)
+        ]
+        return "RETURNING " + ", ".join(columns)
+
     def visit_insert(self, insert_stmt, asfrom=False, **kw):
         """
         used to compile <sql.expression.Insert> expressions.

src/crate/client/sqlalchemy/dialect.py

@@ -165,6 +165,7 @@ class CrateDialect(default.DefaultDialect):
     type_compiler = CrateTypeCompiler
     supports_native_boolean = True
     colspecs = colspecs
+    implicit_returning = True
 
     def __init__(self, *args, **kwargs):
         super(CrateDialect, self).__init__(*args, **kwargs)