diff --git a/database-commons/src/main/java/io/cdap/plugin/db/connector/AbstractDBSpecificConnector.java b/database-commons/src/main/java/io/cdap/plugin/db/connector/AbstractDBSpecificConnector.java index 0308cf7a4..e0407203b 100644 --- a/database-commons/src/main/java/io/cdap/plugin/db/connector/AbstractDBSpecificConnector.java +++ b/database-commons/src/main/java/io/cdap/plugin/db/connector/AbstractDBSpecificConnector.java @@ -104,8 +104,12 @@ public InputFormatProvider getInputFormatProvider(ConnectorContext context, Samp String tableQuery = getTableQuery(path.getDatabase(), path.getSchema(), path.getTable(), request.getLimit(), request.getProperties().get("sampleType"), request.getProperties().get("strata"), sessionID); DataDrivenETLDBInputFormat.setInput(connectionConfigAccessor.getConfiguration(), getDBRecordType(), - tableQuery, null, false); + tableQuery, null, isAutoCommitEnabled()); connectionConfigAccessor.setConnectionArguments(Maps.fromProperties(config.getConnectionArgumentsProperties())); + String isolationLevel = getTransactionIsolationLevel(); + if (isolationLevel != null) { + connectionConfigAccessor.setTransactionIsolationLevel(isolationLevel); + } connectionConfigAccessor.getConfiguration().setInt(MRJobConfig.NUM_MAPS, 1); Map additionalArguments = config.getAdditionalArguments(); for (Map.Entry argument : additionalArguments.entrySet()) { @@ -221,4 +225,19 @@ protected Schema getTableSchema(Connection connection, String database, protected String generateSessionID() { return UUID.randomUUID().toString().replace('-', '_'); } + + /** + * Returns whether auto-commit should be enabled for this connector. + * By default, it is false. + */ + protected boolean isAutoCommitEnabled() { + return false; + } + /** + * Returns the default transaction isolation level for this connector. + * If null, it falls back to the database driver's default or serializable. + */ + protected String getTransactionIsolationLevel() { + return null; + } } diff --git a/database-commons/src/test/java/io/cdap/plugin/db/source/DataDrivenETLDBInputFormatTest.java b/database-commons/src/test/java/io/cdap/plugin/db/source/DataDrivenETLDBInputFormatTest.java index b369d008b..a07d0dee5 100644 --- a/database-commons/src/test/java/io/cdap/plugin/db/source/DataDrivenETLDBInputFormatTest.java +++ b/database-commons/src/test/java/io/cdap/plugin/db/source/DataDrivenETLDBInputFormatTest.java @@ -17,9 +17,24 @@ package io.cdap.plugin.db.source; import com.google.common.collect.ImmutableList; +import io.cdap.cdap.api.data.batch.InputFormatProvider; +import io.cdap.cdap.api.data.format.StructuredRecord; +import io.cdap.cdap.api.data.schema.Schema; +import io.cdap.cdap.etl.api.connector.ConnectorContext; +import io.cdap.cdap.etl.api.connector.SampleRequest; +import io.cdap.cdap.etl.mock.common.MockConnectorConfigurer; +import io.cdap.cdap.etl.mock.common.MockConnectorContext; +import io.cdap.plugin.common.db.DBConnectorPath; +import io.cdap.plugin.db.ConnectionConfigAccessor; +import io.cdap.plugin.db.DBRecord; +import io.cdap.plugin.db.TransactionIsolationLevel; +import io.cdap.plugin.db.connector.AbstractDBConnectorConfig; +import io.cdap.plugin.db.connector.AbstractDBSpecificConnector; +import org.apache.hadoop.io.LongWritable; import org.apache.hadoop.mapreduce.InputSplit; import org.apache.hadoop.mapreduce.JobContext; import org.apache.hadoop.mapreduce.lib.db.DBConfiguration; +import org.apache.hadoop.mapreduce.lib.db.DBWritable; import org.apache.hadoop.mapreduce.lib.db.DataDrivenDBInputFormat; import org.junit.Assert; import org.junit.Before; @@ -30,8 +45,11 @@ import org.mockito.runners.MockitoJUnitRunner; import java.io.IOException; +import java.sql.Connection; +import java.sql.Driver; import java.util.Collections; import java.util.List; +import java.util.Map; @RunWith(MockitoJUnitRunner.class) public class DataDrivenETLDBInputFormatTest { @@ -118,4 +136,92 @@ public void testGetSplitsDoesNotAddNullSplitIfBaseReturnsEmptyList() throws IOEx (DataDrivenDBInputFormat.DataDrivenDBInputSplit) finalSplits.get(0); Assert.assertEquals("1=1", split.getLowerClause()); } + + @Test + public void testDefaultConnectorInputFormatConfiguration() throws IOException { + TestDBConnector connector = new TestDBConnector(new TestDBConnectorConfig()); + + Assert.assertFalse(connector.getBaseAutoCommitEnabled()); + Assert.assertNull(connector.getBaseTransactionIsolationLevel()); + + ConnectorContext context = new MockConnectorContext(new MockConnectorConfigurer()); + SampleRequest sampleRequest = SampleRequest.builder(10).setPath("db/table").build(); + InputFormatProvider provider = connector.getInputFormatProvider(context, sampleRequest); + Map conf = provider.getInputFormatConfiguration(); + + Assert.assertEquals("false", conf.get(ConnectionConfigAccessor.AUTO_COMMIT_ENABLED)); + Assert.assertNull(conf.get(TransactionIsolationLevel.CONF_KEY)); + } + + @Test + public void testOverridingConnectorInputFormatConfiguration() throws IOException { + TestDBConnector connector = new TestDBConnector(new TestDBConnectorConfig()) { + @Override + protected boolean isAutoCommitEnabled() { + return true; + } + + @Override + protected String getTransactionIsolationLevel() { + return TransactionIsolationLevel.Level.TRANSACTION_READ_UNCOMMITTED.name(); + } + }; + + ConnectorContext context = new MockConnectorContext(new MockConnectorConfigurer()); + SampleRequest sampleRequest = SampleRequest.builder(10).setPath("db/table").build(); + InputFormatProvider provider = connector.getInputFormatProvider(context, sampleRequest); + Map conf = provider.getInputFormatConfiguration(); + + Assert.assertEquals("true", conf.get(ConnectionConfigAccessor.AUTO_COMMIT_ENABLED)); + Assert.assertEquals(TransactionIsolationLevel.Level.TRANSACTION_READ_UNCOMMITTED.name(), + conf.get(TransactionIsolationLevel.CONF_KEY)); + } + + private static class TestDBConnectorConfig extends AbstractDBConnectorConfig { + @Override + public String getConnectionString() { + return "jdbc:test://localhost:1234/db"; + } + } + + private static class TestDBConnector extends AbstractDBSpecificConnector { + TestDBConnector(AbstractDBConnectorConfig config) { + super(config); + this.driverClass = Driver.class; + } + + @Override + public boolean supportSchema() { + return false; + } + + @Override + protected Class getDBRecordType() { + return DBRecord.class; + } + + @Override + public StructuredRecord transform(LongWritable key, DBRecord val) { + return null; + } + + @Override + protected Connection getConnection(DBConnectorPath path) { + return null; + } + + @Override + protected Schema loadTableSchema(Connection connection, String query, + Integer timeoutSec, String sessionID) { + return Schema.recordOf("outputSchema", Schema.Field.of("id", Schema.of(Schema.Type.INT))); + } + + boolean getBaseAutoCommitEnabled() { + return isAutoCommitEnabled(); + } + + String getBaseTransactionIsolationLevel() { + return getTransactionIsolationLevel(); + } + } } diff --git a/databricks-plugin/docs/Databricks-batchsource.md b/databricks-plugin/docs/Databricks-batchsource.md new file mode 100644 index 000000000..f259e6e9d --- /dev/null +++ b/databricks-plugin/docs/Databricks-batchsource.md @@ -0,0 +1,15 @@ +# Databricks Batch Source + +Description +----------- +Reads data from a Databricks table using a configurable SQL query. + +Properties +---------- +* **Use Connection**: Whether to use an existing Databricks connection. +* **Host**: Server Hostname of the Databricks cluster or SQL warehouse. +* **Port**: Database port (default is 443). +* **HTTP Path**: The HTTP Path for the Databricks cluster or SQL warehouse. +* **Reference Name**: Name used to identify this source for lineage. +* **Database / Catalog**: Optional catalog or database name. +* **Import Query**: SQL query to execute against Databricks. diff --git a/databricks-plugin/docs/Databricks-connector.md b/databricks-plugin/docs/Databricks-connector.md new file mode 100644 index 000000000..73d7ca1fd --- /dev/null +++ b/databricks-plugin/docs/Databricks-connector.md @@ -0,0 +1,15 @@ +# Databricks Database Connector + +Description +----------- +Connects to Databricks database / Lakehouse via JDBC. + +Properties +---------- +* **Host**: Server Hostname of the Databricks cluster or SQL warehouse. +* **Port**: Database port (default is 443). +* **HTTP Path**: The HTTP Path for the Databricks cluster or SQL warehouse. +* **Database / Catalog**: Optional catalog or database name to connect to. +* **Username**: Username / token user. +* **Password / Token**: Personal Access Token (PAT) or password. +* **Connection Arguments**: Arbitrary key-value pairs to pass as connection arguments to the JDBC driver (e.g. `AuthMech=11;Auth_Flow=2`). diff --git a/databricks-plugin/icons/Databricks-batchsource.png b/databricks-plugin/icons/Databricks-batchsource.png new file mode 100644 index 000000000..e27f31a4c Binary files /dev/null and b/databricks-plugin/icons/Databricks-batchsource.png differ diff --git a/databricks-plugin/pom.xml b/databricks-plugin/pom.xml new file mode 100644 index 000000000..4d66f360b --- /dev/null +++ b/databricks-plugin/pom.xml @@ -0,0 +1,127 @@ + + + + + database-plugins-parent + io.cdap.plugin + 1.13.0-SNAPSHOT + + + Databricks plugin + databricks-plugin + 4.0.0 + + + 3.4.3 + + + + + io.cdap.cdap + cdap-etl-api + + + io.cdap.plugin + database-commons + ${project.version} + + + io.cdap.plugin + hydrator-common + + + com.google.guava + guava + + + + + com.databricks + databricks-jdbc + ${databricks-jdbc.version} + test + + + io.cdap.plugin + database-commons + ${project.version} + test-jar + test + + + io.cdap.cdap + hydrator-test + + + io.cdap.cdap + cdap-data-pipeline3_2.12 + + + junit + junit + + + org.mockito + mockito-core + test + + + io.cdap.cdap + cdap-api + provided + + + + + + + io.cdap + cdap-maven-plugin + + + org.apache.felix + maven-bundle-plugin + 5.1.2 + true + + + <_exportcontents> + io.cdap.plugin.databricks.*; + io.cdap.plugin.db.source.*; + org.apache.commons.lang; + org.apache.commons.logging.*; + org.codehaus.jackson.* + + *;inline=false;scope=compile + true + lib + + + + + package + + bundle + + + + + + + diff --git a/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnector.java b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnector.java new file mode 100644 index 000000000..cb5774df3 --- /dev/null +++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnector.java @@ -0,0 +1,184 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import io.cdap.cdap.api.annotation.Category; +import io.cdap.cdap.api.annotation.Description; +import io.cdap.cdap.api.annotation.Name; +import io.cdap.cdap.api.annotation.Plugin; +import io.cdap.cdap.api.data.format.StructuredRecord; +import io.cdap.cdap.etl.api.batch.BatchSource; +import io.cdap.cdap.etl.api.connector.Connector; +import io.cdap.cdap.etl.api.connector.ConnectorSpec; +import io.cdap.cdap.etl.api.connector.ConnectorSpecRequest; +import io.cdap.cdap.etl.api.connector.PluginSpec; +import io.cdap.cdap.etl.api.connector.SampleType; +import io.cdap.plugin.common.Constants; +import io.cdap.plugin.common.ReferenceNames; +import io.cdap.plugin.common.db.DBConnectorPath; +import io.cdap.plugin.db.NoOpCommitConnection; +import io.cdap.plugin.db.SchemaReader; +import io.cdap.plugin.db.TransactionIsolationLevel; +import io.cdap.plugin.db.connector.AbstractDBSpecificConnector; +import io.cdap.plugin.db.connector.DBSpecificPath; +import org.apache.hadoop.io.LongWritable; +import org.apache.hadoop.mapreduce.lib.db.DBWritable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.sql.Connection; +import java.sql.SQLException; +import java.util.HashMap; +import java.util.Map; + +/** + * Databricks Database Connector that connects to Databricks database via JDBC. + */ +@Plugin(type = Connector.PLUGIN_TYPE) +@Name(DatabricksConstants.PLUGIN_NAME) +@Description("Connection to access data in Databricks using JDBC.") +@Category("Database") +public class DatabricksConnector extends AbstractDBSpecificConnector { + public static final String NAME = DatabricksConstants.PLUGIN_NAME; + private final DatabricksConnectorConfig config; + + private static final Logger LOG = LoggerFactory.getLogger(DatabricksConnector.class); + + public DatabricksConnector(DatabricksConnectorConfig config) { + super(config); + this.config = config; + } + + @Override + protected DBConnectorPath getDBConnectorPath(String path) throws IOException { + return DBSpecificPath.of(path, supportSchema()); + } + + @Override + protected Connection getConnection(DBConnectorPath path) { + Connection connection = super.getConnection(path); + try { + connection.setTransactionIsolation(Connection.TRANSACTION_READ_UNCOMMITTED); + } catch (SQLException e) { + LOG.warn("Failed to set transaction isolation level to READ_UNCOMMITTED", e); + } + return new NoOpCommitConnection(connection); + } + + @Override + protected Connection getConnection() { + Connection connection = super.getConnection(); + try { + connection.setTransactionIsolation(Connection.TRANSACTION_READ_UNCOMMITTED); + } catch (SQLException e) { + LOG.warn("Failed to set transaction isolation level to READ_UNCOMMITTED", e); + } + return new NoOpCommitConnection(connection); + } + + @Override + public boolean supportSchema() { + return true; + } + + @Override + protected Class getDBRecordType() { + return DatabricksDBRecord.class; + } + + @Override + public StructuredRecord transform(LongWritable longWritable, DatabricksDBRecord record) { + return record.getRecord(); + } + + @Override + protected SchemaReader getSchemaReader(String sessionID) { + return new DatabricksSchemaReader(sessionID); + } + + @Override + protected String getTableName(String database, String schema, String table) { + if (database == null && schema == null) { + return String.format("`%s`", table); + } + if (database == null) { + return String.format("`%s`.`%s`", schema, table); + } + if (schema == null) { + return String.format("`%s`.`%s`", database, table); + } + return String.format("`%s`.`%s`.`%s`", database, schema, table); + } + + @Override + protected String getRandomQuery(String tableName, int limit) { + return String.format("SELECT * FROM %s\n" + + "WHERE rand() < %d.0 / (SELECT COUNT(*) FROM %s)", + tableName, limit, tableName); + } + + @Override + protected String getStratifiedQuery(String tableName, int limit, String strata, String sessionID) { + return String.format("WITH t_%s AS (\n" + + " SELECT *,\n" + + " ROW_NUMBER() OVER (ORDER BY %s, RAND()) AS sqn_%s,\n" + + " COUNT(*) OVER () AS c_%s\n" + + " FROM %s\n" + + " )\n" + + "SELECT * FROM t_%s\n" + + "WHERE MOD(sqn_%s, GREATEST(1, CAST(c_%s / %d AS BIGINT))) = 1\n" + + "ORDER BY %s\n" + + "LIMIT %d", + sessionID, strata, sessionID, sessionID, tableName, sessionID, sessionID, sessionID, + limit, strata, limit); + } + + @Override + protected void setConnectorSpec(ConnectorSpecRequest request, DBConnectorPath path, + ConnectorSpec.Builder builder) { + Map sourceProperties = new HashMap<>(); + setConnectionProperties(sourceProperties, request); + builder + .addRelatedPlugin(new PluginSpec(DatabricksConstants.PLUGIN_NAME, + BatchSource.PLUGIN_TYPE, sourceProperties)) + .addSupportedSampleType(SampleType.RANDOM) + .addSupportedSampleType(SampleType.STRATIFIED); + + String schema = path.getSchema(); + sourceProperties.put(DatabricksSource.DatabricksSourceConfig.NUM_SPLITS, "1"); + sourceProperties.put(DatabricksSource.DatabricksSourceConfig.FETCH_SIZE, + DatabricksSource.DatabricksSourceConfig.DEFAULT_FETCH_SIZE); + String table = path.getTable(); + if (table == null) { + return; + } + sourceProperties.put(DatabricksSource.DatabricksSourceConfig.IMPORT_QUERY, + getTableQuery(path.getDatabase(), schema, table)); + sourceProperties.put(Constants.Reference.REFERENCE_NAME, ReferenceNames.cleanseReferenceName(table)); + } + + @Override + protected boolean isAutoCommitEnabled() { + return true; + } + + @Override + protected String getTransactionIsolationLevel() { + return TransactionIsolationLevel.Level.TRANSACTION_REPEATABLE_READ.name(); + } +} diff --git a/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnectorConfig.java b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnectorConfig.java new file mode 100644 index 000000000..628fd7586 --- /dev/null +++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnectorConfig.java @@ -0,0 +1,130 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import com.google.common.base.Strings; +import io.cdap.cdap.api.annotation.Description; +import io.cdap.cdap.api.annotation.Macro; +import io.cdap.cdap.api.annotation.Name; +import io.cdap.plugin.db.ConnectionConfig; +import io.cdap.plugin.db.connector.AbstractDBConnectorConfig; + +import java.util.Properties; +import javax.annotation.Nullable; + +/** + * Configuration for Databricks connector + */ +public class DatabricksConnectorConfig extends AbstractDBConnectorConfig { + + public static final String HTTP_PATH = "httpPath"; + + @Name(ConnectionConfig.HOST) + @Description("The server hostname of the Databricks cluster or SQL warehouse.") + @Macro + private String host; + + @Name(ConnectionConfig.PORT) + @Description("Database port number. Default is 443.") + @Macro + @Nullable + private Integer port; + + @Name(HTTP_PATH) + @Description("The HTTP Path for the Databricks cluster or SQL warehouse.") + @Macro + private String httpPath; + + @Name(ConnectionConfig.DATABASE) + @Description("Database or Catalog name to connect to.") + @Macro + @Nullable + private String database; + + public DatabricksConnectorConfig(@Nullable @Name(ConnectionConfig.USER) String user, + @Nullable @Name(ConnectionConfig.PASSWORD) String password, + @Name(ConnectionConfig.JDBC_PLUGIN_NAME) String jdbcPluginName, + @Nullable @Name(ConnectionConfig.CONNECTION_ARGUMENTS) String connectionArguments, + @Name(ConnectionConfig.HOST) String host, + @Name(DatabricksConnectorConfig.HTTP_PATH) String httpPath, + @Nullable @Name(ConnectionConfig.DATABASE) String database, + @Nullable @Name(ConnectionConfig.PORT) Integer port) { + this.user = user; + this.password = password; + this.jdbcPluginName = jdbcPluginName; + this.connectionArguments = connectionArguments; + this.host = host; + this.httpPath = httpPath; + this.database = database; + this.port = port; + } + + @Nullable + @Override + public String getUser() { + if (Strings.isNullOrEmpty(user) && !Strings.isNullOrEmpty(password)) { + return "token"; + } + return user; + } + + @Override + public Properties getConnectionArgumentsProperties() { + return getConnectionArgumentsProperties(connectionArguments, getUser(), getPassword()); + } + + @Nullable + public String getDatabase() { + return database; + } + + public String getHost() { + return host; + } + + public int getPort() { + return port == null ? 443 : port; + } + + public String getHttpPath() { + return httpPath; + } + + @Override + public String getConnectionString() { + if (database != null && !database.trim().isEmpty()) { + return String.format( + DatabricksConstants.DATABRICKS_DB_CONNECTION_STRING_FORMAT, + host, + getPort(), + database, + httpPath); + } + return String.format( + DatabricksConstants.DATABRICKS_CONNECTION_STRING_FORMAT, + host, + getPort(), + httpPath); + } + + @Override + public boolean canConnect() { + return super.canConnect() && !containsMacro(ConnectionConfig.HOST) && + !containsMacro(ConnectionConfig.PORT) && !containsMacro(HTTP_PATH) && + !containsMacro(ConnectionConfig.DATABASE); + } +} diff --git a/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConstants.java b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConstants.java new file mode 100644 index 000000000..1e8a1b35c --- /dev/null +++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConstants.java @@ -0,0 +1,30 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +/** Databricks constants. */ +public final class DatabricksConstants { + + private DatabricksConstants() { + } + + public static final String PLUGIN_NAME = "Databricks"; + public static final String DATABRICKS_CONNECTION_STRING_FORMAT = + "jdbc:databricks://%s:%d;HttpPath=%s;"; + public static final String DATABRICKS_DB_CONNECTION_STRING_FORMAT = + "jdbc:databricks://%s:%d;ConnCatalog=%s;HttpPath=%s;"; +} diff --git a/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksDBRecord.java b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksDBRecord.java new file mode 100644 index 000000000..d1c5f6ab6 --- /dev/null +++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksDBRecord.java @@ -0,0 +1,80 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import io.cdap.cdap.api.data.format.StructuredRecord; +import io.cdap.cdap.api.data.schema.Schema; +import io.cdap.plugin.db.DBRecord; +import io.cdap.plugin.db.SchemaReader; + +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Timestamp; +import java.sql.Types; + +/** + * Writable class for Databricks Source + */ +public class DatabricksDBRecord extends DBRecord { + + /** + * Used in map-reduce. Do not remove. + */ + @SuppressWarnings("unused") + public DatabricksDBRecord() { + } + + @Override + protected SchemaReader getSchemaReader() { + return new DatabricksSchemaReader(); + } + + @Override + protected void handleField(ResultSet resultSet, StructuredRecord.Builder recordBuilder, Schema.Field field, + int columnIndex, int sqlType, int sqlPrecision, int sqlScale) throws SQLException { + ResultSetMetaData metadata = resultSet.getMetaData(); + String columnTypeName = metadata.getColumnTypeName(columnIndex); + String normalizedType = columnTypeName != null ? columnTypeName.trim().toUpperCase() : null; + + if (sqlType == Types.NULL || "VOID".equals(normalizedType) || "NULL".equals(normalizedType)) { + recordBuilder.set(field.getName(), null); + return; + } + + if (normalizedType != null && (normalizedType.equals("VARIANT") || + normalizedType.startsWith("ARRAY") || normalizedType.startsWith("MAP") || + normalizedType.startsWith("STRUCT") || normalizedType.equals("OBJECT") || + normalizedType.equals("FILE") || normalizedType.startsWith("INTERVAL") || + normalizedType.startsWith("GEOGRAPHY") || normalizedType.startsWith("GEOMETRY"))) { + Object value = resultSet.getObject(columnIndex); + recordBuilder.set(field.getName(), value != null ? value.toString() : null); + return; + } + + Schema nonNullableSchema = field.getSchema().isNullable() ? + field.getSchema().getNonNullable() : field.getSchema(); + if (Schema.LogicalType.DATETIME.equals(nonNullableSchema.getLogicalType()) || + "TIMESTAMP_NTZ".equals(normalizedType)) { + Timestamp timestamp = resultSet.getTimestamp(columnIndex); + recordBuilder.setDateTime(field.getName(), timestamp != null ? timestamp.toLocalDateTime() : null); + return; + } + + setField(resultSet, recordBuilder, field, columnIndex, sqlType, sqlPrecision, sqlScale); + } +} diff --git a/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksSchemaReader.java b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksSchemaReader.java new file mode 100644 index 000000000..856f58ebe --- /dev/null +++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksSchemaReader.java @@ -0,0 +1,71 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import io.cdap.cdap.api.data.schema.Schema; +import io.cdap.plugin.db.CommonSchemaReader; + +import java.sql.ResultSetMetaData; +import java.sql.SQLException; + +/** + * Databricks Schema Reader class + */ +public class DatabricksSchemaReader extends CommonSchemaReader { + + private final String sessionID; + + public DatabricksSchemaReader() { + this(null); + } + + public DatabricksSchemaReader(String sessionID) { + super(); + this.sessionID = sessionID; + } + + @Override + public Schema getSchema(ResultSetMetaData metadata, int index) throws SQLException { + String typeName = metadata.getColumnTypeName(index); + + if (typeName != null) { + String normalizedType = typeName.trim().toUpperCase(); + if (normalizedType.equals("TIMESTAMP_NTZ")) { + return Schema.of(Schema.LogicalType.DATETIME); + } + if (normalizedType.equals("VARIANT") || normalizedType.startsWith("ARRAY") || + normalizedType.startsWith("MAP") || normalizedType.startsWith("STRUCT") || + normalizedType.equals("OBJECT") || normalizedType.equals("FILE") || + normalizedType.equals("VOID") || normalizedType.equals("NULL") || + normalizedType.startsWith("INTERVAL") || normalizedType.startsWith("GEOGRAPHY") || + normalizedType.startsWith("GEOMETRY")) { + return Schema.of(Schema.Type.STRING); + } + } + + return super.getSchema(metadata, index); + } + + @Override + public boolean shouldIgnoreColumn(ResultSetMetaData metadata, int index) throws SQLException { + if (sessionID == null) { + return false; + } + String columnName = metadata.getColumnName(index); + return ("c_" + sessionID).equals(columnName) || ("sqn_" + sessionID).equals(columnName); + } +} diff --git a/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksSource.java b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksSource.java new file mode 100644 index 000000000..7ccb4cb38 --- /dev/null +++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksSource.java @@ -0,0 +1,153 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import com.google.common.annotations.VisibleForTesting; +import io.cdap.cdap.api.annotation.Description; +import io.cdap.cdap.api.annotation.Macro; +import io.cdap.cdap.api.annotation.Metadata; +import io.cdap.cdap.api.annotation.MetadataProperty; +import io.cdap.cdap.api.annotation.Name; +import io.cdap.cdap.api.annotation.Plugin; +import io.cdap.cdap.api.data.schema.Schema; +import io.cdap.cdap.etl.api.FailureCollector; +import io.cdap.cdap.etl.api.batch.BatchSource; +import io.cdap.cdap.etl.api.batch.BatchSourceContext; +import io.cdap.cdap.etl.api.connector.Connector; +import io.cdap.plugin.common.Asset; +import io.cdap.plugin.common.ConfigUtil; +import io.cdap.plugin.common.LineageRecorder; +import io.cdap.plugin.db.ConnectionConfigAccessor; +import io.cdap.plugin.db.SchemaReader; +import io.cdap.plugin.db.TransactionIsolationLevel; +import io.cdap.plugin.db.config.AbstractDBSpecificSourceConfig; +import io.cdap.plugin.db.source.AbstractDBSource; +import io.cdap.plugin.util.DBUtils; +import org.apache.hadoop.mapreduce.lib.db.DBWritable; + +import java.io.IOException; +import java.util.Collections; +import java.util.Map; +import javax.annotation.Nullable; + +/** + * Batch source to read from a Databricks database. + */ +@Plugin(type = BatchSource.PLUGIN_TYPE) +@Name(DatabricksConstants.PLUGIN_NAME) +@Description( + "Reads from a Databricks table using a configurable SQL query." + + " Outputs one record for each row returned by the query.") +@Metadata(properties = {@MetadataProperty(key = Connector.PLUGIN_TYPE, value = DatabricksConnector.NAME)}) +public class DatabricksSource extends AbstractDBSource { + + private final DatabricksSourceConfig databricksSourceConfig; + + public DatabricksSource(DatabricksSourceConfig databricksSourceConfig) { + super(databricksSourceConfig); + this.databricksSourceConfig = databricksSourceConfig; + } + + @Override + protected SchemaReader getSchemaReader() { + return new DatabricksSchemaReader(); + } + + @Override + protected Class getDBRecordType() { + return DatabricksDBRecord.class; + } + + @Override + protected String createConnectionString() { + DatabricksConnectorConfig connection = databricksSourceConfig.getConnection(); + return connection == null ? null : connection.getConnectionString(); + } + + @Override + protected LineageRecorder getLineageRecorder(BatchSourceContext context) { + DatabricksConnectorConfig connection = databricksSourceConfig.getConnection(); + String host = connection == null ? null : connection.getHost(); + int port = connection == null ? 443 : connection.getPort(); + String database = connection == null ? null : connection.getDatabase(); + String fqn = DBUtils.constructFQN("databricks", host, port, database, + databricksSourceConfig.getReferenceName()); + Asset.Builder assetBuilder = Asset.builder(databricksSourceConfig.getReferenceName()).setFqn(fqn); + return new LineageRecorder(context, assetBuilder.build()); + } + + @Override + public ConnectionConfigAccessor getConnectionConfigAccessor(String driverClassName, + Schema schemaFromDB, + FailureCollector collector) throws IOException { + ConnectionConfigAccessor configAccessor = + super.getConnectionConfigAccessor(driverClassName, schemaFromDB, collector); + configAccessor.setAutoCommitEnabled(true); + return configAccessor; + } + + /** + * Databricks source config. + */ + public static class DatabricksSourceConfig extends AbstractDBSpecificSourceConfig { + + @Name(ConfigUtil.NAME_USE_CONNECTION) + @Nullable + @Description("Whether to use an existing connection.") + private Boolean useConnection; + + @Name(ConfigUtil.NAME_CONNECTION) + @Macro + @Nullable + @Description("The existing connection to use.") + private DatabricksConnectorConfig connection; + + @Override + public Map getDBSpecificArguments() { + return Collections.emptyMap(); + } + + @VisibleForTesting + public DatabricksSourceConfig(@Nullable Boolean useConnection, + @Nullable DatabricksConnectorConfig connection) { + this.useConnection = useConnection; + this.connection = connection; + } + + @Override + public String getTransactionIsolationLevel() { + return TransactionIsolationLevel.Level.TRANSACTION_REPEATABLE_READ.name(); + } + + @Override + public Integer getFetchSize() { + Integer fetchSize = super.getFetchSize(); + return fetchSize == null ? Integer.parseInt(DEFAULT_FETCH_SIZE) : fetchSize; + } + + @Override + protected DatabricksConnectorConfig getConnection() { + return connection; + } + + @Override + public void validate(FailureCollector collector) { + ConfigUtil.validateConnection(this, useConnection, connection, collector); + super.validate(collector); + } + } +} diff --git a/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorTest.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorTest.java new file mode 100644 index 000000000..e8fe79122 --- /dev/null +++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorTest.java @@ -0,0 +1,39 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import io.cdap.plugin.db.connector.DBSpecificConnectorBaseTest; +import org.junit.Test; + +import java.io.IOException; + +/** + * Integration tests for {@link DatabricksConnector}. + */ +public class DatabricksConnectorTest extends DBSpecificConnectorBaseTest { + + private static final String JDBC_DRIVER_CLASS_NAME = "com.databricks.client.jdbc.Driver"; + + @Test + public void test() throws IOException, ClassNotFoundException, InstantiationException, IllegalAccessException { + String httpPath = System.getProperty("http.path", "sql/1.0/warehouses/test"); + test(new DatabricksConnector( + new DatabricksConnectorConfig(username, password, JDBC_PLUGIN_NAME, connectionArguments, + host, httpPath, database, port)), + JDBC_DRIVER_CLASS_NAME, DatabricksConstants.PLUGIN_NAME); + } +} diff --git a/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorUnitTest.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorUnitTest.java new file mode 100644 index 000000000..99b8457fe --- /dev/null +++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorUnitTest.java @@ -0,0 +1,86 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import org.junit.Assert; +import org.junit.Test; + +/** + * Unit tests for {@link DatabricksConnector} + */ +public class DatabricksConnectorUnitTest { + + private static final DatabricksConnector CONNECTOR = new DatabricksConnector(new DatabricksConnectorConfig( + "token", "password", "jdbc", "", "dbc-xxx.cloud.databricks.com", + "sql/1.0/warehouses/xxx", "main", 443)); + + @Test + public void testGetTableName() { + Assert.assertEquals("`main`.`default`.`my_table`", + CONNECTOR.getTableName("main", "default", "my_table")); + Assert.assertEquals("`default`.`my_table`", + CONNECTOR.getTableName(null, "default", "my_table")); + Assert.assertEquals("`my_table`", + CONNECTOR.getTableName(null, null, "my_table")); + } + + @Test + public void testGetRandomQuery() { + Assert.assertEquals("SELECT * FROM `main`.`default`.`my_table`\n" + + "WHERE rand() < 10.0 / (SELECT COUNT(*) FROM `main`.`default`.`my_table`)", + CONNECTOR.getRandomQuery("`main`.`default`.`my_table`", 10)); + } + + @Test + public void testGetStratifiedQuery() { + Assert.assertEquals("WITH t_s1 AS (\n" + + " SELECT *,\n" + + " ROW_NUMBER() OVER (ORDER BY id, RAND()) AS sqn_s1,\n" + + " COUNT(*) OVER () AS c_s1\n" + + " FROM `main`.`default`.`my_table`\n" + + " )\n" + + "SELECT * FROM t_s1\n" + + "WHERE MOD(sqn_s1, GREATEST(1, CAST(c_s1 / 10 AS BIGINT))) = 1\n" + + "ORDER BY id\n" + + "LIMIT 10", + CONNECTOR.getStratifiedQuery("`main`.`default`.`my_table`", 10, "id", "s1")); + } + + @Test + public void testGetDBRecordType() { + Assert.assertEquals("class io.cdap.plugin.databricks.DatabricksDBRecord", + CONNECTOR.getDBRecordType().toString()); + } + + @Test + public void testConnectionString() { + DatabricksConnectorConfig config = new DatabricksConnectorConfig( + "token", "secret", "jdbc", "", "dbc-xxx.cloud.databricks.com", + "sql/1.0/warehouses/xxx", "main", 443); + Assert.assertEquals( + "jdbc:databricks://dbc-xxx.cloud.databricks.com:443;ConnCatalog=main;HttpPath=sql/1.0/warehouses/xxx;", + config.getConnectionString()); + + DatabricksConnectorConfig configNoDb = new DatabricksConnectorConfig( + "token", "secret", "jdbc", "", "dbc-xxx.cloud.databricks.com", + "sql/1.0/warehouses/xxx", null, 443); + Assert.assertEquals( + "jdbc:databricks://dbc-xxx.cloud.databricks.com:443;HttpPath=sql/1.0/warehouses/xxx;", + configNoDb.getConnectionString()); + } + +} diff --git a/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksDBRecordUnitTest.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksDBRecordUnitTest.java new file mode 100644 index 000000000..b49a3608b --- /dev/null +++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksDBRecordUnitTest.java @@ -0,0 +1,145 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import io.cdap.cdap.api.data.format.StructuredRecord; +import io.cdap.cdap.api.data.schema.Schema; +import io.cdap.plugin.util.DBUtils; +import org.junit.Assert; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mockito; +import org.mockito.junit.MockitoJUnitRunner; + +import java.math.BigDecimal; +import java.sql.Date; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Timestamp; +import java.sql.Types; +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.time.OffsetDateTime; +import java.time.ZoneId; +import java.time.ZoneOffset; + +/** + * Unit tests for {@link DatabricksDBRecord}. + */ +@RunWith(MockitoJUnitRunner.class) +public class DatabricksDBRecordUnitTest { + + @Test + public void testHandleFieldTimestampAndTimestampNtz() throws SQLException { + OffsetDateTime offsetDateTime = OffsetDateTime.of(2026, 4, 16, 10, 30, 0, 123456000, ZoneOffset.UTC); + LocalDateTime localDateTime = LocalDateTime.of(2026, 4, 16, 10, 30, 0, 123456000); + + ResultSetMetaData metaData = Mockito.mock(ResultSetMetaData.class); + Mockito.when(metaData.getColumnTypeName(1)).thenReturn("TIMESTAMP"); + Mockito.when(metaData.getColumnTypeName(2)).thenReturn("TIMESTAMP_NTZ"); + + Timestamp ts = Timestamp.from(offsetDateTime.toInstant()); + ResultSet resultSet = Mockito.mock(ResultSet.class); + Mockito.when(resultSet.getMetaData()).thenReturn(metaData); + Mockito.when(resultSet.getObject(1)).thenReturn(ts); + Mockito.when(resultSet.getTimestamp(1, DBUtils.PURE_GREGORIAN_CALENDAR)).thenReturn(ts); + Mockito.when(resultSet.getTimestamp(2)).thenReturn(Timestamp.valueOf(localDateTime)); + + Schema.Field tsField = Schema.Field.of("ts_col", Schema.of(Schema.LogicalType.TIMESTAMP_MICROS)); + Schema.Field tsNtzField = Schema.Field.of("ts_ntz_col", Schema.of(Schema.LogicalType.DATETIME)); + Schema schema = Schema.recordOf("dbRecord", tsField, tsNtzField); + StructuredRecord.Builder builder = StructuredRecord.builder(schema); + + DatabricksDBRecord dbRecord = new DatabricksDBRecord(); + dbRecord.handleField(resultSet, builder, tsField, 1, Types.TIMESTAMP, 0, 0); + dbRecord.handleField(resultSet, builder, tsNtzField, 2, Types.TIMESTAMP, 0, 0); + + StructuredRecord record = builder.build(); + Assert.assertEquals(offsetDateTime.toInstant(), record.getTimestamp("ts_col", ZoneId.of("UTC")).toInstant()); + Assert.assertEquals(localDateTime, record.getDateTime("ts_ntz_col")); + } + + @Test + public void testHandleFieldDateAndDecimal() throws SQLException { + LocalDate expectedDate = LocalDate.of(2026, 4, 16); + Date sqlDate = Date.valueOf(expectedDate); + BigDecimal expectedDecimal = new BigDecimal("12345678901234567890.1234567890"); + + ResultSetMetaData metaData = Mockito.mock(ResultSetMetaData.class); + Mockito.when(metaData.getColumnTypeName(1)).thenReturn("DATE"); + Mockito.when(metaData.getColumnTypeName(2)).thenReturn("DECIMAL"); + + ResultSet resultSet = Mockito.mock(ResultSet.class); + Mockito.when(resultSet.getMetaData()).thenReturn(metaData); + Mockito.when(resultSet.getObject(1)).thenReturn(sqlDate); + Mockito.when(resultSet.getDate(1)).thenReturn(sqlDate); + Mockito.when(resultSet.getObject(2)).thenReturn(expectedDecimal); + + Schema.Field dateField = Schema.Field.of("date_col", Schema.of(Schema.LogicalType.DATE)); + Schema.Field decimalField = Schema.Field.of("dec_col", Schema.decimalOf(38, 10)); + Schema schema = Schema.recordOf("dbRecord", dateField, decimalField); + StructuredRecord.Builder builder = StructuredRecord.builder(schema); + + DatabricksDBRecord dbRecord = new DatabricksDBRecord(); + dbRecord.handleField(resultSet, builder, dateField, 1, Types.DATE, 0, 0); + dbRecord.handleField(resultSet, builder, decimalField, 2, Types.DECIMAL, 38, 10); + + StructuredRecord record = builder.build(); + Assert.assertEquals(expectedDate, record.getDate("date_col")); + Assert.assertEquals(expectedDecimal, record.getDecimal("dec_col")); + } + + @Test + public void testHandleFieldComplexAndNullTypes() throws SQLException { + ResultSetMetaData metaData = Mockito.mock(ResultSetMetaData.class); + Mockito.when(metaData.getColumnTypeName(1)).thenReturn("ARRAY"); + Mockito.when(metaData.getColumnTypeName(2)).thenReturn("MAP"); + Mockito.when(metaData.getColumnTypeName(3)).thenReturn("STRUCT"); + Mockito.when(metaData.getColumnTypeName(4)).thenReturn("VARIANT"); + Mockito.when(metaData.getColumnTypeName(5)).thenReturn("VOID"); + + ResultSet resultSet = Mockito.mock(ResultSet.class); + Mockito.when(resultSet.getMetaData()).thenReturn(metaData); + Mockito.when(resultSet.getObject(1)).thenReturn("[10001,10002]"); + Mockito.when(resultSet.getObject(2)).thenReturn("{\"pickup\":10001}"); + Mockito.when(resultSet.getObject(3)).thenReturn("{\"distance\":2.5,\"fare\":15.5}"); + Mockito.when(resultSet.getObject(4)).thenReturn("{\"vendor\":1}"); + + Schema.Field arrayField = Schema.Field.of("array_col", Schema.of(Schema.Type.STRING)); + Schema.Field mapField = Schema.Field.of("map_col", Schema.of(Schema.Type.STRING)); + Schema.Field structField = Schema.Field.of("struct_col", Schema.of(Schema.Type.STRING)); + Schema.Field variantField = Schema.Field.of("variant_col", Schema.of(Schema.Type.STRING)); + Schema.Field voidField = Schema.Field.of("void_col", Schema.nullableOf(Schema.of(Schema.Type.STRING))); + Schema schema = Schema.recordOf("dbRecord", arrayField, mapField, structField, variantField, voidField); + StructuredRecord.Builder builder = StructuredRecord.builder(schema); + + DatabricksDBRecord dbRecord = new DatabricksDBRecord(); + dbRecord.handleField(resultSet, builder, arrayField, 1, Types.ARRAY, 0, 0); + dbRecord.handleField(resultSet, builder, mapField, 2, Types.OTHER, 0, 0); + dbRecord.handleField(resultSet, builder, structField, 3, Types.STRUCT, 0, 0); + dbRecord.handleField(resultSet, builder, variantField, 4, Types.OTHER, 0, 0); + dbRecord.handleField(resultSet, builder, voidField, 5, Types.NULL, 0, 0); + + StructuredRecord record = builder.build(); + Assert.assertEquals("[10001,10002]", record.get("array_col")); + Assert.assertEquals("{\"pickup\":10001}", record.get("map_col")); + Assert.assertEquals("{\"distance\":2.5,\"fare\":15.5}", record.get("struct_col")); + Assert.assertEquals("{\"vendor\":1}", record.get("variant_col")); + Assert.assertNull(record.get("void_col")); + } +} diff --git a/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksFailedConnectionTest.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksFailedConnectionTest.java new file mode 100644 index 000000000..3292acbe1 --- /dev/null +++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksFailedConnectionTest.java @@ -0,0 +1,49 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import io.cdap.plugin.db.connector.DBSpecificFailedConnectionTest; +import org.junit.Assume; +import org.junit.Test; + +import java.io.IOException; + +/** + * Test failed connection handling for {@link DatabricksConnector}. + */ +public class DatabricksFailedConnectionTest extends DBSpecificFailedConnectionTest { + private static final String JDBC_DRIVER_CLASS_NAME = "com.databricks.client.jdbc.Driver"; + + @Test + public void test() throws ClassNotFoundException, IOException { + DatabricksConnector connector = new DatabricksConnector( + new DatabricksConnectorConfig("token", "password", "jdbc", "", "localhost", + "sql/1.0/warehouses/test", "db", 443)); + + try { + super.test(JDBC_DRIVER_CLASS_NAME, connector, + "Failed to create connection to database via connection string: " + + "jdbc:databricks://localhost:443;ConnCatalog=db;HttpPath=sql/1.0/warehouses/test; " + + "and arguments: {user=token}. Error: DatabricksHttpException: " + + "Caught error while executing http request: [https://localhost:443/sql/1.0/warehouses/test]. " + + "Error Message: [com.databricks.internal.apache.http.conn.HttpHostConnectException: " + + "Connect to localhost:443 [localhost/127.0.0.1] failed: Connection refused (Connection refused)]."); + } catch (UnsupportedClassVersionError e) { + Assume.assumeNoException(e); + } + } +} diff --git a/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksPluginTestBase.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksPluginTestBase.java new file mode 100644 index 000000000..50a9ba284 --- /dev/null +++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksPluginTestBase.java @@ -0,0 +1,205 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import com.google.common.base.Charsets; +import com.google.common.base.Throwables; +import com.google.common.collect.Sets; +import io.cdap.cdap.api.artifact.ArtifactSummary; +import io.cdap.cdap.api.plugin.PluginClass; +import io.cdap.cdap.datapipeline.DataPipelineApp; +import io.cdap.cdap.proto.id.ArtifactId; +import io.cdap.cdap.proto.id.NamespaceId; +import io.cdap.plugin.db.ConnectionConfig; +import io.cdap.plugin.db.DBRecord; +import io.cdap.plugin.db.batch.DatabasePluginTestBase; +import io.cdap.plugin.db.sink.ETLDBOutputFormat; +import io.cdap.plugin.db.source.DataDrivenETLDBInputFormat; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.math.BigDecimal; +import java.sql.Connection; +import java.sql.Date; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.SQLException; +import java.sql.Statement; +import java.sql.Timestamp; +import java.util.Arrays; +import java.util.Calendar; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.TimeZone; + +/** + * Base test class for Databricks plugins. + */ +public abstract class DatabricksPluginTestBase extends DatabasePluginTestBase { + private static final Logger LOGGER = LoggerFactory.getLogger(DatabricksPluginTestBase.class); + protected static final ArtifactId DATAPIPELINE_ARTIFACT_ID = NamespaceId.DEFAULT.artifact("data-pipeline", "3.2.0"); + protected static final ArtifactSummary DATAPIPELINE_ARTIFACT = new ArtifactSummary("data-pipeline", "3.2.0"); + protected static final long CURRENT_TS = System.currentTimeMillis(); + + protected static final String DRIVER_CLASS = "com.databricks.client.jdbc.Driver"; + protected static final String JDBC_DRIVER_NAME = "databricks"; + protected static final Map BASE_PROPS = new HashMap<>(); + + protected static String connectionUrl; + protected static int year; + protected static final int PRECISION = 10; + protected static final int SCALE = 6; + private static int startCount; + + @BeforeClass + public static void setupTest() throws Exception { + if (startCount++ > 0) { + return; + } + + getProperties(); + + Calendar calendar = Calendar.getInstance(); + calendar.setTime(new Date(CURRENT_TS)); + year = calendar.get(Calendar.YEAR); + + setupBatchArtifacts(DATAPIPELINE_ARTIFACT_ID, DataPipelineApp.class); + + addPluginArtifact(NamespaceId.DEFAULT.artifact(JDBC_DRIVER_NAME, "1.0.0"), + DATAPIPELINE_ARTIFACT_ID, + DatabricksSource.class, DatabricksDBRecord.class, DBRecord.class, + ETLDBOutputFormat.class, DataDrivenETLDBInputFormat.class); + + Class driverClass = Class.forName(DRIVER_CLASS); + + PluginClass databricksDriver = new PluginClass(ConnectionConfig.JDBC_PLUGIN_TYPE, JDBC_DRIVER_NAME, + "databricks driver class", driverClass.getName(), + null, Collections.emptyMap()); + addPluginArtifact(NamespaceId.DEFAULT.artifact("databricks-jdbc-connector", "1.0.0"), + DATAPIPELINE_ARTIFACT_ID, + Sets.newHashSet(databricksDriver), driverClass); + + TimeZone.setDefault(TimeZone.getTimeZone("UTC")); + + connectionUrl = String.format(DatabricksConstants.DATABRICKS_DB_CONNECTION_STRING_FORMAT, + BASE_PROPS.get(ConnectionConfig.HOST), + Integer.parseInt(BASE_PROPS.get(ConnectionConfig.PORT)), + BASE_PROPS.get(ConnectionConfig.DATABASE), + BASE_PROPS.get(DatabricksConnectorConfig.HTTP_PATH)); + Connection conn = createConnection(); + createTestTables(conn); + prepareTestData(conn); + } + + private static void getProperties() { + BASE_PROPS.put(ConnectionConfig.HOST, getPropertyOrSkip("databricks.host")); + BASE_PROPS.put(ConnectionConfig.PORT, getPropertyOrSkip("databricks.port")); + BASE_PROPS.put(ConnectionConfig.DATABASE, getPropertyOrSkip("databricks.database")); + BASE_PROPS.put(DatabricksConnectorConfig.HTTP_PATH, getPropertyOrSkip("databricks.httpPath")); + BASE_PROPS.put(ConnectionConfig.USER, getPropertyOrSkip("databricks.username")); + BASE_PROPS.put(ConnectionConfig.PASSWORD, getPropertyOrSkip("databricks.password")); + BASE_PROPS.put(ConnectionConfig.JDBC_PLUGIN_NAME, JDBC_DRIVER_NAME); + } + + protected static void createTestTables(Connection conn) throws SQLException { + try (Statement stmt = conn.createStatement()) { + stmt.execute("CREATE TABLE IF NOT EXISTS my_table (" + + "ID INT NOT NULL, " + + "NAME VARCHAR(40) NOT NULL, " + + "SCORE FLOAT, " + + "GRADUATED BOOLEAN, " + + "NOT_IMPORTED VARCHAR(30), " + + "SMALLINT_COL SMALLINT, " + + "BIG BIGINT, " + + "NUMERIC_COL DECIMAL(" + PRECISION + "," + SCALE + "), " + + "DECIMAL_COL DECIMAL(" + PRECISION + "," + SCALE + "), " + + "DOUBLE_PREC_COL DOUBLE, " + + "DATE_COL DATE, " + + "TIMESTAMP_COL TIMESTAMP, " + + "TIMESTAMP_NTZ_COL TIMESTAMP_NTZ, " + + "TEXT_COL STRING, " + + "CHAR_COL CHAR(100), " + + "BYTEA_COL BINARY" + + ")"); + stmt.execute("CREATE TABLE IF NOT EXISTS your_table AS SELECT * FROM my_table"); + } + } + + protected static void prepareTestData(Connection conn) throws SQLException { + try ( + PreparedStatement pStmt1 = + conn.prepareStatement("INSERT INTO my_table " + + "VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"); + PreparedStatement pStmt2 = + conn.prepareStatement("INSERT INTO your_table " + + "VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)")) { + populateData(pStmt1, pStmt2); + } + } + + private static void populateData(PreparedStatement... stmts) throws SQLException { + for (PreparedStatement pStmt : stmts) { + for (int i = 1; i <= 5; i++) { + String name = "user" + i; + pStmt.setInt(1, i); + pStmt.setString(2, name); + pStmt.setFloat(3, 123.45f + i); + pStmt.setBoolean(4, (i % 2 == 0)); + pStmt.setString(5, "random" + i); + pStmt.setShort(6, (short) i); + pStmt.setLong(7, (long) i); + pStmt.setBigDecimal(8, new BigDecimal("123.45").add(new BigDecimal(i))); + pStmt.setBigDecimal(9, new BigDecimal("123.45").add(new BigDecimal(i))); + pStmt.setDouble(10, 123.45 + i); + pStmt.setDate(11, new Date(CURRENT_TS)); + pStmt.setTimestamp(12, new Timestamp(CURRENT_TS)); + pStmt.setTimestamp(13, new Timestamp(CURRENT_TS)); + pStmt.setString(14, name); + pStmt.setString(15, "char" + i); + pStmt.setBytes(16, name.getBytes(Charsets.UTF_8)); + pStmt.executeUpdate(); + } + } + } + + public static Connection createConnection() { + try { + Class.forName(DRIVER_CLASS); + return DriverManager.getConnection(connectionUrl, BASE_PROPS.get(ConnectionConfig.USER), + BASE_PROPS.get(ConnectionConfig.PASSWORD)); + } catch (Exception e) { + throw Throwables.propagate(e); + } + } + + @AfterClass + public static void tearDownDB() { + if (connectionUrl == null) { + return; + } + try (Connection conn = createConnection(); + Statement stmt = conn.createStatement()) { + executeCleanup(Arrays.asList(() -> stmt.execute("DROP TABLE IF EXISTS my_table"), + () -> stmt.execute("DROP TABLE IF EXISTS your_table")), LOGGER); + } catch (Exception e) { + LOGGER.warn("Fail to tear down.", e); + } + } +} diff --git a/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksPluginTestSuite.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksPluginTestSuite.java new file mode 100644 index 000000000..888dee57f --- /dev/null +++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksPluginTestSuite.java @@ -0,0 +1,31 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import io.cdap.cdap.common.test.TestSuite; +import org.junit.runner.RunWith; +import org.junit.runners.Suite; + +/** + * This is a test suite that runs all the tests for Databricks plugins. + */ +@RunWith(TestSuite.class) +@Suite.SuiteClasses({ + DatabricksSourceTestRun.class, +}) +public class DatabricksPluginTestSuite extends DatabricksPluginTestBase { +} diff --git a/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksSchemaReaderTest.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksSchemaReaderTest.java new file mode 100644 index 000000000..1f1313b70 --- /dev/null +++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksSchemaReaderTest.java @@ -0,0 +1,93 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import io.cdap.cdap.api.data.schema.Schema; +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Types; + +public class DatabricksSchemaReaderTest { + + private void mockColumnType(ResultSetMetaData metadata, int index, String typeName, int sqlType) throws SQLException { + Mockito.when(metadata.getColumnTypeName(index)).thenReturn(typeName); + Mockito.when(metadata.getColumnType(index)).thenReturn(sqlType); + Mockito.when(metadata.isSigned(index)).thenReturn(true); + } + + @Test + public void testGetSchemaDatabricksTypes() throws SQLException { + DatabricksSchemaReader schemaReader = new DatabricksSchemaReader(); + ResultSetMetaData metadata = Mockito.mock(ResultSetMetaData.class); + + mockColumnType(metadata, 1, "INT", Types.INTEGER); + mockColumnType(metadata, 2, "BIGINT", Types.BIGINT); + mockColumnType(metadata, 3, "TIMESTAMP", Types.TIMESTAMP); + mockColumnType(metadata, 4, "TIMESTAMP_NTZ", Types.TIMESTAMP); + mockColumnType(metadata, 5, "DATE", Types.DATE); + mockColumnType(metadata, 6, "VARIANT", Types.OTHER); + mockColumnType(metadata, 7, "STRUCT", Types.STRUCT); + mockColumnType(metadata, 8, "ARRAY", Types.ARRAY); + mockColumnType(metadata, 9, "MAP", Types.OTHER); + mockColumnType(metadata, 10, "SMALLINT", Types.SMALLINT); + mockColumnType(metadata, 11, "TINYINT", Types.TINYINT); + mockColumnType(metadata, 12, "TIME", Types.TIME); + mockColumnType(metadata, 13, "INTERVAL", Types.OTHER); + mockColumnType(metadata, 14, "VOID", Types.NULL); + mockColumnType(metadata, 15, "GEOGRAPHY", Types.OTHER); + mockColumnType(metadata, 16, "GEOMETRY", Types.OTHER); + mockColumnType(metadata, 17, "FILE", Types.OTHER); + mockColumnType(metadata, 18, "OBJECT", Types.OTHER); + + Assert.assertEquals(Schema.of(Schema.Type.INT), schemaReader.getSchema(metadata, 1)); + Assert.assertEquals(Schema.of(Schema.Type.LONG), schemaReader.getSchema(metadata, 2)); + Assert.assertEquals(Schema.of(Schema.LogicalType.TIMESTAMP_MICROS), schemaReader.getSchema(metadata, 3)); + Assert.assertEquals(Schema.of(Schema.LogicalType.DATETIME), schemaReader.getSchema(metadata, 4)); + Assert.assertEquals(Schema.of(Schema.LogicalType.DATE), schemaReader.getSchema(metadata, 5)); + Assert.assertEquals(Schema.of(Schema.Type.STRING), schemaReader.getSchema(metadata, 6)); + Assert.assertEquals(Schema.of(Schema.Type.STRING), schemaReader.getSchema(metadata, 7)); + Assert.assertEquals(Schema.of(Schema.Type.STRING), schemaReader.getSchema(metadata, 8)); + Assert.assertEquals(Schema.of(Schema.Type.STRING), schemaReader.getSchema(metadata, 9)); + Assert.assertEquals(Schema.of(Schema.Type.INT), schemaReader.getSchema(metadata, 10)); + Assert.assertEquals(Schema.of(Schema.Type.INT), schemaReader.getSchema(metadata, 11)); + Assert.assertEquals(Schema.of(Schema.LogicalType.TIME_MICROS), schemaReader.getSchema(metadata, 12)); + Assert.assertEquals(Schema.of(Schema.Type.STRING), schemaReader.getSchema(metadata, 13)); + Assert.assertEquals(Schema.of(Schema.Type.STRING), schemaReader.getSchema(metadata, 14)); + Assert.assertEquals(Schema.of(Schema.Type.STRING), schemaReader.getSchema(metadata, 15)); + Assert.assertEquals(Schema.of(Schema.Type.STRING), schemaReader.getSchema(metadata, 16)); + Assert.assertEquals(Schema.of(Schema.Type.STRING), schemaReader.getSchema(metadata, 17)); + Assert.assertEquals(Schema.of(Schema.Type.STRING), schemaReader.getSchema(metadata, 18)); + } + + @Test + public void testShouldIgnoreColumn() throws SQLException { + DatabricksSchemaReader schemaReader = new DatabricksSchemaReader("sessionID"); + ResultSetMetaData metadata = Mockito.mock(ResultSetMetaData.class); + + Mockito.when(metadata.getColumnName(1)).thenReturn("c_sessionID"); + Mockito.when(metadata.getColumnName(2)).thenReturn("sqn_sessionID"); + Mockito.when(metadata.getColumnName(3)).thenReturn("columnName"); + + Assert.assertTrue(schemaReader.shouldIgnoreColumn(metadata, 1)); + Assert.assertTrue(schemaReader.shouldIgnoreColumn(metadata, 2)); + Assert.assertFalse(schemaReader.shouldIgnoreColumn(metadata, 3)); + } +} diff --git a/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksSourceTest.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksSourceTest.java new file mode 100644 index 000000000..85c1b2d08 --- /dev/null +++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksSourceTest.java @@ -0,0 +1,103 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import io.cdap.cdap.etl.api.batch.BatchSourceContext; +import io.cdap.plugin.common.LineageRecorder; +import io.cdap.plugin.db.SchemaReader; +import io.cdap.plugin.db.TransactionIsolationLevel; +import org.apache.hadoop.mapreduce.lib.db.DBWritable; +import org.junit.Assert; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mockito; +import org.mockito.junit.MockitoJUnitRunner; + +import java.util.Map; + +/** + * Unit tests for {@link DatabricksSource}. + */ +@RunWith(MockitoJUnitRunner.class) +public class DatabricksSourceTest { + + private DatabricksConnectorConfig createConnectorConfig() { + return new DatabricksConnectorConfig("token", "password", "jdbcPluginName", "connectionArguments", + "dbc-xxx.cloud.databricks.com", "sql/1.0/warehouses/xxx", "main", 443); + } + + @Test + public void testGetDBSpecificArguments() { + DatabricksSource.DatabricksSourceConfig config = + new DatabricksSource.DatabricksSourceConfig(false, createConnectorConfig()); + Map dbSpecificArguments = config.getDBSpecificArguments(); + Assert.assertEquals(0, dbSpecificArguments.size()); + } + + @Test + public void testGetFetchSize() { + DatabricksSource.DatabricksSourceConfig config = + new DatabricksSource.DatabricksSourceConfig(false, createConnectorConfig()); + Integer fetchSize = config.getFetchSize(); + Assert.assertEquals(1000, fetchSize.intValue()); + } + + @Test + public void testGetTransactionIsolationLevel() { + DatabricksSource.DatabricksSourceConfig config = + new DatabricksSource.DatabricksSourceConfig(false, createConnectorConfig()); + Assert.assertEquals(TransactionIsolationLevel.Level.TRANSACTION_REPEATABLE_READ.name(), + config.getTransactionIsolationLevel()); + } + + @Test + public void testGetSchemaReader() { + DatabricksSource source = + new DatabricksSource(new DatabricksSource.DatabricksSourceConfig(false, createConnectorConfig())); + SchemaReader schemaReader = source.getSchemaReader(); + Assert.assertTrue(schemaReader instanceof DatabricksSchemaReader); + } + + @Test + public void testGetDBRecordType() { + DatabricksSource source = + new DatabricksSource(new DatabricksSource.DatabricksSourceConfig(false, createConnectorConfig())); + Class dbRecordType = source.getDBRecordType(); + Assert.assertEquals(DatabricksDBRecord.class, dbRecordType); + } + + @Test + public void testCreateConnectionString() { + DatabricksSource.DatabricksSourceConfig config = + new DatabricksSource.DatabricksSourceConfig(false, createConnectorConfig()); + DatabricksSource source = new DatabricksSource(config); + Assert.assertEquals( + "jdbc:databricks://dbc-xxx.cloud.databricks.com:443;ConnCatalog=main;HttpPath=sql/1.0/warehouses/xxx;", + source.createConnectionString()); + } + + @Test + public void testGetLineageRecorder() { + BatchSourceContext context = Mockito.mock(BatchSourceContext.class); + DatabricksSource.DatabricksSourceConfig config = + new DatabricksSource.DatabricksSourceConfig(false, createConnectorConfig()); + DatabricksSource source = new DatabricksSource(config); + + LineageRecorder lineageRecorder = source.getLineageRecorder(context); + Assert.assertNotNull(lineageRecorder); + } +} diff --git a/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksSourceTestRun.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksSourceTestRun.java new file mode 100644 index 000000000..1a5ed22ff --- /dev/null +++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksSourceTestRun.java @@ -0,0 +1,283 @@ +/* + * Copyright © 2026 Cask Data, Inc. + * + * 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 io.cdap.plugin.databricks; + +import com.google.common.collect.ImmutableMap; +import io.cdap.cdap.api.common.Bytes; +import io.cdap.cdap.api.data.format.StructuredRecord; +import io.cdap.cdap.api.dataset.table.Table; +import io.cdap.cdap.etl.api.batch.BatchSource; +import io.cdap.cdap.etl.mock.batch.MockSink; +import io.cdap.cdap.etl.proto.v2.ETLBatchConfig; +import io.cdap.cdap.etl.proto.v2.ETLPlugin; +import io.cdap.cdap.etl.proto.v2.ETLStage; +import io.cdap.cdap.proto.artifact.AppRequest; +import io.cdap.cdap.proto.id.ApplicationId; +import io.cdap.cdap.proto.id.NamespaceId; +import io.cdap.cdap.test.ApplicationManager; +import io.cdap.cdap.test.DataSetManager; +import io.cdap.plugin.common.Constants; +import io.cdap.plugin.db.ConnectionConfig; +import io.cdap.plugin.db.DBConfig; +import io.cdap.plugin.db.source.AbstractDBSource; +import org.junit.Assert; +import org.junit.Test; + +import java.math.BigDecimal; +import java.math.MathContext; +import java.nio.ByteBuffer; +import java.sql.Date; +import java.text.SimpleDateFormat; +import java.time.Instant; +import java.time.LocalDate; +import java.time.ZoneId; +import java.time.ZoneOffset; +import java.time.ZonedDateTime; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Test for Databricks source plugin. + */ +public class DatabricksSourceTestRun extends DatabricksPluginTestBase { + + @Test + @SuppressWarnings("ConstantConditions") + public void testDBMacroSupport() throws Exception { + String importQuery = "SELECT * FROM my_table WHERE DATE_COL <= '${logicalStartTime(yyyy-MM-dd,1d)}' " + + "AND $CONDITIONS"; + String boundingQuery = "SELECT MIN(ID),MAX(ID) from my_table"; + String splitBy = "ID"; + + ImmutableMap sourceProps = ImmutableMap.builder() + .putAll(BASE_PROPS) + .put(AbstractDBSource.DBSourceConfig.IMPORT_QUERY, importQuery) + .put(AbstractDBSource.DBSourceConfig.BOUNDING_QUERY, boundingQuery) + .put(AbstractDBSource.DBSourceConfig.SPLIT_BY, splitBy) + .put(Constants.Reference.REFERENCE_NAME, "DBTestSource").build(); + + ETLPlugin sourceConfig = new ETLPlugin( + DatabricksConstants.PLUGIN_NAME, + BatchSource.PLUGIN_TYPE, + sourceProps + ); + + ETLPlugin sinkConfig = MockSink.getPlugin("macroOutputTable"); + + ApplicationManager appManager = deployETL(sourceConfig, sinkConfig, + DATAPIPELINE_ARTIFACT, "testDBMacro"); + runETLOnce(appManager, ImmutableMap.of("logical.start.time", String.valueOf(CURRENT_TS))); + + DataSetManager outputManager = getDataset("macroOutputTable"); + Assert.assertTrue(MockSink.readOutput(outputManager).isEmpty()); + } + + @Test + @SuppressWarnings("ConstantConditions") + public void testDBSource() throws Exception { + String importQuery = "SELECT ID, NAME, SCORE, GRADUATED, SMALLINT_COL, BIG, " + + "NUMERIC_COL, CHAR_COL, DECIMAL_COL, BYTEA_COL, DATE_COL, TIMESTAMP_COL, " + + "TIMESTAMP_NTZ_COL, TEXT_COL, DOUBLE_PREC_COL FROM my_table " + + "WHERE ID < 3 AND $CONDITIONS"; + String boundingQuery = "SELECT MIN(ID),MAX(ID) from my_table"; + String splitBy = "ID"; + ETLPlugin sourceConfig = new ETLPlugin( + DatabricksConstants.PLUGIN_NAME, + BatchSource.PLUGIN_TYPE, + ImmutableMap.builder() + .putAll(BASE_PROPS) + .put(AbstractDBSource.DBSourceConfig.IMPORT_QUERY, importQuery) + .put(AbstractDBSource.DBSourceConfig.BOUNDING_QUERY, boundingQuery) + .put(AbstractDBSource.DBSourceConfig.SPLIT_BY, splitBy) + .put(Constants.Reference.REFERENCE_NAME, "DBSourceTest") + .build(), + null + ); + + String outputDatasetName = "output-dbsourcetest"; + ETLPlugin sinkConfig = MockSink.getPlugin(outputDatasetName); + + ApplicationManager appManager = deployETL(sourceConfig, sinkConfig, + DATAPIPELINE_ARTIFACT, "testDBSource"); + runETLOnce(appManager); + + DataSetManager
outputManager = getDataset(outputDatasetName); + List outputRecords = MockSink.readOutput(outputManager); + + Assert.assertEquals(2, outputRecords.size()); + String userid = outputRecords.get(0).get("NAME"); + StructuredRecord row1 = "user1".equals(userid) ? outputRecords.get(0) : outputRecords.get(1); + StructuredRecord row2 = "user1".equals(userid) ? outputRecords.get(1) : outputRecords.get(0); + + Assert.assertEquals("user1", row1.get("NAME")); + Assert.assertEquals("user2", row2.get("NAME")); + Assert.assertEquals("user1", row1.get("TEXT_COL")); + Assert.assertEquals("user2", row2.get("TEXT_COL")); + Assert.assertEquals("char1", ((String) row1.get("CHAR_COL")).trim()); + Assert.assertEquals("char2", ((String) row2.get("CHAR_COL")).trim()); + Assert.assertEquals(124.45f, ((Float) row1.get("SCORE")).doubleValue(), 0.000001); + Assert.assertEquals(125.45f, ((Float) row2.get("SCORE")).doubleValue(), 0.000001); + Assert.assertEquals(false, row1.get("GRADUATED")); + Assert.assertEquals(true, row2.get("GRADUATED")); + Assert.assertNull(row1.get("NOT_IMPORTED")); + Assert.assertNull(row2.get("NOT_IMPORTED")); + + Assert.assertEquals(1, (int) row1.get("SMALLINT_COL")); + Assert.assertEquals(2, (int) row2.get("SMALLINT_COL")); + Assert.assertEquals(1, (long) row1.get("BIG")); + Assert.assertEquals(2, (long) row2.get("BIG")); + + Assert.assertEquals(new BigDecimal("124.45", new MathContext(PRECISION)).setScale(SCALE), + row1.getDecimal("NUMERIC_COL")); + Assert.assertEquals(new BigDecimal("125.45", new MathContext(PRECISION)).setScale(SCALE), + row2.getDecimal("NUMERIC_COL")); + Assert.assertEquals(new BigDecimal("124.45", new MathContext(PRECISION)).setScale(SCALE), + row1.getDecimal("DECIMAL_COL")); + + Assert.assertEquals(124.45, (double) row1.get("DOUBLE_PREC_COL"), 0.000001); + Assert.assertEquals(125.45, (double) row2.get("DOUBLE_PREC_COL"), 0.000001); + + Date date = new Date(CURRENT_TS); + SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd"); + LocalDate expectedDate = Date.valueOf(sdf.format(date)).toLocalDate(); + ZonedDateTime expectedTs = Instant.ofEpochMilli(CURRENT_TS).atZone(ZoneId.ofOffset("UTC", ZoneOffset.UTC)); + Assert.assertEquals(expectedDate, row1.getDate("DATE_COL")); + Assert.assertEquals(expectedTs, row1.getTimestamp("TIMESTAMP_COL", ZoneId.ofOffset("UTC", ZoneOffset.UTC))); + Assert.assertNotNull(row1.getDateTime("TIMESTAMP_NTZ_COL")); + + Assert.assertEquals("user1", Bytes.toString(((ByteBuffer) row1.get("BYTEA_COL")).array(), 0, 5)); + Assert.assertEquals("user2", Bytes.toString(((ByteBuffer) row2.get("BYTEA_COL")).array(), 0, 5)); + } + + @Test + public void testDbSourceMultipleTables() throws Exception { + String importQuery = "SELECT my_table.ID, your_table.NAME FROM my_table, your_table " + + "WHERE my_table.ID < 3 and my_table.ID = your_table.ID and $CONDITIONS"; + String boundingQuery = "SELECT LEAST(MIN(my_table.ID), MIN(your_table.ID)), " + + "GREATEST(MAX(my_table.ID), MAX(your_table.ID))"; + String splitBy = "my_table.ID"; + ETLPlugin sourceConfig = new ETLPlugin( + DatabricksConstants.PLUGIN_NAME, + BatchSource.PLUGIN_TYPE, + ImmutableMap.builder() + .putAll(BASE_PROPS) + .put(AbstractDBSource.DBSourceConfig.IMPORT_QUERY, importQuery) + .put(AbstractDBSource.DBSourceConfig.BOUNDING_QUERY, boundingQuery) + .put(AbstractDBSource.DBSourceConfig.SPLIT_BY, splitBy) + .put(Constants.Reference.REFERENCE_NAME, "DBMultipleTest") + .build(), + null + ); + + String outputDatasetName = "output-multitabletest"; + ETLPlugin sinkConfig = MockSink.getPlugin(outputDatasetName); + + ApplicationManager appManager = deployETL(sourceConfig, sinkConfig, + DATAPIPELINE_ARTIFACT, "testDBSourceWithMultipleTables"); + runETLOnce(appManager); + + DataSetManager
outputManager = getDataset(outputDatasetName); + List outputRecords = MockSink.readOutput(outputManager); + Assert.assertEquals(2, outputRecords.size()); + String userid = outputRecords.get(0).get("NAME"); + StructuredRecord row1 = "user1".equals(userid) ? outputRecords.get(0) : outputRecords.get(1); + StructuredRecord row2 = "user1".equals(userid) ? outputRecords.get(1) : outputRecords.get(0); + Assert.assertEquals("user1", row1.get("NAME")); + Assert.assertEquals("user2", row2.get("NAME")); + Assert.assertEquals(1, row1.get("ID").intValue()); + Assert.assertEquals(2, row2.get("ID").intValue()); + } + + @Test + public void testUserNamePasswordCombinations() throws Exception { + String importQuery = "SELECT * FROM my_table WHERE $CONDITIONS"; + String boundingQuery = "SELECT MIN(ID),MAX(ID) from my_table"; + String splitBy = "ID"; + + ETLPlugin sinkConfig = MockSink.getPlugin("outputTable"); + + Map baseSourceProps = ImmutableMap.builder() + .put(ConnectionConfig.HOST, BASE_PROPS.get(ConnectionConfig.HOST)) + .put(ConnectionConfig.PORT, BASE_PROPS.get(ConnectionConfig.PORT)) + .put(ConnectionConfig.DATABASE, BASE_PROPS.get(ConnectionConfig.DATABASE)) + .put(DatabricksConnectorConfig.HTTP_PATH, BASE_PROPS.get(DatabricksConnectorConfig.HTTP_PATH)) + .put(ConnectionConfig.JDBC_PLUGIN_NAME, JDBC_DRIVER_NAME) + .put(AbstractDBSource.DBSourceConfig.IMPORT_QUERY, importQuery) + .put(AbstractDBSource.DBSourceConfig.BOUNDING_QUERY, boundingQuery) + .put(AbstractDBSource.DBSourceConfig.SPLIT_BY, splitBy) + .put(Constants.Reference.REFERENCE_NAME, "UserPassDBTest") + .build(); + + ApplicationId appId = NamespaceId.DEFAULT.app("dbTest"); + + // null user name, null password. Should succeed. + ETLPlugin dbConfig = new ETLPlugin(DatabricksConstants.PLUGIN_NAME, BatchSource.PLUGIN_TYPE, + baseSourceProps, null); + ETLStage table = new ETLStage("uniqueTableSink", sinkConfig); + ETLStage database = new ETLStage("databaseSource", dbConfig); + ETLBatchConfig etlConfig = ETLBatchConfig.builder() + .addStage(database) + .addStage(table) + .addConnection(database.getName(), table.getName()) + .build(); + AppRequest appRequest = new AppRequest<>(DATAPIPELINE_ARTIFACT, etlConfig); + deployApplication(appId, appRequest); + + // null user name, non-null password. In Databricks, getUser() defaults to "token", so deployment succeeds. + Map tokenAuth = new HashMap<>(baseSourceProps); + tokenAuth.put(DBConfig.PASSWORD, BASE_PROPS.get(ConnectionConfig.PASSWORD)); + database = new ETLStage("databaseSource", new ETLPlugin(DatabricksConstants.PLUGIN_NAME, + BatchSource.PLUGIN_TYPE, tokenAuth, null)); + etlConfig = ETLBatchConfig.builder() + .addStage(database) + .addStage(table) + .addConnection(database.getName(), table.getName()) + .build(); + appRequest = new AppRequest<>(DATAPIPELINE_ARTIFACT, etlConfig); + deployApplication(appId, appRequest); + } + + @Test + public void testNonExistentDBTable() throws Exception { + String importQuery = "SELECT ID, NAME FROM dummy WHERE ID < 3 AND $CONDITIONS"; + String boundingQuery = "SELECT MIN(ID),MAX(ID) FROM dummy"; + String splitBy = "ID"; + ETLPlugin sinkConfig = MockSink.getPlugin("table"); + ETLPlugin sourceBadNameConfig = new ETLPlugin( + DatabricksConstants.PLUGIN_NAME, + BatchSource.PLUGIN_TYPE, + ImmutableMap.builder() + .putAll(BASE_PROPS) + .put(AbstractDBSource.DBSourceConfig.IMPORT_QUERY, importQuery) + .put(AbstractDBSource.DBSourceConfig.BOUNDING_QUERY, boundingQuery) + .put(AbstractDBSource.DBSourceConfig.SPLIT_BY, splitBy) + .put(Constants.Reference.REFERENCE_NAME, "DBNonExistentTest") + .build(), + null); + ETLStage sink = new ETLStage("sink", sinkConfig); + ETLStage sourceBadName = new ETLStage("sourceBadName", sourceBadNameConfig); + + ETLBatchConfig etlConfig = ETLBatchConfig.builder() + .addStage(sourceBadName) + .addStage(sink) + .addConnection(sourceBadName.getName(), sink.getName()) + .build(); + ApplicationId appId = NamespaceId.DEFAULT.app("dbSourceNonExistingTest"); + assertDeployAppFailure(appId, etlConfig, DATAPIPELINE_ARTIFACT); + } +} diff --git a/databricks-plugin/widgets/Databricks-batchsource.json b/databricks-plugin/widgets/Databricks-batchsource.json new file mode 100644 index 000000000..d6e0de95d --- /dev/null +++ b/databricks-plugin/widgets/Databricks-batchsource.json @@ -0,0 +1,279 @@ +{ + "metadata": { + "spec-version": "1.5" + }, + "display-name": "Databricks", + "configuration-groups": [ + { + "label": "Connection", + "properties": [ + { + "widget-type": "toggle", + "label": "Use connection", + "name": "useConnection", + "widget-attributes": { + "on": { + "value": "true", + "label": "YES" + }, + "off": { + "value": "false", + "label": "NO" + }, + "default": "false" + } + }, + { + "widget-type": "connection-select", + "label": "Connection", + "name": "connection", + "widget-attributes": { + "connectionType": "Databricks" + } + }, + { + "widget-type": "plugin-list", + "label": "JDBC Driver name", + "name": "jdbcPluginName", + "widget-attributes": { + "plugin-type": "jdbc" + } + }, + { + "widget-type": "textbox", + "label": "Host", + "name": "host", + "widget-attributes": { + "placeholder": "Databricks server hostname." + } + }, + { + "widget-type": "number", + "label": "Port", + "name": "port", + "widget-attributes": { + "default": "443" + } + }, + { + "widget-type": "textbox", + "label": "HTTP Path", + "name": "httpPath", + "widget-attributes": { + "placeholder": "e.g., sql/1.0/warehouses/xxxx" + } + }, + { + "widget-type": "textbox", + "label": "Username", + "name": "user" + }, + { + "widget-type": "password", + "label": "Password / Token", + "name": "password" + }, + { + "widget-type": "keyvalue", + "label": "Connection Arguments", + "name": "connectionArguments", + "widget-attributes": { + "showDelimiter": "false", + "key-placeholder": "Key", + "value-placeholder": "Value", + "kv-delimiter": "=", + "delimiter": ";" + } + } + ] + }, + { + "label": "Basic", + "properties": [ + { + "widget-type": "textbox", + "label": "Reference Name", + "name": "referenceName", + "widget-attributes": { + "placeholder": "Name used to identify this source for lineage. Typically, the name of the table/view." + } + }, + { + "widget-type": "textbox", + "label": "Database / Catalog", + "name": "database" + }, + { + "widget-type": "connection-browser", + "widget-category": "plugin", + "widget-attributes": { + "connectionType": "Databricks", + "label": "Browse Database" + } + } + ] + }, + { + "label": "SQL Query", + "properties": [ + { + "widget-type": "textarea", + "label": "Import Query", + "name": "importQuery", + "widget-attributes": { + "rows": "4" + } + }, + { + "widget-type": "get-schema", + "widget-category": "plugin" + } + ] + }, + { + "label": "Advanced", + "properties": [ + { + "widget-type": "textarea", + "label": "Bounding Query", + "name": "boundingQuery", + "widget-attributes": { + "rows": "4" + } + }, + { + "widget-type": "textbox", + "label": "Split-By Field Name", + "name": "splitBy" + }, + { + "widget-type": "textbox", + "label": "Number of Splits", + "name": "numSplits", + "widget-attributes": { + "default": "1" + } + }, + { + "widget-type": "number", + "label": "Fetch Size", + "name": "fetchSize", + "widget-attributes": { + "default": "1000", + "minimum": "0" + } + } + ] + }, + { + "properties": [ + { + "widget-type": "hidden", + "label": "Initial Retry Duration (sec)", + "name": "initialRetryDuration", + "widget-attributes": { + "default": 5, + "minimum": 0 + } + }, + { + "widget-type": "hidden", + "label": "Maximum Retry Duration (sec)", + "name": "maxRetryDuration", + "widget-attributes": { + "default": 80, + "minimum": 0 + } + }, + { + "widget-type": "hidden", + "label": "Maximum Retry Count", + "name": "maxRetryCount", + "widget-attributes": { + "default": 5, + "minimum": 0 + } + } + ] + } + ], + "outputs": [ + { + "name": "schema", + "widget-type": "schema", + "widget-attributes": { + "schema-types": [ + "boolean", + "int", + "long", + "float", + "double", + "bytes", + "string" + ], + "schema-default-type": "string" + } + } + ], + "filters": [ + { + "name": "showConnectionProperties", + "condition": { + "expression": "useConnection == false" + }, + "show": [ + { + "type": "property", + "name": "jdbcPluginName" + }, + { + "type": "property", + "name": "host" + }, + { + "type": "property", + "name": "port" + }, + { + "type": "property", + "name": "httpPath" + }, + { + "type": "property", + "name": "user" + }, + { + "type": "property", + "name": "password" + }, + { + "type": "property", + "name": "database" + }, + { + "type": "property", + "name": "connectionArguments" + } + ] + }, + { + "name": "showConnectionId", + "condition": { + "expression": "useConnection == true" + }, + "show": [ + { + "type": "property", + "name": "connection" + } + ] + } + ], + "jump-config": { + "datasets": [ + { + "ref-property-name": "referenceName" + } + ] + } +} diff --git a/databricks-plugin/widgets/Databricks-connector.json b/databricks-plugin/widgets/Databricks-connector.json new file mode 100644 index 000000000..15325289c --- /dev/null +++ b/databricks-plugin/widgets/Databricks-connector.json @@ -0,0 +1,114 @@ +{ + "metadata": { + "spec-version": "1.0" + }, + "display-name": "Databricks", + "configuration-groups": [ + { + "label": "Basic", + "properties": [ + { + "widget-type": "plugin-list", + "label": "JDBC Driver name", + "name": "jdbcPluginName", + "widget-attributes": { + "plugin-type": "jdbc" + } + }, + { + "widget-type": "textbox", + "label": "Host", + "name": "host", + "widget-attributes": { + "placeholder": "e.g., dbc-xxxx.cloud.databricks.com" + } + }, + { + "widget-type": "number", + "label": "Port", + "name": "port", + "widget-attributes": { + "default": "443" + } + }, + { + "widget-type": "textbox", + "label": "HTTP Path", + "name": "httpPath", + "widget-attributes": { + "placeholder": "e.g., sql/1.0/warehouses/xxxx" + } + }, + { + "widget-type": "textbox", + "label": "Database / Catalog", + "name": "database" + } + ] + }, + { + "label": "Credentials", + "properties": [ + { + "widget-type": "textbox", + "label": "Username", + "name": "user" + }, + { + "widget-type": "password", + "label": "Password / Token", + "name": "password" + } + ] + }, + { + "label": "Advanced", + "properties": [ + { + "widget-type": "keyvalue", + "label": "Connection Arguments", + "name": "connectionArguments", + "widget-attributes": { + "showDelimiter": "false", + "key-placeholder": "Key", + "value-placeholder": "Value", + "kv-delimiter": "=", + "delimiter": ";" + } + } + ] + }, + { + "properties": [ + { + "widget-type": "hidden", + "label": "Initial Retry Duration (sec)", + "name": "initialRetryDuration", + "widget-attributes": { + "default": 5, + "minimum": 0 + } + }, + { + "widget-type": "hidden", + "label": "Maximum Retry Duration (sec)", + "name": "maxRetryDuration", + "widget-attributes": { + "default": 80, + "minimum": 0 + } + }, + { + "widget-type": "hidden", + "label": "Maximum Retry Count", + "name": "maxRetryCount", + "widget-attributes": { + "default": 5, + "minimum": 0 + } + } + ] + } + ], + "outputs": [] +} diff --git a/pom.xml b/pom.xml index 54e6ef09e..c739d8855 100644 --- a/pom.xml +++ b/pom.xml @@ -45,6 +45,7 @@ teradata-plugin generic-db-argument-setter amazon-redshift-plugin + databricks-plugin