Skip to content

Commit 9a4d4fc

Browse files
vikasrathee-csvikasrathee-cs
committed
Addition of Databricks Plugin
1 parent e1fe69b commit 9a4d4fc

16 files changed

Lines changed: 1350 additions & 1 deletion

‎database-commons/src/main/java/io/cdap/plugin/db/connector/AbstractDBSpecificConnector.java‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,8 +104,12 @@ public InputFormatProvider getInputFormatProvider(ConnectorContext context, Samp
104104
String tableQuery = getTableQuery(path.getDatabase(), path.getSchema(), path.getTable(), request.getLimit(),
105105
request.getProperties().get("sampleType"), request.getProperties().get("strata"), sessionID);
106106
DataDrivenETLDBInputFormat.setInput(connectionConfigAccessor.getConfiguration(), getDBRecordType(),
107-
tableQuery, null, false);
107+
tableQuery, null, isAutoCommitEnabled());
108108
connectionConfigAccessor.setConnectionArguments(Maps.fromProperties(config.getConnectionArgumentsProperties()));
109+
String isolationLevel = getTransactionIsolationLevel();
110+
if (isolationLevel != null) {
111+
connectionConfigAccessor.setTransactionIsolationLevel(isolationLevel);
112+
}
109113
connectionConfigAccessor.getConfiguration().setInt(MRJobConfig.NUM_MAPS, 1);
110114
Map<String, String> additionalArguments = config.getAdditionalArguments();
111115
for (Map.Entry<String, String> argument : additionalArguments.entrySet()) {
@@ -221,4 +225,19 @@ protected Schema getTableSchema(Connection connection, String database,
221225
protected String generateSessionID() {
222226
return UUID.randomUUID().toString().replace('-', '_');
223227
}
228+
229+
/**
230+
* Returns whether auto-commit should be enabled for this connector.
231+
* By default, it is false.
232+
*/
233+
protected boolean isAutoCommitEnabled() {
234+
return false;
235+
}
236+
/**
237+
* Returns the default transaction isolation level for this connector.
238+
* If null, it falls back to the database driver's default or serializable.
239+
*/
240+
protected String getTransactionIsolationLevel() {
241+
return null;
242+
}
224243
}

‎databricks-plugin/docs/Databricks-batchsource.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Databricks Batch Source
2+
3+
Description
4+
-----------
5+
Reads data from a Databricks table using a configurable SQL query.
6+
7+
Properties
8+
----------
9+
* **Use Connection**: Whether to use an existing Databricks connection.
10+
* **Host**: Server Hostname of the Databricks cluster or SQL warehouse.
11+
* **Port**: Database port (default is 443).
12+
* **HTTP Path**: The HTTP Path for the Databricks cluster or SQL warehouse.
13+
* **Reference Name**: Name used to identify this source for lineage.
14+
* **Database / Catalog**: Optional catalog or database name.
15+
* **Import Query**: SQL query to execute against Databricks.

‎databricks-plugin/docs/Databricks-connector.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# Databricks Database Connector
2+
3+
Description
4+
-----------
5+
Connects to Databricks database / Lakehouse via JDBC.
6+
7+
Properties
8+
----------
9+
* **Host**: Server Hostname of the Databricks cluster or SQL warehouse.
10+
* **Port**: Database port (default is 443).
11+
* **HTTP Path**: The HTTP Path for the Databricks cluster or SQL warehouse.
12+
* **Database / Catalog**: Optional catalog or database name to connect to.
13+
* **Username**: Username / token user.
14+
* **Password / Token**: Personal Access Token (PAT) or password.
15+
* **Connection Arguments**: Arbitrary key-value pairs to pass as connection arguments to the JDBC driver (e.g. `AuthMech=11;Auth_Flow=2`).

‎databricks-plugin/icons/Databricks-batchsource.png‎

27.9 KB
Loading

‎databricks-plugin/pom.xml‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<!--
3+
Copyright © 2026 CDAP
4+
5+
Licensed under the Apache License, Version 2.0 (the "License"); you may not
6+
use this file except in compliance with the License. You may obtain a copy of
7+
the License at
8+
9+
http://www.apache.org/licenses/LICENSE-2.0
10+
11+
Unless required by applicable law or agreed to in writing, software
12+
distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
13+
WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
14+
License for the specific language governing permissions and limitations under
15+
the License.
16+
-->
17+
<project xmlns="http://maven.apache.org/POM/4.0.0"
18+
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
19+
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
20+
<parent>
21+
<artifactId>database-plugins-parent</artifactId>
22+
<groupId>io.cdap.plugin</groupId>
23+
<version>1.13.0-SNAPSHOT</version>
24+
</parent>
25+
26+
<name>Databricks plugin</name>
27+
<artifactId>databricks-plugin</artifactId>
28+
<modelVersion>4.0.0</modelVersion>
29+
30+
<properties>
31+
<databricks-jdbc.version>3.4.1</databricks-jdbc.version>
32+
</properties>
33+
34+
<dependencies>
35+
<dependency>
36+
<groupId>io.cdap.cdap</groupId>
37+
<artifactId>cdap-etl-api</artifactId>
38+
</dependency>
39+
<dependency>
40+
<groupId>io.cdap.plugin</groupId>
41+
<artifactId>database-commons</artifactId>
42+
<version>${project.version}</version>
43+
</dependency>
44+
<dependency>
45+
<groupId>io.cdap.plugin</groupId>
46+
<artifactId>hydrator-common</artifactId>
47+
</dependency>
48+
<dependency>
49+
<groupId>com.google.guava</groupId>
50+
<artifactId>guava</artifactId>
51+
</dependency>
52+
53+
<!-- test dependencies -->
54+
<dependency>
55+
<groupId>com.databricks</groupId>
56+
<artifactId>databricks-jdbc</artifactId>
57+
<version>${databricks-jdbc.version}</version>
58+
<scope>test</scope>
59+
</dependency>
60+
<dependency>
61+
<groupId>io.cdap.plugin</groupId>
62+
<artifactId>database-commons</artifactId>
63+
<version>${project.version}</version>
64+
<type>test-jar</type>
65+
<scope>test</scope>
66+
</dependency>
67+
<dependency>
68+
<groupId>io.cdap.cdap</groupId>
69+
<artifactId>hydrator-test</artifactId>
70+
</dependency>
71+
<dependency>
72+
<groupId>io.cdap.cdap</groupId>
73+
<artifactId>cdap-data-pipeline3_2.12</artifactId>
74+
</dependency>
75+
<dependency>
76+
<groupId>junit</groupId>
77+
<artifactId>junit</artifactId>
78+
</dependency>
79+
<dependency>
80+
<groupId>org.mockito</groupId>
81+
<artifactId>mockito-core</artifactId>
82+
<scope>test</scope>
83+
</dependency>
84+
<dependency>
85+
<groupId>io.cdap.cdap</groupId>
86+
<artifactId>cdap-api</artifactId>
87+
<scope>provided</scope>
88+
</dependency>
89+
</dependencies>
90+
91+
<build>
92+
<plugins>
93+
<plugin>
94+
<groupId>io.cdap</groupId>
95+
<artifactId>cdap-maven-plugin</artifactId>
96+
</plugin>
97+
<plugin>
98+
<groupId>org.apache.felix</groupId>
99+
<artifactId>maven-bundle-plugin</artifactId>
100+
<version>5.1.2</version>
101+
<extensions>true</extensions>
102+
<configuration>
103+
<instructions>
104+
<_exportcontents>
105+
io.cdap.plugin.databricks.*;
106+
io.cdap.plugin.db.source.*;
107+
org.apache.commons.lang;
108+
org.apache.commons.logging.*;
109+
org.codehaus.jackson.*
110+
</_exportcontents>
111+
<Embed-Dependency>*;inline=false;scope=compile</Embed-Dependency>
112+
<Embed-Transitive>true</Embed-Transitive>
113+
<Embed-Directory>lib</Embed-Directory>
114+
</instructions>
115+
</configuration>
116+
<executions>
117+
<execution>
118+
<phase>package</phase>
119+
<goals>
120+
<goal>bundle</goal>
121+
</goals>
122+
</execution>
123+
</executions>
124+
</plugin>
125+
</plugins>
126+
</build>
127+
</project>

‎databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnector.java‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
/*
2+
* Copyright © 2026 Cask Data, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License"); you may not
5+
* use this file except in compliance with the License. You may obtain a copy of
6+
* the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12+
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13+
* License for the specific language governing permissions and limitations under
14+
* the License.
15+
*/
16+
17+
package io.cdap.plugin.databricks;
18+
19+
import io.cdap.cdap.api.annotation.Category;
20+
import io.cdap.cdap.api.annotation.Description;
21+
import io.cdap.cdap.api.annotation.Name;
22+
import io.cdap.cdap.api.annotation.Plugin;
23+
import io.cdap.cdap.api.data.format.StructuredRecord;
24+
import io.cdap.cdap.etl.api.batch.BatchSource;
25+
import io.cdap.cdap.etl.api.connector.Connector;
26+
import io.cdap.cdap.etl.api.connector.ConnectorSpec;
27+
import io.cdap.cdap.etl.api.connector.ConnectorSpecRequest;
28+
import io.cdap.cdap.etl.api.connector.PluginSpec;
29+
import io.cdap.plugin.common.Constants;
30+
import io.cdap.plugin.common.ReferenceNames;
31+
import io.cdap.plugin.common.db.DBConnectorPath;
32+
import io.cdap.plugin.db.NoOpCommitConnection;
33+
import io.cdap.plugin.db.SchemaReader;
34+
import io.cdap.plugin.db.TransactionIsolationLevel;
35+
import io.cdap.plugin.db.connector.AbstractDBSpecificConnector;
36+
import io.cdap.plugin.db.connector.DBSpecificPath;
37+
import org.apache.hadoop.io.LongWritable;
38+
import org.apache.hadoop.mapreduce.lib.db.DBWritable;
39+
import org.slf4j.Logger;
40+
import org.slf4j.LoggerFactory;
41+
42+
import java.io.IOException;
43+
import java.sql.Connection;
44+
import java.sql.SQLException;
45+
import java.util.HashMap;
46+
import java.util.Map;
47+
48+
/**
49+
* Databricks Database Connector that connects to Databricks database via JDBC.
50+
*/
51+
@Plugin(type = Connector.PLUGIN_TYPE)
52+
@Name(DatabricksConstants.PLUGIN_NAME)
53+
@Description("Connection to access data in Databricks using JDBC.")
54+
@Category("Database")
55+
public class DatabricksConnector extends AbstractDBSpecificConnector<DatabricksDBRecord> {
56+
public static final String NAME = DatabricksConstants.PLUGIN_NAME;
57+
private final DatabricksConnectorConfig config;
58+
59+
private static final Logger LOG = LoggerFactory.getLogger(DatabricksConnector.class);
60+
61+
public DatabricksConnector(DatabricksConnectorConfig config) {
62+
super(config);
63+
this.config = config;
64+
}
65+
66+
@Override
67+
protected DBConnectorPath getDBConnectorPath(String path) throws IOException {
68+
return DBSpecificPath.of(path, supportSchema());
69+
}
70+
71+
@Override
72+
protected Connection getConnection(DBConnectorPath path) {
73+
Connection connection = super.getConnection(path);
74+
try {
75+
connection.setTransactionIsolation(Connection.TRANSACTION_READ_UNCOMMITTED);
76+
} catch (SQLException e) {
77+
LOG.warn("Failed to set transaction isolation level to READ_UNCOMMITTED", e);
78+
}
79+
return new NoOpCommitConnection(connection);
80+
}
81+
82+
@Override
83+
protected Connection getConnection() {
84+
Connection connection = super.getConnection();
85+
try {
86+
connection.setTransactionIsolation(Connection.TRANSACTION_READ_UNCOMMITTED);
87+
} catch (SQLException e) {
88+
LOG.warn("Failed to set transaction isolation level to READ_UNCOMMITTED", e);
89+
}
90+
return new NoOpCommitConnection(connection);
91+
}
92+
93+
@Override
94+
public boolean supportSchema() {
95+
return true;
96+
}
97+
98+
@Override
99+
protected Class<? extends DBWritable> getDBRecordType() {
100+
return DatabricksDBRecord.class;
101+
}
102+
103+
@Override
104+
public StructuredRecord transform(LongWritable longWritable, DatabricksDBRecord record) {
105+
return record.getRecord();
106+
}
107+
108+
@Override
109+
protected SchemaReader getSchemaReader(String sessionID) {
110+
return new DatabricksSchemaReader(sessionID);
111+
}
112+
113+
@Override
114+
protected String getTableName(String database, String schema, String table) {
115+
if (database == null && schema == null) {
116+
return String.format("`%s`", table);
117+
}
118+
if (database == null) {
119+
return String.format("`%s`.`%s`", schema, table);
120+
}
121+
if (schema == null) {
122+
return String.format("`%s`.`%s`", database, table);
123+
}
124+
return String.format("`%s`.`%s`.`%s`", database, schema, table);
125+
}
126+
127+
@Override
128+
protected String getRandomQuery(String tableName, int limit) {
129+
return String.format("SELECT * FROM %s LIMIT %d", tableName, limit);
130+
}
131+
132+
@Override
133+
protected void setConnectorSpec(ConnectorSpecRequest request, DBConnectorPath path,
134+
ConnectorSpec.Builder builder) {
135+
Map<String, String> sourceProperties = new HashMap<>();
136+
setConnectionProperties(sourceProperties, request);
137+
builder.addRelatedPlugin(new PluginSpec(DatabricksConstants.PLUGIN_NAME,
138+
BatchSource.PLUGIN_TYPE, sourceProperties));
139+
140+
String schema = path.getSchema();
141+
sourceProperties.put(DatabricksSource.DatabricksSourceConfig.NUM_SPLITS, "1");
142+
sourceProperties.put(DatabricksSource.DatabricksSourceConfig.FETCH_SIZE,
143+
DatabricksSource.DatabricksSourceConfig.DEFAULT_FETCH_SIZE);
144+
String table = path.getTable();
145+
if (table == null) {
146+
return;
147+
}
148+
sourceProperties.put(DatabricksSource.DatabricksSourceConfig.IMPORT_QUERY,
149+
getTableQuery(path.getDatabase(), schema, table));
150+
sourceProperties.put(Constants.Reference.REFERENCE_NAME, ReferenceNames.cleanseReferenceName(table));
151+
}
152+
153+
@Override
154+
protected boolean isAutoCommitEnabled() {
155+
return true;
156+
}
157+
158+
@Override
159+
protected String getTransactionIsolationLevel() {
160+
return TransactionIsolationLevel.Level.TRANSACTION_READ_UNCOMMITTED.name();
161+
}
162+
}

0 commit comments

Comments
 (0)