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/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..d87a0ab61 --- /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.1 + + + + + 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..fc54dd9e3 --- /dev/null +++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnector.java @@ -0,0 +1,162 @@ +/* + * 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.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 LIMIT %d", tableName, 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)); + + 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_READ_UNCOMMITTED.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..88756dc06 --- /dev/null +++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnectorConfig.java @@ -0,0 +1,129 @@ +/* + * 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 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 java.util.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..b82d88234 --- /dev/null +++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConstants.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; + +/** Databricks constants. */ +public final class DatabricksConstants { + + private DatabricksConstants() { + } + + public static final String PLUGIN_NAME = "Databricks"; + public static final String DRIVER_CLASS_NAME = "com.databricks.client.jdbc.Driver"; + 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/%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..2940c7db6 --- /dev/null +++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksDBRecord.java @@ -0,0 +1,65 @@ +/* + * 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; + +/** + * 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); + + if (columnTypeName != null && (columnTypeName.equalsIgnoreCase("VARIANT") || + columnTypeName.equalsIgnoreCase("ARRAY") || columnTypeName.equalsIgnoreCase("MAP") || + columnTypeName.equalsIgnoreCase("STRUCT") || columnTypeName.equalsIgnoreCase("JSON"))) { + Object value = resultSet.getObject(columnIndex); + if (value != null) { + recordBuilder.set(field.getName(), value.toString()); + } else { + recordBuilder.set(field.getName(), 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..9dbf8bfd6 --- /dev/null +++ b/databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksSchemaReader.java @@ -0,0 +1,79 @@ +/* + * 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; +import java.sql.Types; + +/** + * 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) { + if (typeName.equalsIgnoreCase("INT") || typeName.equalsIgnoreCase("INTEGER") || + typeName.equalsIgnoreCase("SMALLINT") || typeName.equalsIgnoreCase("TINYINT")) { + return Schema.of(Schema.Type.INT); + } + if (typeName.equalsIgnoreCase("BIGINT")) { + return Schema.of(Schema.Type.LONG); + } + if (typeName.equalsIgnoreCase("TIMESTAMP") || typeName.equalsIgnoreCase("TIMESTAMP_NTZ") || + typeName.equalsIgnoreCase("TIMESTAMPTZ")) { + return Schema.of(Schema.LogicalType.DATETIME); + } + if (typeName.equalsIgnoreCase("DATE")) { + return Schema.of(Schema.LogicalType.DATE); + } + if (typeName.equalsIgnoreCase("VARIANT") || typeName.equalsIgnoreCase("ARRAY") || + typeName.equalsIgnoreCase("MAP") || typeName.equalsIgnoreCase("STRUCT") || + typeName.equalsIgnoreCase("JSON")) { + 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..51aa810f8 --- /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_READ_UNCOMMITTED.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/DatabricksConnectorUnitTest.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorUnitTest.java new file mode 100644 index 000000000..8a71e8679 --- /dev/null +++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorUnitTest.java @@ -0,0 +1,70 @@ +/* + * 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` LIMIT 10", + CONNECTOR.getRandomQuery("`main`.`default`.`my_table`", 10)); + } + + @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/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/DatabricksSchemaReaderTest.java b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksSchemaReaderTest.java new file mode 100644 index 000000000..2bf486099 --- /dev/null +++ b/databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksSchemaReaderTest.java @@ -0,0 +1,90 @@ +/* + * 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 java.lang.reflect.Proxy; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.util.Map; + +public class DatabricksSchemaReaderTest { + + private ResultSetMetaData createMockMetadata(Map columnTypeNames, + Map columnNames) { + return (ResultSetMetaData) Proxy.newProxyInstance( + ResultSetMetaData.class.getClassLoader(), + new Class[]{ResultSetMetaData.class}, + (proxy, method, args) -> { + if ("getColumnTypeName".equals(method.getName())) { + int index = (Integer) args[0]; + return columnTypeNames.get(index); + } + if ("getColumnName".equals(method.getName())) { + int index = (Integer) args[0]; + return columnNames.get(index); + } + return null; + } + ); + } + + @Test + public void testGetSchemaDatabricksTypes() throws SQLException { + DatabricksSchemaReader schemaReader = new DatabricksSchemaReader(); + Map typeNames = new java.util.HashMap<>(); + typeNames.put(1, "INT"); + typeNames.put(2, "BIGINT"); + typeNames.put(3, "TIMESTAMP"); + typeNames.put(4, "TIMESTAMP_NTZ"); + typeNames.put(5, "DATE"); + typeNames.put(6, "VARIANT"); + typeNames.put(7, "STRUCT"); + typeNames.put(8, "ARRAY"); + typeNames.put(9, "MAP"); + + ResultSetMetaData metadata = createMockMetadata(typeNames, java.util.Collections.emptyMap()); + + 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.DATETIME), 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)); + } + + @Test + public void testShouldIgnoreColumn() throws SQLException { + DatabricksSchemaReader schemaReader = new DatabricksSchemaReader("sessionID"); + Map names = new java.util.HashMap<>(); + names.put(1, "c_sessionID"); + names.put(2, "sqn_sessionID"); + names.put(3, "columnName"); + + ResultSetMetaData metadata = createMockMetadata(java.util.Collections.emptyMap(), names); + + Assert.assertTrue(schemaReader.shouldIgnoreColumn(metadata, 1)); + Assert.assertTrue(schemaReader.shouldIgnoreColumn(metadata, 2)); + Assert.assertFalse(schemaReader.shouldIgnoreColumn(metadata, 3)); + } +} 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