From 6757f7fd26ba64b3e0d2219408ead1f19b30d75c Mon Sep 17 00:00:00 2001
From: Niklas Schmidtmer <niklas@crate.io>
Date: Fri, 2 Dec 2022 09:14:34 +0100
Subject: [PATCH] Support computed columns in SQLAlchemy ORM

---
 CHANGES.txt                                   |  1 +
 docs/sqlalchemy.rst                           |  5 ++++-
 src/crate/client/sqlalchemy/compiler.py       | 15 +++++++++++++
 .../sqlalchemy/tests/create_table_test.py     | 22 +++++++++++++++++++
 4 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/CHANGES.txt b/CHANGES.txt
index 6d080add..32927084 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -11,6 +11,7 @@ Unreleased
   that SQL statement clauses like ``LIMIT -1`` could have been generated. Both
   PostgreSQL and CrateDB only accept ``LIMIT ALL`` instead.
 
+- Added support for computed columns in the SQLAlchemy ORM
 
 2022/10/10 0.27.2
 =================
diff --git a/docs/sqlalchemy.rst b/docs/sqlalchemy.rst
index f53ea5ed..72a73d78 100644
--- a/docs/sqlalchemy.rst
+++ b/docs/sqlalchemy.rst
@@ -179,9 +179,10 @@ system`_::
     ...
     ...     id = sa.Column(sa.String, primary_key=True, default=gen_key)
     ...     name = sa.Column(sa.String)
+    ...     name_normalized = sa.Column(sa.String, sa.Computed("lower(name)"))
     ...     quote = sa.Column(sa.String)
     ...     details = sa.Column(types.Object)
-    ...     more_details = sa.Column(ObjectArray)
+    ...     more_details = sa.Column(types.ObjectArray)
     ...     name_ft = sa.Column(sa.String)
     ...     quote_ft = sa.Column(sa.String)
     ...
@@ -197,6 +198,8 @@ In this example, we:
 - Use the ``gen_key`` function to provide a default value for the ``id`` column
   (which is also the primary key)
 - Use standard SQLAlchemy types for the ``id``, ``name``, and ``quote`` columns
+- Define a computed column ``name_normalized`` (based on ``name``) that
+  translates into a generated column
 - Use the `Object`_ extension type for the ``details`` column
 - Use the `ObjectArray`_ extension type for the ``more_details`` column
 - Set up the ``name_ft`` and ``quote_ft`` fulltext indexes, but exclude them from
diff --git a/src/crate/client/sqlalchemy/compiler.py b/src/crate/client/sqlalchemy/compiler.py
index a747bbe5..b55e5343 100644
--- a/src/crate/client/sqlalchemy/compiler.py
+++ b/src/crate/client/sqlalchemy/compiler.py
@@ -108,8 +108,23 @@ def get_column_specification(self, column, **kwargs):
         colspec = self.preparer.format_column(column) + " " + \
             self.dialect.type_compiler.process(column.type)
         # TODO: once supported add default / NOT NULL here
+
+        if column.computed is not None:
+            colspec += " " + self.process(column.computed)
+
         return colspec
 
+    def visit_computed_column(self, generated):
+        if generated.persisted is False:
+            raise sa.exc.CompileError(
+                "Virtual computed columns are not supported, set "
+                "'persisted' to None or True"
+            )
+
+        return "GENERATED ALWAYS AS (%s)" % self.sql_compiler.process(
+            generated.sqltext, include_table=False, literal_binds=True
+        )
+
     def post_create_table(self, table):
         special_options = ''
         clustered_options = defaultdict(str)
diff --git a/src/crate/client/sqlalchemy/tests/create_table_test.py b/src/crate/client/sqlalchemy/tests/create_table_test.py
index 0a7dbe8b..bee6b929 100644
--- a/src/crate/client/sqlalchemy/tests/create_table_test.py
+++ b/src/crate/client/sqlalchemy/tests/create_table_test.py
@@ -97,6 +97,28 @@ class DummyTable(self.Base):
              ') CLUSTERED BY (p)\n\n'),
             ())
 
+    def test_with_computed_column(self):
+        class DummyTable(self.Base):
+            __tablename__ = 't'
+            ts = sa.Column(sa.BigInteger, primary_key=True)
+            p = sa.Column(sa.BigInteger, sa.Computed("date_trunc('day', ts)"))
+        self.Base.metadata.create_all()
+        fake_cursor.execute.assert_called_with(
+            ('\nCREATE TABLE t (\n\t'
+             'ts LONG, \n\t'
+             'p LONG GENERATED ALWAYS AS (date_trunc(\'day\', ts)), \n\t'
+             'PRIMARY KEY (ts)\n'
+             ')\n\n'),
+            ())
+
+    def test_with_virtual_computed_column(self):
+        class DummyTable(self.Base):
+            __tablename__ = 't'
+            ts = sa.Column(sa.BigInteger, primary_key=True)
+            p = sa.Column(sa.BigInteger, sa.Computed("date_trunc('day', ts)", persisted=False))
+        with self.assertRaises(sa.exc.CompileError):
+            self.Base.metadata.create_all()
+
     def test_with_partitioned_by(self):
         class DummyTable(self.Base):
             __tablename__ = 't'