Support computed columns in SQLAlchemy ORM

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
 =================

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

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)

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'