Merge pull request feast-dev#19 from farfetch-external/databricks_job…

…_manager Databricks job manager
Yanson · Jun 17, 2020 · b999861 · b999861
2 parents 76deb09 + 12b241b
commit b999861
Show file tree

Hide file tree

Showing 30 changed files with 1,100 additions and 106 deletions.
diff --git a/core/pom.xml b/core/pom.xml
@@ -240,6 +240,11 @@
             <artifactId>hibernate-validator-annotation-processor</artifactId>
             <version>6.1.2.Final</version>
         </dependency>
+        <dependency>
+            <groupId>dev.feast</groupId>
+            <artifactId>databricks-types</artifactId>
+            <version>${project.version}</version>
+        </dependency>
 
     </dependencies>
 </project>
diff --git a/core/src/main/java/feast/core/config/JobConfig.java b/core/src/main/java/feast/core/config/JobConfig.java
@@ -21,11 +21,14 @@
 import com.google.protobuf.util.JsonFormat;
 import feast.core.config.FeastProperties.JobProperties;
 import feast.core.job.JobManager;
+import feast.core.job.databricks.DatabricksJobManager;
 import feast.core.job.dataflow.DataflowJobManager;
 import feast.core.job.direct.DirectJobRegistry;
 import feast.core.job.direct.DirectRunnerJobManager;
+import feast.proto.core.RunnerProto.DatabricksRunnerConfigOptions;
 import feast.proto.core.RunnerProto.DataflowRunnerConfigOptions;
 import feast.proto.core.RunnerProto.DirectRunnerConfigOptions;
+import java.net.http.HttpClient;
 import java.util.Map;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -36,6 +39,7 @@
 @Slf4j
 @Configuration
 public class JobConfig {
+
   private final Gson gson = new Gson();
 
   /**
@@ -67,6 +71,13 @@ public JobManager getJobManager(FeastProperties feastProperties)
         JsonFormat.parser().merge(configJson, directRunnerConfigOptions);
         return new DirectRunnerJobManager(
             directRunnerConfigOptions.build(), new DirectJobRegistry(), metrics);
+      case DATABRICKS:
+        DatabricksRunnerConfigOptions.Builder databricksRunnerConfigOptions =
+            DatabricksRunnerConfigOptions.newBuilder();
+        JsonFormat.parser().merge(configJson, databricksRunnerConfigOptions);
+        return new DatabricksJobManager(
+            databricksRunnerConfigOptions.build(), HttpClient.newHttpClient());
+
       default:
         throw new IllegalArgumentException("Unsupported runner: " + runner);
     }

diff --git a/core/src/main/java/feast/core/job/Runner.java b/core/src/main/java/feast/core/job/Runner.java
@@ -26,6 +26,7 @@
 public enum Runner {
   DATAFLOW("DataflowRunner"),
   FLINK("FlinkRunner"),
+  DATABRICKS("DatabricksRunner"),
   DIRECT("DirectRunner");
 
   private final String humanName;