测试 Databricks JDBC 驱动程序 (Simba)

注释

此页面适用于低于版本 3 的 Databricks JDBC 驱动程序版本。 有关版本 3 及更高版本,请参阅 Databricks JDBC 驱动程序

本页介绍如何测试使用 Databricks JDBC 驱动程序的代码。

将任何测试框架用于 JDBC 兼容的语言。 以下示例使用 JUnitMockito 测试 JDBC 驱动程序连接。 此代码基于 Databricks JDBC 驱动程序(Simba)的身份验证设置中的示例。

帮助函数

该文件 Helpers.java 包含用于处理 JDBC 连接的实用工具函数:

  • CreateConnection:打开与 Azure Databricks 计算资源的连接。
  • SelectNYCTaxis:从samples.nyctaxi.trips查询指定的行数。
  • PrintResultSet:将结果集内容打印到控制台。
// 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(")]");
    }
  }
}

Main 类

该文件 Main.java 调用帮助程序函数来连接和查询数据:

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);
  }
}

使用模拟进行单元测试

该文件 HelpersTest.java 使用 JUnit 和 Mockito 测试 SelectNYCTaxis 函数。 通过模拟(Mocking)来仿真数据库连接,无需使用实际的计算资源,因此测试可以在数秒内完成,而不会对您的 Azure Databricks 工作区造成影响。

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);
  }
}

由于 SelectNYCTaxis 只运行 SELECT 语句,因此这里不需要严格的模拟。 但是,模拟在测试修改数据的函数 (INSERT INTOUPDATEDELETE FROM) 时特别有用,因为可以重复运行测试,而不会影响表状态。