Add latest changes from gitlab-org/gitlab@master

1 files changed, 33 insertions, 15 deletions

diff --git a/spec/support/database/click_house/hooks.rb b/spec/support/database/click_house/hooks.rb
index b970d3daf84..c13778f9c36 100644
--- a/spec/support/database/click_house/hooks.rb
+++ b/spec/support/database/click_house/hooks.rb
@@ -9,6 +9,8 @@ class ClickHouseTestRunner
         "(SELECT '#{table}' AS table FROM #{table} LIMIT 1)"
       end.join(' UNION ALL ')
 
+      next if query.empty?
+
       tables_with_data = ClickHouse::Client.select(query, db).pluck('table')
       tables_with_data.each do |table|
         ClickHouse::Client.execute("TRUNCATE TABLE #{table}", db)
@@ -16,20 +18,27 @@ class ClickHouseTestRunner
     end
   end
 
-  def ensure_schema
-    return if @ensure_schema
-
-    ClickHouse::Client.configuration.databases.each_key do |db|
+  def clear_db(configuration = ClickHouse::Client.configuration)
+    configuration.databases.each_key do |db|
       # drop all tables
-      lookup_tables(db).each do |table|
-        ClickHouse::Client.execute("DROP TABLE IF EXISTS #{table}", db)
+      lookup_tables(db, configuration).each do |table|
+        ClickHouse::Client.execute("DROP TABLE IF EXISTS #{table}", db, configuration)
       end
 
-      # run the schema SQL files
-      Dir[Rails.root.join("db/click_house/#{db}/*.sql")].each do |file|
-        ClickHouse::Client.execute(File.read(file), db)
-      end
+      ClickHouse::MigrationSupport::SchemaMigration.create_table(db, configuration)
     end
+  end
+
+  def ensure_schema
+    return if @ensure_schema
+
+    clear_db
+
+    # run the schema SQL files
+    migrations_paths = ClickHouse::MigrationSupport::Migrator.migrations_paths
+    schema_migration = ClickHouse::MigrationSupport::SchemaMigration
+    migration_context = ClickHouse::MigrationSupport::MigrationContext.new(migrations_paths, schema_migration)
+    migration_context.up
 
     @ensure_schema = true
   end
@@ -38,11 +47,11 @@ class ClickHouseTestRunner
 
   def tables_for(db)
     @tables ||= {}
-    @tables[db] ||= lookup_tables(db)
+    @tables[db] ||= lookup_tables(db) - [ClickHouse::MigrationSupport::SchemaMigration.table_name]
   end
 
-  def lookup_tables(db)
-    ClickHouse::Client.select('SHOW TABLES', db).pluck('name')
+  def lookup_tables(db, configuration = ClickHouse::Client.configuration)
+    ClickHouse::Client.select('SHOW TABLES', db, configuration).pluck('name')
   end
 end
 # rubocop: enable Gitlab/NamespacedClass
@@ -52,10 +61,19 @@ RSpec.configure do |config|
 
   config.around(:each, :click_house) do |example|
     with_net_connect_allowed do
-      click_house_test_runner.ensure_schema
-      click_house_test_runner.truncate_tables
+      was_verbose = ClickHouse::Migration.verbose
+      ClickHouse::Migration.verbose = false
+
+      if example.example.metadata[:click_house] == :without_migrations
+        click_house_test_runner.clear_db
+      else
+        click_house_test_runner.ensure_schema
+        click_house_test_runner.truncate_tables
+      end
 
       example.run
+    ensure
+      ClickHouse::Migration.verbose = was_verbose
     end
   end
 end