[ZEPPELIN-2679] JDBC. precode for session

docs/interpreter/jdbc.md

@@ -124,6 +124,11 @@ The JDBC interpreter properties are defined by default like below.
     <td></td>
     <td>Some SQL which executes every time after initialization of the interpreter (see <a href="../usage/interpreter/overview.html#interpreter-binding-mode">Binding mode</a>)</td>
   </tr>
+  <tr>
+    <td>default.statementPrecode</td>
+    <td></td>
+    <td>SQL code which executed before the SQL from paragraph, in the same database session (database connection)</td>
+  </tr>
   <tr>
     <td>default.completer.schemaFilters</td>
     <td></td>
@@ -306,7 +311,7 @@ Returns value of `search_path` which is set in the *default.precode*.
 
 
 ```sql
-%jdbc(mysql)
+%mysql
 select @v
 ```
 Returns value of `v` which is set in the *mysql.precode*.

jdbc/src/main/java/org/apache/zeppelin/jdbc/JDBCInterpreter.java

@@ -103,13 +103,15 @@ public class JDBCInterpreter extends KerberosInterpreter {
   static final String USER_KEY = "user";
   static final String PASSWORD_KEY = "password";
   static final String PRECODE_KEY = "precode";
+  static final String STATEMENT_PRECODE_KEY = "statementPrecode";
   static final String COMPLETER_SCHEMA_FILTERS_KEY = "completer.schemaFilters";
   static final String COMPLETER_TTL_KEY = "completer.ttlInSeconds";
   static final String DEFAULT_COMPLETER_TTL = "120";
   static final String SPLIT_QURIES_KEY = "splitQueries";
   static final String JDBC_JCEKS_FILE = "jceks.file";
   static final String JDBC_JCEKS_CREDENTIAL_KEY = "jceks.credentialKey";
   static final String PRECODE_KEY_TEMPLATE = "%s.precode";
+  static final String STATEMENT_PRECODE_KEY_TEMPLATE = "%s.statementPrecode";
   static final String DOT = ".";
 
   private static final char WHITESPACE = ' ';
@@ -125,6 +127,7 @@ public class JDBCInterpreter extends KerberosInterpreter {
   static final String DEFAULT_USER = DEFAULT_KEY + DOT + USER_KEY;
   static final String DEFAULT_PASSWORD = DEFAULT_KEY + DOT + PASSWORD_KEY;
   static final String DEFAULT_PRECODE = DEFAULT_KEY + DOT + PRECODE_KEY;
+  static final String DEFAULT_STATEMENT_PRECODE = DEFAULT_KEY + DOT + STATEMENT_PRECODE_KEY;
 
   static final String EMPTY_COLUMN_VALUE = "";
 
@@ -702,6 +705,13 @@ private InterpreterResult executeSql(String propertyKey, String sql,
         try {
           getJDBCConfiguration(user).saveStatement(paragraphId, statement);
 
+          String statementPrecode =
+              getProperty(String.format(STATEMENT_PRECODE_KEY_TEMPLATE, propertyKey));
+
+          if (StringUtils.isNotBlank(statementPrecode)) {
+            statement.execute(statementPrecode);
+          }
+
           boolean isResultSetAvailable = statement.execute(sqlToExecute);
           getJDBCConfiguration(user).setConnectionInDBDriverPoolSuccessful(propertyKey);
           if (isResultSetAvailable) {

jdbc/src/main/resources/interpreter-setting.json

@@ -48,11 +48,17 @@
       },
       "default.precode": {
         "envName": null,
-        "propertyName": "zeppelin.jdbc.precode",
+        "propertyName": "default.precode",
         "defaultValue": "",
         "description": "SQL which executes while opening connection",
         "type": "textarea"
       },
+      "default.statementPrecode": {
+        "envName": null,
+        "propertyName": "default.statementPrecode",
+        "defaultValue": "",
+        "description": "Runs before each run of the paragraph, in the same connection"
+      },
       "default.splitQueries": {
         "envName": null,
         "propertyName": "default.splitQueries",

jdbc/src/test/java/org/apache/zeppelin/jdbc/JDBCInterpreterTest.java

@@ -24,13 +24,16 @@
 import static org.apache.zeppelin.jdbc.JDBCInterpreter.COMMON_MAX_LINE;
 import static org.apache.zeppelin.jdbc.JDBCInterpreter.DEFAULT_DRIVER;
 import static org.apache.zeppelin.jdbc.JDBCInterpreter.DEFAULT_PASSWORD;
-import static org.apache.zeppelin.jdbc.JDBCInterpreter.DEFAULT_PRECODE;
-import static org.apache.zeppelin.jdbc.JDBCInterpreter.DEFAULT_URL;
+import static org.apache.zeppelin.jdbc.JDBCInterpreter.DEFAULT_STATEMENT_PRECODE;
 import static org.apache.zeppelin.jdbc.JDBCInterpreter.DEFAULT_USER;
+import static org.apache.zeppelin.jdbc.JDBCInterpreter.DEFAULT_URL;
+import static org.apache.zeppelin.jdbc.JDBCInterpreter.DEFAULT_PRECODE;
 import static org.apache.zeppelin.jdbc.JDBCInterpreter.PRECODE_KEY_TEMPLATE;
 
 import org.junit.Before;
 import org.junit.Test;
+import static org.apache.zeppelin.jdbc.JDBCInterpreter.STATEMENT_PRECODE_KEY_TEMPLATE;
+
 
 import java.io.IOException;
 import java.nio.file.Files;
@@ -538,6 +541,66 @@ public void testPrecodeWithAnotherPrefix() throws SQLException, IOException {
     assertEquals("ID\n2\n", interpreterResult.message().get(0).getData());
   }
 
+  @Test
+  public void testStatementPrecode() throws SQLException, IOException {
+    Properties properties = new Properties();
+    properties.setProperty("default.driver", "org.h2.Driver");
+    properties.setProperty("default.url", getJdbcConnection());
+    properties.setProperty("default.user", "");
+    properties.setProperty("default.password", "");
+    properties.setProperty(DEFAULT_STATEMENT_PRECODE, "set @v='statement'");
+    JDBCInterpreter jdbcInterpreter = new JDBCInterpreter(properties);
+    jdbcInterpreter.open();
+
+    String sqlQuery = "select @v";
+
+    InterpreterResult interpreterResult = jdbcInterpreter.interpret(sqlQuery, interpreterContext);
+
+    assertEquals(InterpreterResult.Code.SUCCESS, interpreterResult.code());
+    assertEquals(InterpreterResult.Type.TABLE, interpreterResult.message().get(0).getType());
+    assertEquals("@V\nstatement\n", interpreterResult.message().get(0).getData());
+  }
+
+  @Test
+  public void testIncorrectStatementPrecode() throws SQLException, IOException {
+    Properties properties = new Properties();
+    properties.setProperty("default.driver", "org.h2.Driver");
+    properties.setProperty("default.url", getJdbcConnection());
+    properties.setProperty("default.user", "");
+    properties.setProperty("default.password", "");
+    properties.setProperty(DEFAULT_STATEMENT_PRECODE, "set incorrect");
+    JDBCInterpreter jdbcInterpreter = new JDBCInterpreter(properties);
+    jdbcInterpreter.open();
+
+    String sqlQuery = "select 1";
+
+    InterpreterResult interpreterResult = jdbcInterpreter.interpret(sqlQuery, interpreterContext);
+
+    assertEquals(InterpreterResult.Code.ERROR, interpreterResult.code());
+    assertEquals(InterpreterResult.Type.TEXT, interpreterResult.message().get(0).getType());
+  }
+
+  @Test
+  public void testStatementPrecodeWithAnotherPrefix() throws SQLException, IOException {
+    Properties properties = new Properties();
+    properties.setProperty("anotherPrefix.driver", "org.h2.Driver");
+    properties.setProperty("anotherPrefix.url", getJdbcConnection());
+    properties.setProperty("anotherPrefix.user", "");
+    properties.setProperty("anotherPrefix.password", "");
+    properties.setProperty(String.format(STATEMENT_PRECODE_KEY_TEMPLATE, "anotherPrefix"),
+            "set @v='statementAnotherPrefix'");
+    JDBCInterpreter jdbcInterpreter = new JDBCInterpreter(properties);
+    jdbcInterpreter.open();
+
+    String sqlQuery = "(anotherPrefix) select @v";
+
+    InterpreterResult interpreterResult = jdbcInterpreter.interpret(sqlQuery, interpreterContext);
+
+    assertEquals(InterpreterResult.Code.SUCCESS, interpreterResult.code());
+    assertEquals(InterpreterResult.Type.TABLE, interpreterResult.message().get(0).getType());
+    assertEquals("@V\nstatementAnotherPrefix\n", interpreterResult.message().get(0).getData());
+  }
+
   @Test
   public void testSplitSqlQueryWithComments() throws SQLException, IOException {
     Properties properties = new Properties();