From a152c0abe14032a189db860f46caedbfad2bef53 Mon Sep 17 00:00:00 2001
From: Fynn Starke <fynsta1904@gmail.com>
Date: Thu, 4 Apr 2024 18:00:21 +0200
Subject: [PATCH] Actually fix the template file

---
 lib/squasher/cleaner.rb |  2 +-
 lib/squasher/config.rb  | 21 +++++++++++++++------
 lib/squasher/render.rb  |  5 +++--
 lib/squasher/worker.rb  |  5 +++--
 4 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/lib/squasher/cleaner.rb b/lib/squasher/cleaner.rb
index 9e651f4..5878aa7 100644
--- a/lib/squasher/cleaner.rb
+++ b/lib/squasher/cleaner.rb
@@ -26,7 +26,7 @@ def process_database(database = nil)
         FileUtils.rm(prev_migration)
       end
       File.open(migration_file, 'wb') do |stream|
-        stream << ::Squasher::Render.render(MIGRATION_NAME, config)
+        stream << ::Squasher::Render.render(MIGRATION_NAME, config, database)
       end
 
       if database.nil?
diff --git a/lib/squasher/config.rb b/lib/squasher/config.rb
index 51807c8..525ad37 100644
--- a/lib/squasher/config.rb
+++ b/lib/squasher/config.rb
@@ -34,7 +34,7 @@ def inspect
       end
     end
 
-    attr_reader :schema_file, :migration_version, :multi_db_format, :databases
+    attr_reader :migration_version, :multi_db_format, :databases
 
     def initialize
       @root_path = Dir.pwd.freeze
@@ -59,9 +59,6 @@ def set(key, value)
       elsif key == :migration
         Squasher.error(:invalid_migration_version, value: value) unless value.to_s =~ /\A\d.\d\z/
         @migration_version = "[#{value}]"
-      elsif key == :sql
-        @schema_file = File.join(@app_path, 'db', 'structure.sql')
-        @flags << key
       elsif key == :multi_db_format
         Squasher.error(:invalid_multi_db_format, value: value) unless %w[rails multiverse].include?(value)
         @multi_db_format = value
@@ -76,6 +73,19 @@ def set?(k)
       @flags.include?(k)
     end
 
+    def schema_files
+      return [schema_file] unless @multi_db_format == 'rails'
+
+      @databases.map { |db| schema_file(db) }
+    end
+
+    def schema_file(database = nil)
+      prefix = database.nil? || database == 'primary' ? '' : "#{ database }_"
+      file = set?(:sql) ? 'structure.sql' : 'schema.rb'
+
+      File.join(@app_path, 'db', "#{ prefix }#{ file }")
+    end
+
     def migration_files(database = nil)
       Dir.glob(File.join(migrations_folder(database), '**.rb'))
     end
@@ -112,7 +122,7 @@ def dbconfig?
     def stub_dbconfig
       return unless dbconfig?
 
-      list = [dbconfig_file, schema_file]
+      list = [dbconfig_file, *schema_files]
       list.each do |file|
         next unless File.exist?(file)
         FileUtils.mv file, "#{ file }.sq"
@@ -179,7 +189,6 @@ def dbconfig
 
     def set_app_path(path)
       @app_path = path
-      @schema_file = File.join(path, 'db', 'schema.rb')
       @dbconfig_file = File.join(path, 'config', 'database.yml')
     end
   end
diff --git a/lib/squasher/render.rb b/lib/squasher/render.rb
index 57e6326..f77938b 100644
--- a/lib/squasher/render.rb
+++ b/lib/squasher/render.rb
@@ -8,9 +8,10 @@ def self.render(*args)
 
     attr_reader :name, :config
 
-    def initialize(name, config)
+    def initialize(name, config, database = nil)
       @name = name
       @config = config
+      @database = database
     end
 
     def render
@@ -22,7 +23,7 @@ def render
     end
 
     def each_schema_line(&block)
-      File.open(config.schema_file, 'r') do |stream|
+      File.open(config.schema_file(@database), 'r') do |stream|
         if @config.set?(:sql)
           stream_structure(stream, &block)
         else
diff --git a/lib/squasher/worker.rb b/lib/squasher/worker.rb
index b894c53..c177b42 100644
--- a/lib/squasher/worker.rb
+++ b/lib/squasher/worker.rb
@@ -1,4 +1,5 @@
 require 'fileutils'
+require 'pry'
 
 module Squasher
   class Worker
@@ -72,8 +73,8 @@ def get_timestamp(file)
 
     def clean_migrations(database = nil)
       path = config.migration_file(finish_timestamp(database), :init_schema, database)
-      File.open(path, 'wb') { |io| io << Render.render(:init_schema, config) }
-      migrations(database).each { |file| FileUtils.rm(file) }
+      migrations(database).each { |file| FileUtils.rm(file) } # Remove all migrations before creating the new one
+      File.open(path, 'wb') { |io| io << Render.render(:init_schema, config, database) }
     end
 
     def before_date?(timestamp)