Merge pull request rails#49346 from fractaledmind/ar-sqlite-virtual-columns

Add support for generated columns in SQLite3 adapter

Author: byroot

activerecord/CHANGELOG.md

@@ -1,3 +1,18 @@
+*   Add support for generated columns in SQLite3 adapter
+
+    Generated columns (both stored and dynamic) are supported since version 3.31.0 of SQLite.
+    This adds support for those to the SQLite3 adapter.
+
+    ```ruby
+    create_table :users do |t|
+      t.string :name
+      t.virtual :name_upper, type: :string, as: 'UPPER(name)'
+      t.virtual :name_lower, type: :string, as: 'LOWER(name)', stored: true
+    end
+    ```
+
+    *Stephen Margheim*
+
 *   TrilogyAdapter: ignore `host` if `socket` parameter is set.
 
     This allows to configure a connection on a UNIX socket via DATABASE_URL:

activerecord/lib/active_record/connection_adapters/sqlite3/column.rb

@@ -6,10 +6,11 @@ module SQLite3
       class Column < ConnectionAdapters::Column # :nodoc:
         attr_reader :rowid
 
-        def initialize(*, auto_increment: nil, rowid: false, **)
+        def initialize(*, auto_increment: nil, rowid: false, generated_type: nil, **)
           super
           @auto_increment = auto_increment
           @rowid = rowid
+          @generated_type = generated_type
         end
 
         def auto_increment?
@@ -20,6 +21,18 @@ def auto_incremented_by_db?
           auto_increment? || rowid
         end
 
+        def virtual?
+          !@generated_type.nil?
+        end
+
+        def virtual_stored?
+          virtual? && @generated_type == :stored
+        end
+
+        def has_default?
+          super && !virtual?
+        end
+
         def init_with(coder)
           @auto_increment = coder["auto_increment"]
           super

activerecord/lib/active_record/connection_adapters/sqlite3/schema_creation.rb

@@ -25,6 +25,16 @@ def add_column_options!(sql, options)
             if options[:collation]
               sql << " COLLATE \"#{options[:collation]}\""
             end
+
+            if as = options[:as]
+              sql << " GENERATED ALWAYS AS (#{as})"
+
+              if options[:stored]
+                sql << " STORED"
+              else
+                sql << " VIRTUAL"
+              end
+            end
             super
           end
       end

activerecord/lib/active_record/connection_adapters/sqlite3/schema_definitions.rb

@@ -16,10 +16,23 @@ def references(*args, **options)
         end
         alias :belongs_to :references
 
+        def new_column_definition(name, type, **options) # :nodoc:
+          case type
+          when :virtual
+            type = options[:type]
+          end
+
+          super
+        end
+
         private
           def integer_like_primary_key_type(type, options)
             :primary_key
           end
+
+          def valid_column_definition_options
+            super + [:as, :type, :stored]
+          end
       end
     end
   end

activerecord/lib/active_record/connection_adapters/sqlite3/schema_dumper.rb

@@ -12,6 +12,22 @@ def default_primary_key?(column)
           def explicit_primary_key_default?(column)
             column.bigint?
           end
+
+          def prepare_column_options(column)
+            spec = super
+
+            if @connection.supports_virtual_columns? && column.virtual?
+              spec[:as] = extract_expression_for_virtual_column(column)
+              spec[:stored] = column.virtual_stored?
+              spec = { type: schema_type(column).inspect }.merge!(spec)
+            end
+
+            spec
+          end
+
+          def extract_expression_for_virtual_column(column)
+            column.default_function.inspect
+          end
       end
     end
   end

activerecord/lib/active_record/connection_adapters/sqlite3/schema_statements.rb

@@ -147,7 +147,14 @@ def new_column_from_field(table_name, field, definitions)
 
             type_metadata = fetch_type_metadata(field["type"])
             default_value = extract_value_from_default(default)
-            default_function = extract_default_function(default_value, default)
+            generated_type = extract_generated_type(field)
+
+            if generated_type.present?
+              default_function = default
+            else
+              default_function = extract_default_function(default_value, default)
+            end
+
             rowid = is_column_the_rowid?(field, definitions)
 
             Column.new(
@@ -158,7 +165,8 @@ def new_column_from_field(table_name, field, definitions)
               default_function,
               collation: field["collation"],
               auto_increment: field["auto_increment"],
-              rowid: rowid
+              rowid: rowid,
+              generated_type: generated_type
             )
           end
 
@@ -201,6 +209,13 @@ def assert_valid_deferrable(deferrable)
 
             raise ArgumentError, "deferrable must be `:immediate` or `:deferred`, got: `#{deferrable.inspect}`"
           end
+
+          def extract_generated_type(field)
+            case field["hidden"]
+            when 2 then :virtual
+            when 3 then :stored
+            end
+          end
       end
     end
   end

activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb

@@ -186,6 +186,10 @@ def supports_concurrent_connections?
         !@memory_database
       end
 
+      def supports_virtual_columns?
+        database_version >= "3.31.0"
+      end
+
       def connected?
         !(@raw_connection.nil? || @raw_connection.closed?)
       end
@@ -282,6 +286,7 @@ def rename_table(table_name, new_name, **options)
       end
 
       def add_column(table_name, column_name, type, **options) # :nodoc:
+        type = type.to_sym
         if invalid_alter_table_type?(type, options)
           alter_table(table_name) do |definition|
             definition.column(column_name, type, **options)
@@ -462,7 +467,11 @@ def bind_params_length
         end
 
         def table_structure(table_name)
-          structure = internal_exec_query("PRAGMA table_info(#{quote_table_name(table_name)})", "SCHEMA")
+          structure = if supports_virtual_columns?
+            internal_exec_query("PRAGMA table_xinfo(#{quote_table_name(table_name)})", "SCHEMA")
+          else
+            internal_exec_query("PRAGMA table_info(#{quote_table_name(table_name)})", "SCHEMA")
+          end
           raise(ActiveRecord::StatementInvalid, "Could not find table '#{table_name}'") if structure.empty?
           table_structure_with_collation(table_name, structure)
         end
@@ -502,8 +511,9 @@ def has_default_function?(default_value, default)
         # See: https://www.sqlite.org/lang_altertable.html
         # SQLite has an additional restriction on the ALTER TABLE statement
         def invalid_alter_table_type?(type, options)
-          type.to_sym == :primary_key || options[:primary_key] ||
-            options[:null] == false && options[:default].nil?
+          type == :primary_key || options[:primary_key] ||
+            options[:null] == false && options[:default].nil? ||
+            (type == :virtual && options[:stored])
         end
 
         def alter_table(
@@ -651,10 +661,12 @@ def translate_exception(exception, message:, sql:, binds:)
 
         COLLATE_REGEX = /.*"(\w+)".*collate\s+"(\w+)".*/i
         PRIMARY_KEY_AUTOINCREMENT_REGEX = /.*"(\w+)".+PRIMARY KEY AUTOINCREMENT/i
+        GENERATED_ALWAYS_AS_REGEX = /.*"(\w+)".+GENERATED ALWAYS AS \((.+)\) (?:STORED|VIRTUAL)/i
 
         def table_structure_with_collation(table_name, basic_structure)
           collation_hash = {}
           auto_increments = {}
+          generated_columns = {}
 
           column_strings = table_structure_sql(table_name)
 
@@ -664,6 +676,7 @@ def table_structure_with_collation(table_name, basic_structure)
               # the value in $1 and $2 respectively.
               collation_hash[$1] = $2 if COLLATE_REGEX =~ column_string
               auto_increments[$1] = true if PRIMARY_KEY_AUTOINCREMENT_REGEX =~ column_string
+              generated_columns[$1] = $2 if GENERATED_ALWAYS_AS_REGEX =~ column_string
             end
 
             basic_structure.map do |column|
@@ -677,6 +690,10 @@ def table_structure_with_collation(table_name, basic_structure)
                 column["auto_increment"] = true
               end
 
+              if generated_columns.has_key?(column_name)
+                column["dflt_value"] = generated_columns[column_name]
+              end
+
               column
             end
           else

activerecord/test/cases/adapters/sqlite3/sqlite3_adapter_test.rb

@@ -263,7 +263,7 @@ def test_insert_logged
           sql = "INSERT INTO ex (number) VALUES (10)"
           name = "foo"
 
-          pragma_query = ["PRAGMA table_info(\"ex\")", "SCHEMA", []]
+          pragma_query = ["PRAGMA table_xinfo(\"ex\")", "SCHEMA", []]
           schema_query = ["SELECT sql FROM (SELECT * FROM sqlite_master UNION ALL SELECT * FROM sqlite_temp_master) WHERE type = 'table' AND name = 'ex'", "SCHEMA", []]
           modified_insert_query = [(sql + ' RETURNING "id"'), name, []]
           assert_logged [pragma_query, schema_query, modified_insert_query] do

activerecord/test/cases/adapters/sqlite3/virtual_column_test.rb

@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+
+require "cases/helper"
+require "support/schema_dumping_helper"
+
+if ActiveRecord::Base.connection.supports_virtual_columns?
+  class SQLite3VirtualColumnTest < ActiveRecord::SQLite3TestCase
+    include SchemaDumpingHelper
+
+    class VirtualColumn < ActiveRecord::Base
+    end
+
+    def setup
+      @connection = ActiveRecord::Base.connection
+      @connection.create_table :virtual_columns, force: true do |t|
+        t.string  :name
+        t.virtual :upper_name, type: :string, as: "UPPER(name)", stored: true
+        t.virtual :lower_name, type: :string, as: "LOWER(name)", stored: false
+        t.virtual :octet_name, type: :integer, as: "LENGTH(name)"
+        t.integer :column1
+      end
+      VirtualColumn.create(name: "Rails", column1: 10)
+    end
+
+    def teardown
+      @connection.drop_table :virtual_columns, if_exists: true
+      VirtualColumn.reset_column_information
+    end
+
+    def test_virtual_column_with_full_inserts
+      partial_inserts_was = VirtualColumn.partial_inserts
+      VirtualColumn.partial_inserts = false
+      assert_nothing_raised do
+        VirtualColumn.create!(name: "Rails")
+      end
+    ensure
+      VirtualColumn.partial_inserts = partial_inserts_was
+    end
+
+    def test_stored_column
+      column = VirtualColumn.columns_hash["upper_name"]
+      assert_predicate column, :virtual?
+      assert_predicate column, :virtual_stored?
+      assert_equal "RAILS", VirtualColumn.take.upper_name
+    end
+
+    def test_explicit_virtual_column
+      column = VirtualColumn.columns_hash["lower_name"]
+      assert_predicate column, :virtual?
+      assert_not_predicate column, :virtual_stored?
+      assert_equal "rails", VirtualColumn.take.lower_name
+    end
+
+    def test_implicit_virtual_column
+      column = VirtualColumn.columns_hash["octet_name"]
+      assert_predicate column, :virtual?
+      assert_not_predicate column, :virtual_stored?
+      assert_equal 5, VirtualColumn.take.octet_name
+    end
+
+    def test_change_table_with_stored_generated_column
+      @connection.change_table :virtual_columns do |t|
+        t.virtual :decr_column1, type: :integer, as: "column1 - 1", stored: true
+      end
+      VirtualColumn.reset_column_information
+      column = VirtualColumn.columns_hash["decr_column1"]
+      assert_predicate column, :virtual?
+      assert_predicate column, :virtual_stored?
+      assert_equal 9, VirtualColumn.take.decr_column1
+    end
+
+    def test_change_table_with_explicit_virtual_generated_column
+      @connection.change_table :virtual_columns do |t|
+        t.virtual :incr_column1, type: :integer, as: "column1 + 1", stored: false
+      end
+      VirtualColumn.reset_column_information
+      column = VirtualColumn.columns_hash["incr_column1"]
+      assert_predicate column, :virtual?
+      assert_not_predicate column, :virtual_stored?
+      assert_equal 11, VirtualColumn.take.incr_column1
+    end
+
+    def test_change_table_with_implicit_virtual_generated_column
+      @connection.change_table :virtual_columns do |t|
+        t.virtual :sqr_column1, type: :integer, as: "pow(column1, 2)"
+      end
+      VirtualColumn.reset_column_information
+      column = VirtualColumn.columns_hash["sqr_column1"]
+      assert_predicate column, :virtual?
+      assert_not_predicate column, :virtual_stored?
+      assert_equal 100, VirtualColumn.take.sqr_column1
+    end
+
+    def test_schema_dumping
+      output = dump_table_schema("virtual_columns")
+      assert_match(/t\.virtual\s+"upper_name",\s+type: :string,\s+as: "UPPER\(name\)", stored: true$/i, output)
+    end
+
+    def test_build_fixture_sql
+      ActiveRecord::FixtureSet.create_fixtures(FIXTURES_ROOT, :virtual_columns)
+    end
+  end
+end

activerecord/test/cases/migration/invalid_options_test.rb

@@ -14,6 +14,8 @@ def invalid_add_column_option_exception_message(key)
           default_keys.concat([":auto_increment", ":charset", ":as", ":size", ":unsigned", ":first", ":after", ":type", ":stored"])
         elsif current_adapter?(:PostgreSQLAdapter)
           default_keys.concat([":array", ":using", ":cast_as", ":as", ":type", ":enum_type", ":stored"])
+        elsif current_adapter?(:SQLite3Adapter)
+          default_keys.concat([":as", ":type", ":stored"])
         end
 
         "Unknown key: :#{key}. Valid keys are: #{default_keys.join(", ")}"