Skip to content

Commit 80cabb2

Browse files
vikasrathee-csvikasrathee-cs
committed
Addition of Databricks Plugin
1 parent e1fe69b commit 80cabb2

16 files changed

Lines changed: 1344 additions & 2 deletions

‎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: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
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+
40+
import java.io.IOException;
41+
import java.sql.Connection;
42+
import java.sql.SQLException;
43+
import java.util.HashMap;
44+
import java.util.Map;
45+
46+
/**
47+
* Databricks Database Connector that connects to Databricks database via JDBC.
48+
*/
49+
@Plugin(type = Connector.PLUGIN_TYPE)
50+
@Name(DatabricksConstants.PLUGIN_NAME)
51+
@Description("Connection to access data in Databricks using JDBC.")
52+
@Category("Database")
53+
public class DatabricksConnector extends AbstractDBSpecificConnector<DatabricksDBRecord> {
54+
public static final String NAME = DatabricksConstants.PLUGIN_NAME;
55+
private final DatabricksConnectorConfig config;
56+
57+
public DatabricksConnector(DatabricksConnectorConfig config) {
58+
super(config);
59+
this.config = config;
60+
}
61+
62+
@Override
63+
protected DBConnectorPath getDBConnectorPath(String path) throws IOException {
64+
return DBSpecificPath.of(path, supportSchema());
65+
}
66+
67+
@Override
68+
protected Connection getConnection(DBConnectorPath path) {
69+
Connection connection = super.getConnection(path);
70+
try {
71+
connection.setTransactionIsolation(Connection.TRANSACTION_READ_UNCOMMITTED);
72+
} catch (SQLException e) {
73+
throw new RuntimeException("Failed to set transaction isolation level to READ_UNCOMMITTED", e);
74+
}
75+
return new NoOpCommitConnection(connection);
76+
}
77+
78+
@Override
79+
protected Connection getConnection() {
80+
Connection connection = super.getConnection();
81+
try {
82+
connection.setTransactionIsolation(Connection.TRANSACTION_READ_UNCOMMITTED);
83+
} catch (SQLException e) {
84+
throw new RuntimeException("Failed to set transaction isolation level to READ_UNCOMMITTED", e);
85+
}
86+
return new NoOpCommitConnection(connection);
87+
}
88+
89+
@Override
90+
public boolean supportSchema() {
91+
return true;
92+
}
93+
94+
@Override
95+
protected Class<? extends DBWritable> getDBRecordType() {
96+
return DatabricksDBRecord.class;
97+
}
98+
99+
@Override
100+
public StructuredRecord transform(LongWritable longWritable, DatabricksDBRecord record) {
101+
return record.getRecord();
102+
}
103+
104+
@Override
105+
protected SchemaReader getSchemaReader(String sessionID) {
106+
return new DatabricksSchemaReader(sessionID);
107+
}
108+
109+
@Override
110+
protected String getTableName(String database, String schema, String table) {
111+
if (database == null && schema == null) {
112+
return String.format("`%s`", table);
113+
}
114+
if (database == null) {
115+
return String.format("`%s`.`%s`", schema, table);
116+
}
117+
if (schema == null) {
118+
return String.format("`%s`.`%s`", database, table);
119+
}
120+
return String.format("`%s`.`%s`.`%s`", database, schema, table);
121+
}
122+
123+
@Override
124+
protected String getRandomQuery(String tableName, int limit) {
125+
return String.format("SELECT * FROM %s LIMIT %d", tableName, limit);
126+
}
127+
128+
@Override
129+
protected void setConnectorSpec(ConnectorSpecRequest request, DBConnectorPath path,
130+
ConnectorSpec.Builder builder) {
131+
Map<String, String> sourceProperties = new HashMap<>();
132+
setConnectionProperties(sourceProperties, request);
133+
builder.addRelatedPlugin(new PluginSpec(DatabricksConstants.PLUGIN_NAME,
134+
BatchSource.PLUGIN_TYPE, sourceProperties));
135+
136+
String schema = path.getSchema();
137+
sourceProperties.put(DatabricksSource.DatabricksSourceConfig.NUM_SPLITS, "1");
138+
sourceProperties.put(DatabricksSource.DatabricksSourceConfig.FETCH_SIZE,
139+
DatabricksSource.DatabricksSourceConfig.DEFAULT_FETCH_SIZE);
140+
String table = path.getTable();
141+
if (table == null) {
142+
return;
143+
}
144+
sourceProperties.put(DatabricksSource.DatabricksSourceConfig.IMPORT_QUERY,
145+
getTableQuery(path.getDatabase(), schema, table));
146+
sourceProperties.put(Constants.Reference.REFERENCE_NAME, ReferenceNames.cleanseReferenceName(table));
147+
}
148+
149+
@Override
150+
protected boolean isAutoCommitEnabled() {
151+
return true;
152+
}
153+
154+
@Override
155+
protected String getTransactionIsolationLevel() {
156+
return TransactionIsolationLevel.Level.TRANSACTION_READ_UNCOMMITTED.name();
157+
}
158+
}

0 commit comments

Comments
 (0)