Testing the Databricks JDBC Driver

This article describes how to test code that uses the Databricks JDBC Driver.

To test code that uses the Databricks JDBC Driver along with a collection of connection properties, you can use any test frameworks for programming languages that support JDBC. For instance, the following Java code example uses JUnit and Mockito to automate and test the Databricks JDBC Driver against a collection of connection properties. This example code is based on the example code in Authentication settings for the Databricks JDBC Driver.

The following example code file named Helpers.java contains several functions that automate the Databricks JDBC Driver against a collection of connection properties:

• The CreateConnect function uses a collection of connection properties to open a connection through a Databricks compute resource.

• The SelectNYCTaxis function uses the connection to select the specified number of data rows from the trips table in the samples catalog’s nyctaxi schema.

• The PrintResultSet function prints the data rows’ content to the screen.

// Helpers.java

import java.sql.*;
import java.util.Properties;

public class Helpers {
  static Connection CreateConnection(
    String url,
    Properties p
  ) throws SQLException {
    Connection conn = DriverManager.getConnection(url, p);
    return conn;
  }

  static ResultSet SelectNYCTaxis(
    Connection conn,
    long rows
  ) throws SQLException {
    Statement stmt = conn.createStatement();
    ResultSet rs = stmt.executeQuery("SELECT * FROM samples.nyctaxi.trips LIMIT " + rows);
    return rs;
  }

  static void PrintResultSet(ResultSet rs) throws SQLException {
    ResultSetMetaData md = rs.getMetaData();
    String[] columns = new String[md.getColumnCount()];
    for (int i = 0; i < columns.length; i++) {
      columns[i] = md.getColumnName(i + 1);
    }
    while (rs.next()) {
      System.out.print("Row " + rs.getRow() + "=[");
      for (int i = 0; i < columns.length; i++) {
        if (i != 0) {
          System.out.print(", ");
        }
        System.out.print(columns[i] + "='" + rs.getObject(i + 1) + "'");
      }
      System.out.println(")]");
    }
  }
}

The following example code file named Main.class file calls the functions in the Helpers.class file:

package org.example;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Properties;

public class Main {
  public static void main(String[] args) throws ClassNotFoundException, SQLException {
    Class.forName("com.databricks.client.jdbc.Driver");
    String url = "jdbc:databricks://" + System.getenv("DATABRICKS_SERVER_HOSTNAME") + ":443";
    Properties p = new Properties();
    p.put("httpPath", System.getenv("DATABRICKS_HTTP_PATH"));
    p.put("AuthMech", "3");
    p.put("UID", "token");
    p.put("PWD", System.getenv("DATABRICKS_TOKEN"));

    Connection conn = Helpers.CreateConnection(url, p);
    ResultSet rs = Helpers.SelectNYCTaxis(conn, 2);
    Helpers.PrintResultSet(rs);
  }
}

The following example code file named HelpersTest.class uses JUnit to test the SelectNYCTaxis function in the Helpers.class file. Instead of using the time and cost of actual compute resources to call the function in the Helpers.class file, the following example code uses Mockito to simulate this call. Simulated calls such as this are typically completed in just a few seconds, increasing your confidence in the quality of your code while not changing the state of your existing Databricks accounts or workspaces.

package org.example;

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;
import static org.junit.jupiter.api.Assertions.assertEquals;

public class HelpersTest {
  @Test
  public void testSelectNYCTaxis() throws SQLException {
    Connection mockConnection = Mockito.mock(Connection.class);
    Statement mockStatement = Mockito.mock(Statement.class);
    ResultSet mockResultSet = Mockito.mock(ResultSet.class);

    Mockito.when(mockConnection.createStatement()).thenReturn(mockStatement);
    Mockito.when(mockStatement.executeQuery(Mockito.anyString())).thenReturn(mockResultSet);

    ResultSet rs = Helpers.SelectNYCTaxis(mockConnection, 2);
    assertEquals(mockResultSet, rs);
  }
}

Because the SelectNYCTaxis function contains a SELECT statement and therefore does not change the state of the trips table, mocking is not absolutely required in this example. However, mocking enables you to quickly run your tests without waiting for an actual connection to be made with the compute resource. Also, mocking enables you to run simulated tests multiple times for functions that might change a table’s state, such as INSERT INTO, UPDATE, and DELETE FROM.