diff --git a/server/app/models/connector.rb b/server/app/models/connector.rb index a6f230be7..5d026255d 100644 --- a/server/app/models/connector.rb +++ b/server/app/models/connector.rb @@ -117,7 +117,7 @@ def execute_query(query, limit: 50) query = query.chomp(";") # Check if the query already has a LIMIT clause - has_limit = query.match?(/LIMIT \s*\d+\s*$/i) + has_limit = query.match?(/\bLIMIT\s+\d+/i) # Append LIMIT only if not already present final_query = has_limit ? query : "#{query} LIMIT #{limit}" client.send(:query, db, final_query) diff --git a/server/spec/models/connector_spec.rb b/server/spec/models/connector_spec.rb index 28f2a9c99..44ec82acc 100644 --- a/server/spec/models/connector_spec.rb +++ b/server/spec/models/connector_spec.rb @@ -115,6 +115,49 @@ expect(result).to eq(query_result) end end +<<<<<<< HEAD +======= + + context "when query is for Postgresql connector" do + before do + allow(postgres_connector).to receive(:connector_client).and_return(postgres_client_class) + + allow(postgres_client_class).to receive(:new).and_return(postgres_client_instance) + + allow(postgres_client_instance).to receive(:create_connection) + .with(postgres_connector.configuration.with_indifferent_access) + .and_return(postgres_db_connection) + + allow(postgres_client_instance).to receive(:query) + .with(postgres_db_connection, "SET search_path TO \"public\", \"public\"; #{postgres_query} LIMIT 50") + .and_return(postgres_result) + end + + it "appends a LIMIT clause and executes the query" do + expect(postgres_client_instance) + .to receive(:query) + .with( + postgres_db_connection, + "SET search_path TO \"public\", \"public\"; #{postgres_query} LIMIT 50" + ).and_return(postgres_result) + result = postgres_connector.execute_query(postgres_query) + expect(result).to eq(postgres_result) + end + end + + context "when query has a LIMIT and OFFSET clause" do + let(:query_with_limit_and_offset) { "#{query} LIMIT 50 OFFSET 10" } + + it "executes the query without modifying it" do + expect(client_double) + .to receive(:query) + .with(db_connection, query_with_limit_and_offset) + .and_return(query_result) + result = connector.execute_query(query_with_limit_and_offset) + expect(result).to eq(query_result) + end + end +>>>>>>> e9fcf2dbe (chore(CE): update connector model to handle query pagination safely (#1790)) end describe "#execute_search" do