Skip to content

Add ability to set RequestOptions #99

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zack13532 merged 1 commit into from
Jul 24, 2025
Merged

Add ability to set RequestOptions #99

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 101 additions & 6 deletions src/main/java/ai/tecton/client/request/GetFeaturesBatchRequest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ public class GetFeaturesBatchRequest {
private List<? extends AbstractGetFeaturesRequest> requestList;
private final int microBatchSize;
private final Duration timeout;
private final RequestOptions requestOptions;
private static final String BATCH_ENDPOINT = "/api/v1/feature-service/get-features-batch";
private static JsonAdapter<GetFeaturesMicroBatchRequest.GetFeaturesRequestBatchJson> jsonAdapter =
null;
Expand Down Expand Up @@ -82,7 +83,8 @@ public GetFeaturesBatchRequest(
requestDataList,
RequestConstants.DEFAULT_METADATA_OPTIONS,
RequestConstants.DEFAULT_MICRO_BATCH_SIZE,
RequestConstants.NONE_TIMEOUT);
RequestConstants.NONE_TIMEOUT,
null);
}

/**
Expand Down Expand Up @@ -116,7 +118,8 @@ public GetFeaturesBatchRequest(
requestDataList,
metadataOptions,
RequestConstants.DEFAULT_MICRO_BATCH_SIZE,
RequestConstants.NONE_TIMEOUT);
RequestConstants.NONE_TIMEOUT,
null);
}

/**
Expand Down Expand Up @@ -154,7 +157,8 @@ public GetFeaturesBatchRequest(
requestDataList,
metadataOptions,
microBatchSize,
RequestConstants.NONE_TIMEOUT);
RequestConstants.NONE_TIMEOUT,
null);
}

/**
Expand Down Expand Up @@ -191,8 +195,60 @@ public GetFeaturesBatchRequest(
Set<MetadataOption> metadataOptions,
int microBatchSize,
Duration timeout) {
this(
workspaceName,
featureServiceName,
requestDataList,
metadataOptions,
microBatchSize,
timeout,
null);
}

/**
* Constructor that creates a new GetFeaturesBatchRequest with the specified parameters including
* requestOptions
*
* @param workspaceName Name of the workspace in which the Feature Service is defined
* @param featureServiceName Name of the Feature Service for which the feature vectors are being
* requested
* @param requestDataList a {@link List} of {@link GetFeaturesRequestData} object with joinKeyMap
* and/or requestContextMap
* @param metadataOptions metadataOptions A {@link Set} of {@link MetadataOption} for retrieving
* additional metadata about the feature values. Use {@link
* RequestConstants#ALL_METADATA_OPTIONS} to request all metadata and {@link
* RequestConstants#NONE_METADATA_OPTIONS} to request no metadata respectively. By default,
* {@link RequestConstants#DEFAULT_METADATA_OPTIONS} will be added to each request
* @param microBatchSize an int value between 1 and {@value
* RequestConstants#MAX_MICRO_BATCH_SIZE}. The client splits the GetFeaturesBatchRequest into
* multiple micro batches of this size and executes them parallely. By default, the
* microBatchSize is set to {@value RequestConstants#DEFAULT_MICRO_BATCH_SIZE}
* @param timeout The max time in {@link Duration} for which the client waits for the batch
* requests to complete before canceling the operation and returning the partial list of
* results.
* @param requestOptions {@link RequestOptions} object with request-level options to control
* feature server behavior
* @throws InvalidRequestParameterException when workspaceName or featureServiceName is empty or
* null
* @throws InvalidRequestParameterException when requestDataList is invalid (null/empty or
* contains null/empty elements)
* @throws InvalidRequestParameterException when the microBatchSize is out of bounds of [ 1,
* {@value RequestConstants#MAX_MICRO_BATCH_SIZE} ]
*/
public GetFeaturesBatchRequest(
String workspaceName,
String featureServiceName,
List<GetFeaturesRequestData> requestDataList,
Set<MetadataOption> metadataOptions,
int microBatchSize,
Duration timeout,
RequestOptions requestOptions) {
validateParameters(workspaceName, featureServiceName, requestDataList, microBatchSize);
this.timeout = timeout;
this.requestOptions = requestOptions;

// Create final variable for use in lambda expressions
final RequestOptions finalRequestOptions = requestOptions;

if (microBatchSize > 1 && requestDataList.size() > 1) {
// For batch requests, partition the requestDataList into n sublists of size
Expand All @@ -203,7 +259,11 @@ public GetFeaturesBatchRequest(
.map(
requestData ->
new GetFeaturesMicroBatchRequest(
workspaceName, featureServiceName, requestData, metadataOptions))
workspaceName,
featureServiceName,
requestData,
metadataOptions,
finalRequestOptions))
.collect(Collectors.toList());
this.microBatchSize = microBatchSize;
jsonAdapter = moshi.adapter(GetFeaturesMicroBatchRequest.GetFeaturesRequestBatchJson.class);
Expand All @@ -217,7 +277,11 @@ public GetFeaturesBatchRequest(
.map(
requestData ->
new GetFeaturesRequest(
workspaceName, featureServiceName, requestData, metadataOptions))
workspaceName,
featureServiceName,
requestData,
metadataOptions,
finalRequestOptions))
.collect(Collectors.toList());
this.microBatchSize = microBatchSize;
this.endpoint = GetFeaturesRequest.ENDPOINT;
Expand Down Expand Up @@ -272,6 +336,7 @@ public static class Builder {
private Set<MetadataOption> metadataOptionList = RequestConstants.DEFAULT_METADATA_OPTIONS;
private int microBatchSize = RequestConstants.DEFAULT_MICRO_BATCH_SIZE;
private Duration timeout = RequestConstants.NONE_TIMEOUT;
private RequestOptions requestOptions;

/** Constructs an empty Builder */
public Builder() {
Expand Down Expand Up @@ -369,6 +434,18 @@ public Builder timeout(Duration timeout) {
return this;
}

/**
* Setter for {@link RequestOptions}
*
* @param requestOptions {@link RequestOptions} object with request-level options to control
* feature server behavior
* @return this Builder
*/
public Builder requestOptions(RequestOptions requestOptions) {
this.requestOptions = requestOptions;
return this;
}

/**
* Returns an instance of {@link GetFeaturesBatchRequest} created from the fields set on this
* builder
Expand All @@ -386,7 +463,8 @@ public GetFeaturesBatchRequest build() throws TectonClientException {
requestDataList,
metadataOptionList,
microBatchSize,
timeout);
timeout,
requestOptions);
}
}

Expand Down Expand Up @@ -414,6 +492,7 @@ private static void validateParameters(
static class GetFeaturesMicroBatchRequest extends AbstractGetFeaturesRequest {

private final List<GetFeaturesRequestData> requestDataList;
private final RequestOptions requestOptions;

GetFeaturesMicroBatchRequest(
String workspaceName,
Expand All @@ -422,6 +501,18 @@ static class GetFeaturesMicroBatchRequest extends AbstractGetFeaturesRequest {
Set<MetadataOption> metadataOptions) {
super(workspaceName, featureServiceName, BATCH_ENDPOINT, metadataOptions);
this.requestDataList = requestDataList;
this.requestOptions = null;
}

GetFeaturesMicroBatchRequest(
String workspaceName,
String featureServiceName,
List<GetFeaturesRequestData> requestDataList,
Set<MetadataOption> metadataOptions,
RequestOptions requestOptions) {
super(workspaceName, featureServiceName, BATCH_ENDPOINT, metadataOptions);
this.requestDataList = requestDataList;
this.requestOptions = requestOptions;
}

// Moshi JSON classes
Expand All @@ -438,6 +529,7 @@ static class GetFeaturesBatchFields {
String workspace_name;
List<RequestDataField> request_data;
Map<String, Boolean> metadata_options;
Map<String, Object> request_options;
}

static class RequestDataField {
Expand Down Expand Up @@ -472,6 +564,9 @@ public String requestToJson() {
metadataOptions.stream()
.collect(Collectors.toMap(MetadataOption::getJsonName, (a) -> Boolean.TRUE));
}
if (requestOptions != null && !requestOptions.isEmpty()) {
getFeaturesFields.request_options = requestOptions.getOptions();
}
GetFeaturesRequestBatchJson getFeaturesRequestJson =
new GetFeaturesRequestBatchJson(getFeaturesFields);
try {
Expand Down
70 changes: 67 additions & 3 deletions src/main/java/ai/tecton/client/request/GetFeaturesRequest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ public class GetFeaturesRequest extends AbstractGetFeaturesRequest {
static final String ENDPOINT = "/api/v1/feature-service/get-features";
private final JsonAdapter<GetFeaturesRequestJson> jsonAdapter;
private final GetFeaturesRequestData getFeaturesRequestData;
private final RequestOptions requestOptions;
private final Moshi moshi = new Moshi.Builder().add(SerializeNulls.JSON_ADAPTER_FACTORY).build();

/**
Expand All @@ -40,6 +41,7 @@ public GetFeaturesRequest(
super(workspaceName, featureServiceName, ENDPOINT, RequestConstants.DEFAULT_METADATA_OPTIONS);
validateRequestParameters(getFeaturesRequestData);
this.getFeaturesRequestData = getFeaturesRequestData;
this.requestOptions = null;
jsonAdapter = moshi.adapter(GetFeaturesRequestJson.class);
}

Expand All @@ -66,13 +68,49 @@ public GetFeaturesRequest(
super(workspaceName, featureServiceName, ENDPOINT, metadataOptions);
validateRequestParameters(getFeaturesRequestData);
this.getFeaturesRequestData = getFeaturesRequestData;
this.requestOptions = null;
jsonAdapter = moshi.adapter(GetFeaturesRequestJson.class);
}

/**
* Constructor that creates a new GetFeaturesRequest with the specified parameters including
* requestOptions
*
* @param workspaceName Name of the workspace in which the Feature Service is defined
* @param featureServiceName Name of the Feature Service for which the feature vector is being
* requested
* @param getFeaturesRequestData {@link GetFeaturesRequestData} object with joinKeyMap and/or
* requestContextMap
* @param metadataOptions A {@link Set} of {@link MetadataOption} for retrieving additional
* metadata about the feature values. Use {@link RequestConstants#ALL_METADATA_OPTIONS} to
* request all metadata and {@link RequestConstants#NONE_METADATA_OPTIONS} to request no
* metadata respectively. By default, {@link RequestConstants#DEFAULT_METADATA_OPTIONS} will
* be added to each request
* @param requestOptions {@link RequestOptions} object with request-level options to control
* feature server behavior
*/
public GetFeaturesRequest(
String workspaceName,
String featureServiceName,
GetFeaturesRequestData getFeaturesRequestData,
Set<MetadataOption> metadataOptions,
RequestOptions requestOptions) {

super(workspaceName, featureServiceName, ENDPOINT, metadataOptions);
validateRequestParameters(getFeaturesRequestData);
this.getFeaturesRequestData = getFeaturesRequestData;
this.requestOptions = requestOptions;
jsonAdapter = moshi.adapter(GetFeaturesRequestJson.class);
}

GetFeaturesRequestData getFeaturesRequestData() {
return this.getFeaturesRequestData;
}

RequestOptions getRequestOptions() {
return this.requestOptions;
}

static class GetFeaturesRequestJson {
GetFeaturesFields params;

Expand All @@ -87,6 +125,7 @@ static class GetFeaturesFields {
@SerializeNulls Map<String, String> join_key_map;
@SerializeNulls Map<String, Object> request_context_map;
Map<String, Boolean> metadata_options;
Map<String, Object> request_options;
}

/**
Expand All @@ -110,6 +149,9 @@ public String requestToJson() {
metadataOptions.stream()
.collect(Collectors.toMap(MetadataOption::getJsonName, (a) -> Boolean.TRUE));
}
if (requestOptions != null && !requestOptions.isEmpty()) {
getFeaturesFields.request_options = requestOptions.getOptions();
}
GetFeaturesRequestJson getFeaturesRequestJson = new GetFeaturesRequestJson(getFeaturesFields);
try {
return jsonAdapter.toJson(getFeaturesRequestJson);
Expand All @@ -126,13 +168,14 @@ public boolean equals(Object o) {
if (o == null || getClass() != o.getClass()) return false;
if (!super.equals(o)) return false;
GetFeaturesRequest that = (GetFeaturesRequest) o;
return getFeaturesRequestData.equals(that.getFeaturesRequestData);
return getFeaturesRequestData.equals(that.getFeaturesRequestData)
&& Objects.equals(requestOptions, that.requestOptions);
}

/** Overrides <i>hashCode()</i> in class {@link Object} */
@Override
public int hashCode() {
return Objects.hash(super.hashCode(), getFeaturesRequestData);
return Objects.hash(super.hashCode(), getFeaturesRequestData, requestOptions);
}

/**
Expand All @@ -144,6 +187,7 @@ public static final class Builder {
private String workspaceName;
private String featureServiceName;
private GetFeaturesRequestData getFeaturesRequestData;
private RequestOptions requestOptions;

/** Constructor for instantiating an empty Builder */
public Builder() {
Expand Down Expand Up @@ -200,6 +244,18 @@ public Builder getFeaturesRequestData(GetFeaturesRequestData getFeaturesRequestD
return this;
}

/**
* Setter for {@link RequestOptions}
*
* @param requestOptions {@link RequestOptions} object with request-level options to control
* feature server behavior
* @return this Builder
*/
public Builder requestOptions(RequestOptions requestOptions) {
this.requestOptions = requestOptions;
return this;
}

/**
* Returns an instance of {@link GetFeaturesRequest} created from the fields set on this builder
*
Expand All @@ -208,7 +264,15 @@ public Builder getFeaturesRequestData(GetFeaturesRequestData getFeaturesRequestD
* or empty
*/
public GetFeaturesRequest build() {
if (this.metadataOptions.isEmpty()) {
if (this.requestOptions != null) {
// If requestOptions is set, use the constructor that accepts it
Set<MetadataOption> options =
this.metadataOptions.isEmpty()
? RequestConstants.DEFAULT_METADATA_OPTIONS
: this.metadataOptions;
return new GetFeaturesRequest(
workspaceName, featureServiceName, getFeaturesRequestData, options, requestOptions);
} else if (this.metadataOptions.isEmpty()) {
return new GetFeaturesRequest(workspaceName, featureServiceName, getFeaturesRequestData);
} else {
return new GetFeaturesRequest(
Expand Down
Loading