Minor database tool fixes (#204)

* Minor database tool fixes

* Fix database tool spec

---------

Co-authored-by: Andrei Bondarev <andrei@sourcelabs.io>

lib/langchain/tool/database.rb

@@ -14,7 +14,7 @@ class Database < Base
       The input to this tool should be valid SQL.
     DESC
 
-    attr_reader :db, :requested_tables, :except_tables
+    attr_reader :db, :requested_tables, :excluded_tables
 
     #
     # Establish a database connection
@@ -25,16 +25,15 @@ class Database < Base
 
     # @return [Database] Database object
     #
-    def initialize(connection_string:, tables: [], except_tables: [])
+    def initialize(connection_string:, tables: [], exclude_tables: [])
       depends_on "sequel"
       require "sequel"
-      require "sequel/extensions/schema_dumper"
 
       raise StandardError, "connection_string parameter cannot be blank" if connection_string.empty?
 
       @db = Sequel.connect(connection_string)
       @requested_tables = tables
-      @except_tables = except_tables
+      @excluded_tables = exclude_tables
     end
 
     #
@@ -46,20 +45,31 @@ def dump_schema
       Langchain.logger.info("Dumping schema tables and keys", for: self.class)
       schema = ""
       db.tables.each do |table|
-        next if except_tables.include?(table)
+        next if excluded_tables.include?(table)
         next unless requested_tables.empty? || requested_tables.include?(table)
 
+        primary_key_columns = []
+        primary_key_column_count = db.schema(table).count { |column| column[1][:primary_key] == true }
+
         schema << "CREATE TABLE #{table}(\n"
         db.schema(table).each do |column|
           schema << "#{column[0]} #{column[1][:type]}"
-          schema << " PRIMARY KEY" if column[1][:primary_key] == true
-          schema << "," unless column == db.schema(table).last
-          schema << "\n"
+          if column[1][:primary_key] == true
+            schema << " PRIMARY KEY" if primary_key_column_count == 1
+          else
+            primary_key_columns << column[0]
+          end
+          schema << ",\n" unless column == db.schema(table).last && primary_key_column_count == 1
+        end
+        if primary_key_column_count > 1
+          schema << "PRIMARY KEY (#{primary_key_columns.join(",")})"
         end
-        schema << ");\n"
         db.foreign_key_list(table).each do |fk|
-          schema << "ALTER TABLE #{table} ADD FOREIGN KEY (#{fk[:columns][0]}) REFERENCES #{fk[:table]}(#{fk[:key][0]});\n"
+          schema << ",\n" if fk == db.foreign_key_list(table).first
+          schema << "FOREIGN KEY (#{fk[:columns][0]}) REFERENCES #{fk[:table]}(#{fk[:key][0]})"
+          schema << ",\n" unless fk == db.foreign_key_list(table).last
         end
+        schema << ");\n"
       end
       schema
     end

spec/langchain/tool/database_spec.rb

@@ -39,7 +39,7 @@
     end
 
     it "returns the schema" do
-      expect(subject.dump_schema).to eq("CREATE TABLE users(\nid integer PRIMARY KEY,\nname string,\njob string\n);\nALTER TABLE users ADD FOREIGN KEY (job) REFERENCES jobs(job);\n")
+      expect(subject.dump_schema).to eq("CREATE TABLE users(\nid integer PRIMARY KEY,\nname string,\njob string,\nFOREIGN KEY (job) REFERENCES jobs(job));\n")
     end
   end
 end