Merge pull request #5 from weiguoz/impl_interface

Implement SQLFlow client interface

Author: weiguoz

pom.xml

@@ -35,12 +35,36 @@
             <artifactId>grpc-stub</artifactId>
             <version>${grpc.version}</version>
         </dependency>
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-lang3</artifactId>
+            <version>3.9</version>
+        </dependency>
+        <!-- coding tools -->
+        <dependency>
+            <groupId>org.projectlombok</groupId>
+            <artifactId>lombok</artifactId>
+            <version>1.18.10</version>
+        </dependency>
+        <!-- unit test tools -->
         <dependency>
             <groupId>junit</groupId>
             <artifactId>junit</artifactId>
             <version>${junit.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>io.grpc</groupId>
+            <artifactId>grpc-testing</artifactId>
+            <version>${grpc.version}</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <version>3.1.0</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

src/main/java/org/sqlflow/client/SQLFlow.java

@@ -15,45 +15,48 @@
 
 package org.sqlflow.client;
 
-import java.net.ConnectException;
+import io.grpc.StatusRuntimeException;
+import org.sqlflow.client.model.RequestHeader;
 import proto.Sqlflow.JobStatus;
 
 public interface SQLFlow {
   /**
-   * Open the connection(channel) to the SQLFlow server. The serverUrl argument always ends with a
-   * port.
+   * Open a channel to the SQLFlow server. The serverUrl argument always ends with a port.
    *
    * @param serverUrl an address the SQLFlow server exposed.
    *     <p>Example: "localhost:50051"
-   * @throws ConnectException when encountering the bad network.
    */
-  void open(String serverUrl) throws ConnectException;
+  void init(String serverUrl);
 
   /**
    * Submit a task to SQLFlow server. This method return immediately.
    *
-   * @param sql: sql program. *
+   * @param header: specify datasource, user ...
+   * @param sql: sql program.
    *     <p>Example: "SELECT * FROM iris.test; SELECT * FROM iris.iris TO TRAIN DNNClassifier
    *     COLUMN..." *
    * @return return a job id for tracking.
-   * @throws Exception TODO(weiguo): more precise
+   * @throws IllegalArgumentException header or sql error
+   * @throws StatusRuntimeException
    */
-  String submit(String sql) throws Exception;
+  String submit(RequestHeader header, String sql)
+      throws IllegalArgumentException, StatusRuntimeException;
 
   /**
    * Fetch the job status by job id. The job id always returned by submit. By fetch(), we are able
    * to tracking the job status
    *
    * @param jobId specific the job we are going to track
    * @return see @code proto.JobStatus.Code
-   * @throws Exception TODO(weiguo): more precise
+   * @throws StatusRuntimeException
    */
-  JobStatus fetch(String jobId) throws Exception;
+  JobStatus fetch(String jobId) throws StatusRuntimeException;
 
   /**
-   * Close the opened connection(channel) to SQLFlow server
+   * Close the opened channel to SQLFlow server. Waits for the channel to become terminated, giving
+   * up if the timeout is reached.
    *
-   * @throws Exception TODO(weiguo): more precise
+   * @throws InterruptedException thrown by awaitTermination
    */
-  void close() throws Exception;
+  void shutdown() throws InterruptedException;
 }

src/main/java/org/sqlflow/client/impl/SQLFlowImpl.java

@@ -15,23 +15,78 @@
 
 package org.sqlflow.client.impl;
 
-import java.net.ConnectException;
+import io.grpc.ManagedChannel;
+import io.grpc.ManagedChannelBuilder;
+import io.grpc.StatusRuntimeException;
+import java.util.concurrent.TimeUnit;
+import org.apache.commons.lang3.StringUtils;
 import org.sqlflow.client.SQLFlow;
+import org.sqlflow.client.model.RequestHeader;
+import proto.SQLFlowGrpc;
+import proto.Sqlflow.Job;
 import proto.Sqlflow.JobStatus;
+import proto.Sqlflow.Request;
+import proto.Sqlflow.Session;
 
 public class SQLFlowImpl implements SQLFlow {
-  // private final ManagedChannel channel;
-  // private final SQLFlowGrpc.SQLFlowBlockingStub blockingStub;
+  private ManagedChannel channel;
+  private SQLFlowGrpc.SQLFlowBlockingStub blockingStub;
 
-  public void open(String serverUrl) throws ConnectException {}
+  public void init(String serverUrl) {
+    this.channel = ManagedChannelBuilder.forTarget(serverUrl).usePlaintext().build();
+    blockingStub = SQLFlowGrpc.newBlockingStub(channel);
+  }
 
-  public String submit(String sql) throws Exception {
-    return null;
+  public SQLFlowImpl(ManagedChannel channel) {
+    this.channel = channel;
+    blockingStub = SQLFlowGrpc.newBlockingStub(channel);
   }
 
-  public JobStatus fetch(String jobId) throws Exception {
-    return null;
+  public String submit(RequestHeader header, String sql)
+      throws IllegalArgumentException, StatusRuntimeException {
+    if (header == null || StringUtils.isAnyBlank(header.getDataSource(), header.getUserId())) {
+      throw new IllegalArgumentException("data source and userId are not allowed to be empty");
+    }
+    if (StringUtils.isBlank(sql)) {
+      throw new IllegalArgumentException("sql is empty");
+    }
+
+    Session session =
+        Session.newBuilder()
+            .setDbConnStr(header.getDataSource())
+            .setUserId(header.getUserId())
+            .setExitOnSubmit(header.isExitOnSubmit())
+            .setHiveLocation(StringUtils.defaultString(header.getHiveLocation()))
+            .setHdfsNamenodeAddr(StringUtils.defaultString(header.getHdfsNameNode()))
+            .setHdfsUser(StringUtils.defaultString(header.getHdfsUser()))
+            .setHdfsPass(StringUtils.defaultString(header.getHdfsPassword()))
+            .build();
+    Request req = Request.newBuilder().setSession(session).setSql(sql).build();
+    try {
+      Job job = blockingStub.submit(req);
+      return job.getId();
+    } catch (StatusRuntimeException e) {
+      // TODO(weiguo) logger.error
+      throw e;
+    }
   }
 
-  public void close() throws Exception {}
+  public JobStatus fetch(String jobId) throws StatusRuntimeException {
+    Job req = Job.newBuilder().setId(jobId).build();
+    try {
+      return blockingStub.fetch(req);
+    } catch (StatusRuntimeException e) {
+      // TODO(weiguo) logger.error
+      throw e;
+    }
+  }
+
+  public void shutdown() throws InterruptedException {
+    try {
+      channel.shutdown().awaitTermination(5, TimeUnit.SECONDS);
+    } catch (InterruptedException e) {
+      // TODO(weiguo) logger.error
+      throw e;
+    }
+  }
 }

src/main/java/org/sqlflow/client/model/RequestHeader.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2019 The SQLFlow Authors. All rights reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.sqlflow.client.model;
+
+import lombok.Data;
+
+@Data
+public class RequestHeader {
+  /**
+   * database source, for:
+   *
+   * <p>maxcomputer
+   * maxcompute://{accesskey_id}:{accesskey_secret}@{endpoint}?curr_project={curr_project}&scheme={scheme}
+   *
+   * <p>mysql
+   * mysql://{username}:{password}@tcp({address})/{dbname}[?param1=value1&...&paramN=valueN]
+   *
+   * <p>hive
+   * hive://user:password@ip:port/dbname[?auth=<auth_mechanism>&session.<cfg_key1>=<cfg_value1>...&session<cfg_keyN>=valueN]
+   */
+  private String dataSource;
+
+  /** user who submits the SQL task. */
+  private String userId;
+
+  /* for alps */
+  private boolean exitOnSubmit;
+
+  /* hive */
+  private String hiveLocation;
+  private String hdfsNameNode;
+  private String hdfsUser;
+  private String hdfsPassword;
+}

src/test/java/org/sqlflow/client/SQLFlowTest.java

@@ -15,9 +15,116 @@
 
 package org.sqlflow.client;
 
+import static org.junit.Assert.assertEquals;
+import static org.mockito.AdditionalAnswers.delegatesTo;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+import io.grpc.ManagedChannel;
+import io.grpc.inprocess.InProcessChannelBuilder;
+import io.grpc.inprocess.InProcessServerBuilder;
+import io.grpc.stub.StreamObserver;
+import io.grpc.testing.GrpcCleanupRule;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
+import org.mockito.ArgumentMatchers;
+import org.sqlflow.client.impl.SQLFlowImpl;
+import org.sqlflow.client.model.RequestHeader;
+import proto.SQLFlowGrpc;
+import proto.Sqlflow.Job;
+import proto.Sqlflow.JobStatus;
+import proto.Sqlflow.JobStatus.Code;
+import proto.Sqlflow.Request;
+import proto.Sqlflow.Session;
 
+@RunWith(JUnit4.class)
 public class SQLFlowTest {
+  private SQLFlow client;
+  /**
+   * This rule manages automatic graceful shutdown for the registered servers and channels at the
+   * end of test.
+   */
+  @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule();
+
+  private final SQLFlowGrpc.SQLFlowImplBase grpcService =
+      mock(
+          SQLFlowGrpc.SQLFlowImplBase.class,
+          delegatesTo(
+              new SQLFlowGrpc.SQLFlowImplBase() {
+                public void submit(Request request, StreamObserver<Job> response) {
+                  Session session = request.getSession();
+                  String userId = session.getUserId();
+                  response.onNext(
+                      Job.newBuilder().setId(mockJobId(userId, request.getSql())).build());
+                  response.onCompleted();
+                }
+
+                public void fetch(Job request, StreamObserver<JobStatus> response) {
+                  String jobId = request.getId();
+                  response.onNext(
+                      JobStatus.newBuilder()
+                          .setCodeValue(Code.PENDING_VALUE)
+                          .setMessage(mockMessage(jobId))
+                          .build());
+                  response.onCompleted();
+                }
+              }));
+
+  @Before
+  public void setUp() throws Exception {
+    String serverName = InProcessServerBuilder.generateName();
+    grpcCleanup.register(
+        InProcessServerBuilder.forName(serverName)
+            .directExecutor()
+            .addService(grpcService)
+            .build()
+            .start());
+
+    ManagedChannel channel =
+        grpcCleanup.register(InProcessChannelBuilder.forName(serverName).directExecutor().build());
+    client = new SQLFlowImpl(channel);
+  }
+
   @Test
-  public void testDummy() {}
+  public void testSubmit() {
+    String userId = "314159";
+    String sql = "SELECT * TO TRAIN DNNClassify WITH ... COLUMN ... INTO ..";
+
+    ArgumentCaptor<Request> requestCaptor = ArgumentCaptor.forClass(Request.class);
+    RequestHeader header = new RequestHeader();
+    header.setUserId(userId);
+    header.setDataSource("mysql://root@root@127.0.0.1:3306/iris");
+    String jobId = client.submit(header, sql);
+    assertEquals(mockJobId(userId, sql), jobId);
+    verify(grpcService)
+        .submit(requestCaptor.capture(), ArgumentMatchers.<StreamObserver<Job>>any());
+    assertEquals(sql, requestCaptor.getValue().getSql());
+  }
+
+  private String mockJobId(String userId, String sql) {
+    return userId + "/" + sql;
+  }
+
+  private String mockMessage(String jobId) {
+    return "Hello " + jobId;
+  }
+
+  @Test
+  public void testFetch() {
+    String jobId = "this is a job id";
+    JobStatus jobStatus = client.fetch(jobId);
+
+    assertEquals(Code.PENDING_VALUE, jobStatus.getCode().getNumber());
+    assertEquals(mockMessage(jobId), jobStatus.getMessage());
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    client.shutdown();
+  }
 }