Merge pull request #46 from ThinkThinkAI/db_check

add timeout checks for connections

app/controllers/application_controller.rb

@@ -27,9 +27,15 @@ def check_requirements
 
   def connect_first_data_source
     first_data_source = current_user.data_sources.first
+
     if first_data_source
-      first_data_source.update(connected: true)
-      first_data_source
+      if DatabaseService.test_connection(first_data_source)
+        first_data_source.update(connected: true)
+        first_data_source
+      else
+        redirect_to edit_data_source_path(first_data_source), alert: 'Failed to connect to the DataSource. Please check your configuration.'
+        nil
+      end
     else
       redirect_to new_data_source_path, alert: 'Please create a data source.'
       nil

app/controllers/data_sources_controller.rb

@@ -7,7 +7,7 @@ class DataSourcesController < ApplicationController
   before_action :set_data_source, only: %i[edit update destroy connect]
 
   def show
-    #redirect_to data_sources_path
+    # redirect_to data_sources_path
   end
 
   def new
@@ -17,9 +17,14 @@ def new
   def create
     @data_source = current_user.data_sources.build(data_source_params)
 
-    if @data_source.save
-      redirect_to '/query', notice: 'DataSource was successfully created.'
+    if DatabaseService.test_connection(@data_source)
+      if @data_source.save
+        redirect_to '/query', notice: 'DataSource was successfully created.'
+      else
+        render :new, status: :unprocessable_entity
+      end
     else
+      flash[:alert] = 'Failed to connect to the DataSource. Please check your configuration.'
       render :new, status: :unprocessable_entity
     end
   end
@@ -29,10 +34,17 @@ def edit
   end
 
   def update
-    if @data_source.update(data_source_params)
-      redirect_to '/query', notice: 'DataSource was successfully updated.'
+    @data_source.assign_attributes(data_source_params)
+
+    if DatabaseService.test_connection(@data_source)
+      if @data_source.save
+        redirect_to '/query', notice: 'DataSource was successfully updated.'
+      else
+        render :edit, status: :unprocessable_entity
+      end
     else
-      render :edit
+      flash[:alert] = 'Failed to connect to the DataSource. Please check your configuration.'
+      render :edit, status: :unprocessable_entity
     end
   end
 
@@ -49,28 +61,35 @@ def destroy
   end
 
   def connect
-    updated_status = !@data_source.connected
-
-    puts @data_source.inspect
-    @data_source.update(connected: updated_status)
+    @data_source.assign_attributes(connected: !@data_source.connected)
 
-    @data_source.reload
+    if DatabaseService.test_connection(@data_source)
+      @data_source.save
+      message = 'DataSource connection status was successfully updated.'
 
-    message = 'DataSource connection status was successfully updated.'
-
-    respond_to do |format|
-      format.html { redirect_to '/query', notice: message }
-      format.json { render json: { data_source: @data_source, message: }, status: :ok }
+      respond_to do |format|
+        format.html { redirect_to '/query', notice: message }
+        format.json { render json: { data_source: @data_source, message: }, status: :ok }
+      end
+    else
+      respond_to do |format|
+        format.html do
+          redirect_to request.referer || edit_data_source_path(@data_source),
+                      alert: 'Failed to connect to the DataSource. Please check your configuration.'
+        end
+        format.json do
+          render json: { message: 'Failed to update DataSource connection: Invalid configuration.' },
+                 status: :unprocessable_entity
+        end
+      end
     end
   rescue StandardError => e
-    puts e.inspect
-    puts e.message
-
     respond_to do |format|
-      format.html { redirect_to data_sources_path, alert: "Failed to update DataSource connection: #{e.message}" }
-      format.json do
-        render json: { message: e.message }, status: :unprocessable_entity
+      format.html do
+        redirect_to request.referer || edit_data_source_path(@data_source),
+                    alert: "Failed to update DataSource connection: #{e.message}"
       end
+      format.json { render json: { message: e.message }, status: :unprocessable_entity }
     end
   end
 

app/services/database_service.rb

@@ -1,10 +1,13 @@
 # frozen_string_literal: true
 
+require 'timeout'
+
 # DatabaseService is responsible for creating the appropriate database adapter
 # based on the provided data source configuration.
 class DatabaseService
   ADAPTERS_DIR = File.expand_path('adapters', __dir__)
   ADAPTER_CLASSES = {} # rubocop:disable Style/MutableConstant
+  TIMEOUT_SECONDS = 2
 
   Dir[File.join(ADAPTERS_DIR, '*.rb')].each do |file|
     require_relative file
@@ -21,6 +24,23 @@ def self.build(data_source)
     new(adapter_class.new(data_source), data_source)
   end
 
+  def self.test_connection(data_source)
+    adapter_class = ADAPTER_CLASSES[data_source.adapter]
+    raise "Adapter not supported: #{data_source.adapter}" unless adapter_class
+
+    begin
+      Timeout.timeout(TIMEOUT_SECONDS) do
+        adapter = adapter_class.new(data_source)
+        adapter.run_raw_query('SELECT 1')
+      end
+      true
+    rescue Timeout::Error
+      false
+    rescue StandardError
+      false
+    end
+  end
+
   def initialize(adapter, data_source)
     @adapter = adapter
     @data_source = data_source

app/views/layouts/application.html.erb

@@ -28,6 +28,19 @@
           </div>
         </div>
         <div class="col m-0 p-4 scrollable-content bg-body-tertiary">
+          <% flash.each do |key, value| %>
+            <% alert_class = case key.to_sym
+              when :notice then 'alert alert-info'
+              when :alert  then 'alert alert-danger'
+              when :success then 'alert alert-success'
+              when :warning then 'alert alert-warning'
+              else 'alert alert-primary'
+              end %>
+            <div class="<%= alert_class %> alert-dismissible fade show" role="alert">
+              <%= value %>
+              <button type="button" class="btn-close" data-bs-dismiss="alert" aria-label="Close"></button>
+            </div>
+          <% end %>
           <%= yield %>
         </div>
       </div>

public/404.html

@@ -61,7 +61,6 @@
       <h1>The page you were looking for doesn't exist.</h1>
       <p>You may have mistyped the address or the page may have moved.</p>
     </div>
-    <p>If you are the application owner check the logs for more information.</p>
   </div>
 </body>
 </html>

public/422.html

@@ -61,7 +61,6 @@
       <h1>The change you wanted was rejected.</h1>
       <p>Maybe you tried to change something you didn't have access to.</p>
     </div>
-    <p>If you are the application owner check the logs for more information.</p>
   </div>
 </body>
 </html>

public/500.html

@@ -60,7 +60,6 @@
     <div>
       <h1>We're sorry, but something went wrong.</h1>
     </div>
-    <p>If you are the application owner check the logs for more information.</p>
   </div>
 </body>
 </html>