Skip to content

Commit e2afe65

Browse files
authored
feature: jobs extension (#14)
* Job Run Status Functionsg * Updated test for jobs feature extension * Format Changes with Make Format
1 parent a4a5249 commit e2afe65

File tree

5 files changed

+206
-2
lines changed

5 files changed

+206
-2
lines changed

Makefile

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,50 @@ build-all:
5656
test: build-tests
5757
cd $(BUILD_DIR) && ctest --output-on-failure
5858

59+
# Run only tests related to modified files (based on git status)
60+
.PHONY: test-new
61+
test-new: build-tests
62+
@echo "Running tests for modified files..."
63+
@MODIFIED_FILES=$$(git status --short | grep -E '^\s*M.*\.(cpp|h)' | awk '{print $$2}' | grep -E 'test|jobs|compute|unity' || true); \
64+
if [ -z "$$MODIFIED_FILES" ]; then \
65+
echo "No modified test files detected. Running all tests..."; \
66+
cd $(BUILD_DIR)/tests && ./unit_tests; \
67+
else \
68+
echo "Modified files: $$MODIFIED_FILES"; \
69+
if echo "$$MODIFIED_FILES" | grep -q "jobs"; then \
70+
echo "Running Jobs tests..."; \
71+
cd $(BUILD_DIR)/tests && ./unit_tests --gtest_filter='*Job*'; \
72+
elif echo "$$MODIFIED_FILES" | grep -q "compute"; then \
73+
echo "Running Compute tests..."; \
74+
cd $(BUILD_DIR)/tests && ./unit_tests --gtest_filter='*Compute*:*Cluster*'; \
75+
elif echo "$$MODIFIED_FILES" | grep -q "unity"; then \
76+
echo "Running Unity Catalog tests..."; \
77+
cd $(BUILD_DIR)/tests && ./unit_tests --gtest_filter='*UnityCatalog*'; \
78+
else \
79+
echo "Running all tests for safety..."; \
80+
cd $(BUILD_DIR)/tests && ./unit_tests; \
81+
fi \
82+
fi
83+
84+
# Run tests with custom filter
85+
# Usage: make test-filter FILTER='JobsApiTest.*'
86+
.PHONY: test-filter
87+
test-filter: build-tests
88+
@if [ -z "$(FILTER)" ]; then \
89+
echo "Usage: make test-filter FILTER='pattern'"; \
90+
echo "Example: make test-filter FILTER='JobsApiTest.*'"; \
91+
echo "Example: make test-filter FILTER='*Cancel*'"; \
92+
exit 1; \
93+
fi
94+
@echo "Running tests matching filter: $(FILTER)"
95+
cd $(BUILD_DIR)/tests && ./unit_tests --gtest_filter='$(FILTER)'
96+
97+
# List all available tests
98+
.PHONY: test-list
99+
test-list: build-tests
100+
@echo "Available test cases:"
101+
cd $(BUILD_DIR)/tests && ./unit_tests --gtest_list_tests
102+
59103
# Clean build artifacts
60104
.PHONY: clean
61105
clean:
@@ -166,7 +210,10 @@ help:
166210
@echo " make clean - Remove build artifacts"
167211
@echo ""
168212
@echo "Testing:"
169-
@echo " make test - Run tests"
213+
@echo " make test - Run all tests"
214+
@echo " make test-new - Run tests for modified files (smart)"
215+
@echo " make test-filter FILTER='pattern' - Run tests matching pattern"
216+
@echo " make test-list - List all available test cases"
170217
@echo " make benchmark - Run connection pooling benchmark"
171218
@echo ""
172219
@echo "Documentation:"

include/databricks/jobs/jobs.h

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,26 @@ class Jobs {
9797
*/
9898
uint64_t run_now(uint64_t job_id, const std::map<std::string, std::string>& notebook_params = {});
9999

100+
/**
101+
* @brief Cancel a running or pending job run
102+
*
103+
* @param run_id The unique identifier of the job run to cancel
104+
* @return true if the cancellation request was successful
105+
* @throws std::runtime_error if the API request fails
106+
*/
107+
bool cancel_run(uint64_t run_id);
108+
109+
/**
110+
* @brief Get the output of a completed job run
111+
*
112+
* @param run_id The unique identifier of the job run output
113+
* @return RunOuput object containing execution output and logs
114+
* @throws std::runtime_error if the API request fails, or if the run is not found
115+
*
116+
* @note Only completed runs have an output. Running jobs will return an error
117+
*/
118+
RunOutput get_run_output(uint64_t run_id);
119+
100120
private:
101121
class Impl;
102122
std::unique_ptr<Impl> pimpl_;

include/databricks/jobs/jobs_types.h

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,4 +54,23 @@ struct JobRun {
5454
static JobRun from_json(const std::string& json_str);
5555
};
5656

57+
/**
58+
* @brief Represents a run output object for a Databricks job
59+
*
60+
* This struct contains the key fields from a Databricks job including logs, errors, and metadata.
61+
*/
62+
struct RunOutput {
63+
std::string notebook_output; ///< Output from notebook tasks (JSON)
64+
std::string sql_output; ///< Output from SQL tasks
65+
std::string logs; ///< Execution logs
66+
std::string error; ///< Error message if run failed
67+
std::string error_trace; ///< Stack trace if available
68+
std::map<std::string, std::string> metadata; ///< Additional output metadata
69+
70+
/**
71+
* @brief Parse RunOutput from JSON string
72+
*/
73+
static RunOutput from_json(const std::string& json_str);
74+
};
75+
5776
} // namespace databricks

src/jobs/jobs.cpp

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ std::string build_query_string(const std::map<std::string, std::string>& params)
5353
} // namespace
5454

5555
// ============================================================================
56-
// Job and JobRun JSON Parsing
56+
// JSON Parsing Helper Functions
5757
// ============================================================================
5858

5959
Job Job::from_json(const std::string& json_str) {
@@ -101,6 +101,27 @@ JobRun JobRun::from_json(const std::string& json_str) {
101101
}
102102
}
103103

104+
RunOutput RunOutput::from_json(const std::string& json_str) {
105+
try {
106+
auto j = json::parse(json_str);
107+
RunOutput run_output;
108+
109+
run_output.notebook_output = j.value("notebook_output", "");
110+
run_output.sql_output = j.value("sql_output", "");
111+
run_output.logs = j.value("logs", "");
112+
run_output.error = j.value("error", "");
113+
run_output.error_trace = j.value("error_trace", "");
114+
115+
if (j.contains("metadata")) {
116+
run_output.metadata["raw"] = j["metadata"].dump();
117+
}
118+
119+
return run_output;
120+
} catch (const json::exception& e) {
121+
throw std::runtime_error("Failed to parse RunOutput JSON: " + std::string(e.what()));
122+
}
123+
}
124+
104125
// ============================================================================
105126
// Jobs Constructor and Destructor
106127
// ============================================================================
@@ -192,6 +213,40 @@ uint64_t Jobs::run_now(uint64_t job_id, const std::map<std::string, std::string>
192213
}
193214
}
194215

216+
bool Jobs::cancel_run(uint64_t run_id) {
217+
internal::get_logger()->info("Cancelling run for run_id=" + std::to_string(run_id));
218+
219+
// Build request body
220+
json body_json;
221+
body_json["run_id"] = run_id;
222+
std::string body = body_json.dump();
223+
224+
internal::get_logger()->debug("Request body: " + body);
225+
226+
// Make API request
227+
auto response = pimpl_->http_client_->post("/jobs/runs/cancel", body);
228+
pimpl_->http_client_->check_response(response, "cancelJob");
229+
internal::get_logger()->info("Successfully cancelled run for run_id=" + std::to_string(run_id));
230+
231+
return true;
232+
}
233+
234+
RunOutput Jobs::get_run_output(uint64_t run_id) {
235+
internal::get_logger()->info("Retrieving the output for run_id=" + std::to_string(run_id));
236+
237+
// Build query parameters
238+
std::map<std::string, std::string> params;
239+
params["run_id"] = std::to_string(run_id);
240+
241+
// Make API request
242+
std::string query = build_query_string(params);
243+
auto response = pimpl_->http_client_->get("/jobs/runs/get-output" + query);
244+
pimpl_->http_client_->check_response(response, "getRunOutput");
245+
246+
internal::get_logger()->debug("Job run output response: " + response.body);
247+
return RunOutput::from_json(response.body);
248+
}
249+
195250
// ============================================================================
196251
// Private Helper Methods
197252
// ============================================================================

tests/unit/jobs/test_jobs.cpp

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,15 @@
11
// Copyright (c) 2025 Calvin Min
22
// SPDX-License-Identifier: MIT
3+
#include "../../mocks/mock_http_client.h"
4+
35
#include <databricks/core/config.h>
46
#include <databricks/jobs/jobs.h>
57
#include <gtest/gtest.h>
68

9+
using databricks::test::MockHttpClient;
10+
using ::testing::_;
11+
using ::testing::Return;
12+
713
// Test fixture for Jobs tests
814
class JobsTest : public ::testing::Test {
915
protected:
@@ -16,6 +22,18 @@ class JobsTest : public ::testing::Test {
1622
}
1723
};
1824

25+
// Test fixture for Jobs API tests with mocks
26+
class JobsApiTest : public ::testing::Test {
27+
protected:
28+
databricks::AuthConfig auth;
29+
30+
void SetUp() override {
31+
auth.host = "https://test.databricks.com";
32+
auth.set_token("test_token");
33+
auth.timeout_seconds = 30;
34+
}
35+
};
36+
1937
// Test: Jobs client construction
2038
TEST_F(JobsTest, ConstructorCreatesValidClient) {
2139
ASSERT_NO_THROW({ databricks::Jobs jobs(auth); });
@@ -227,3 +245,48 @@ TEST_F(JobsTest, MultipleClientsCanCoexist) {
227245
// Both should coexist without issues
228246
});
229247
}
248+
249+
// Test: Cancel Run
250+
TEST_F(JobsApiTest, CancelRunReturnsTrueAndCallsApi) {
251+
// Setup
252+
auto mock_client = std::make_shared<MockHttpClient>();
253+
EXPECT_CALL(*mock_client, post("/jobs/runs/cancel", _))
254+
.WillOnce(Return(MockHttpClient::success_response(R"({"result":"OK"})")));
255+
256+
// check_response should be called and not throw
257+
EXPECT_CALL(*mock_client, check_response(_, "cancelJob")).Times(1);
258+
259+
// Execute call with Mock Client
260+
databricks::Jobs jobs(mock_client);
261+
EXPECT_TRUE(jobs.cancel_run(42));
262+
}
263+
264+
// Test: Get Run Output of Completed Run
265+
TEST_F(JobsApiTest, GetRunOutputCompleted) {
266+
// Setup
267+
auto mock_client = std::make_shared<MockHttpClient>();
268+
269+
EXPECT_CALL(*mock_client, get("/jobs/runs/get-output?run_id=123"))
270+
.WillOnce(Return(MockHttpClient::success_response(R"({"notebook_output":"success"} )")));
271+
272+
EXPECT_CALL(*mock_client, check_response(_, "getRunOutput")).Times(1);
273+
274+
databricks::Jobs jobs(mock_client);
275+
auto output = jobs.get_run_output(123);
276+
EXPECT_EQ(output.notebook_output, "success");
277+
}
278+
279+
// Test: Get Run Output of a Failed Run
280+
TEST_F(JobsApiTest, GetRunOutputFailedRun) {
281+
// Setup
282+
auto mock_client = std::make_shared<MockHttpClient>();
283+
EXPECT_CALL(*mock_client, get("/jobs/runs/get-output?run_id=123"))
284+
.WillOnce(Return(MockHttpClient::success_response(R"({"error":"error message","notebook_output":"failed"})")));
285+
286+
EXPECT_CALL(*mock_client, check_response(_, "getRunOutput")).Times(1);
287+
288+
databricks::Jobs jobs(mock_client);
289+
auto output = jobs.get_run_output(123);
290+
EXPECT_EQ(output.error, "error message");
291+
EXPECT_EQ(output.notebook_output, "failed");
292+
}

0 commit comments

Comments
 (0)