Add support for maxcompute as source type (batch)

caraml-store-protobuf/src/main/proto/feast/core/DataSource.proto

@@ -16,6 +16,7 @@ message DataSource {
     BATCH_FILE = 1;
     BATCH_BIGQUERY = 2;
     STREAM_KAFKA = 3;
+    BATCH_MAXCOMPUTE = 4;
   }
   SourceType type = 1;
 
@@ -56,6 +57,11 @@ message DataSource {
     SparkOverride spark_override = 2;
   }
 
+  // Defines options for DataSource that sources features from a MaxCompute Query
+  message MaxComputeOptions {
+    string table_ref = 1;
+  }
+
   // Defines options for DataSource that sources features from Kafka messages.
   // Each message should be a Protobuf that can be decoded with the generated
   // Java Protobuf class at the given class path
@@ -78,5 +84,6 @@ message DataSource {
     FileOptions file_options = 11;
     BigQueryOptions bigquery_options = 12;
     KafkaOptions kafka_options = 13;
+    MaxComputeOptions maxcompute_options = 14;
   }
 }

caraml-store-registry/src/main/java/dev/caraml/store/feature/DataSource.java

@@ -10,6 +10,7 @@
 import dev.caraml.store.protobuf.core.DataSourceProto.DataSource.BigQueryOptions;
 import dev.caraml.store.protobuf.core.DataSourceProto.DataSource.FileOptions;
 import dev.caraml.store.protobuf.core.DataSourceProto.DataSource.KafkaOptions;
+import dev.caraml.store.protobuf.core.DataSourceProto.DataSource.MaxComputeOptions;
 import dev.caraml.store.protobuf.core.DataSourceProto.DataSource.SourceType;
 import dev.caraml.store.protobuf.core.SparkOverrideProto.SparkOverride;
 import java.util.HashMap;
@@ -86,6 +87,9 @@ public static DataSource fromProto(DataSourceProto.DataSource spec) {
         populateDatasourceConfigMapWithSparkOverride(
             dataSourceConfigMap, spec.getBigqueryOptions().getSparkOverride());
       }
+      case BATCH_MAXCOMPUTE -> {
+        dataSourceConfigMap.put("table_ref", spec.getMaxcomputeOptions().getTableRef());
+      }
       case STREAM_KAFKA -> {
         dataSourceConfigMap.put("bootstrap_servers", spec.getKafkaOptions().getBootstrapServers());
         dataSourceConfigMap.put(
@@ -167,6 +171,11 @@ public DataSourceProto.DataSource toProto() {
             parseDatasourceConfigMapToSparkOverride(dataSourceConfigMap));
         spec.setBigqueryOptions(bigQueryOptions.build());
       }
+      case BATCH_MAXCOMPUTE -> {
+        MaxComputeOptions.Builder maxComputeOptions = MaxComputeOptions.newBuilder();
+        maxComputeOptions.setTableRef(dataSourceConfigMap.get("table_ref"));
+        spec.setMaxcomputeOptions(maxComputeOptions.build());
+      }
       case STREAM_KAFKA -> {
         KafkaOptions.Builder kafkaOptions = KafkaOptions.newBuilder();
         kafkaOptions.setBootstrapServers(dataSourceConfigMap.get("bootstrap_servers"));

caraml-store-registry/src/main/java/dev/caraml/store/feature/DataSourceValidator.java

@@ -17,6 +17,7 @@ public static void validate(DataSource spec) {
       }
       case BATCH_BIGQUERY -> Matchers.checkValidBigQueryTableRef(
           spec.getBigqueryOptions().getTableRef(), "FeatureTable");
+      case BATCH_MAXCOMPUTE -> {} // validation for table_ref to be added
       case STREAM_KAFKA -> {
         StreamFormat.FormatCase messageFormat =
             spec.getKafkaOptions().getMessageFormat().getFormatCase();

caraml-store-registry/src/main/java/dev/caraml/store/sparkjob/JobService.java

@@ -397,6 +397,7 @@ private Job createOrUpdateIngestionJob(
           case BATCH_FILE -> dataSource.getFileOptions().getSparkOverride();
           case BATCH_BIGQUERY -> dataSource.getBigqueryOptions().getSparkOverride();
           case STREAM_KAFKA -> dataSource.getKafkaOptions().getSparkOverride();
+          case BATCH_MAXCOMPUTE -> dataSource.getBigqueryOptions().getSparkOverride();
           default -> throw new IllegalArgumentException(
               String.format("%s is not a valid data source", dataSource.getType()));
         };

caraml-store-registry/src/main/java/dev/caraml/store/sparkjob/adapter/DataSourceConverter.java

@@ -52,6 +52,19 @@ static class BigQuerySource extends DataSource {
     @Nullable String datePartitionColumn;
   }
 
+  @RequiredArgsConstructor
+  @Getter
+  @Setter
+  static class MaxComputeSource extends DataSource {
+    private final String project;
+    private final String dataset;
+    private final String table;
+    private final String eventTimestampColumn;
+    private final Map<String, String> fieldMapping;
+    @Nullable String createdTimestampColumn;
+    @Nullable String datePartitionColumn;
+  }
+
   @Getter
   @Setter
   @AllArgsConstructor
@@ -117,6 +130,37 @@ public Map<String, DataSource> sourceToArgument(DataSourceProto.DataSource sourc
         yield Map.of("bq", bqSource);
       }
 
+      case MAXCOMPUTE_OPTIONS -> {
+        DataSourceProto.DataSource.MaxComputeOptions options =
+            sourceProtobuf.getMaxcomputeOptions();
+        Pattern pattern = Pattern.compile("(?<project>[^:]+):(?<dataset>[^.]+).(?<table>.+)");
+        Matcher matcher = pattern.matcher(options.getTableRef());
+        matcher.find();
+        if (!matcher.matches()) {
+          throw new IllegalArgumentException(
+              String.format(
+                  "Table ref '%s' is not in the form of <project>:<dataset>.<table>",
+                  options.getTableRef()));
+        }
+        String project = matcher.group("project");
+        String dataset = matcher.group("dataset");
+        String table = matcher.group("table");
+        MaxComputeSource maxComputeSource =
+            new MaxComputeSource(
+                project,
+                dataset,
+                table,
+                sourceProtobuf.getEventTimestampColumn(),
+                sourceProtobuf.getFieldMappingMap());
+        if (!sourceProtobuf.getDatePartitionColumn().isEmpty()) {
+          maxComputeSource.setDatePartitionColumn(sourceProtobuf.getDatePartitionColumn());
+        }
+        if (!sourceProtobuf.getCreatedTimestampColumn().isEmpty()) {
+          maxComputeSource.setCreatedTimestampColumn(sourceProtobuf.getCreatedTimestampColumn());
+        }
+        yield Map.of("maxCompute", maxComputeSource);
+      }
+
       case KAFKA_OPTIONS -> {
         DataSourceProto.DataSource.KafkaOptions options = sourceProtobuf.getKafkaOptions();
         StreamFormat messageFormat = options.getMessageFormat();

caraml-store-registry/src/main/java/dev/caraml/store/sparkjob/crd/SparkDriverSpec.java

@@ -25,4 +25,5 @@ public class SparkDriverSpec {
   private Map<String, String> annotations;
   private Map<String, String> serviceAnnotations;
   private List<SecretInfo> secrets;
+  private Map<String, Map<String, String>> envSecretKeyRefs;
 }

caraml-store-registry/src/main/java/dev/caraml/store/sparkjob/crd/SparkExecutorSpec.java

@@ -22,4 +22,5 @@ public class SparkExecutorSpec {
   private Map<String, String> labels;
   private Map<String, String> annotations;
   private List<SecretInfo> secrets;
+  private Map<String, Map<String, String>> envSecretKeyRefs;
 }