Skip to content

Commit 4a41ca6

Browse files
gopalldbgopalldb
authored
Add integration tests and manual tests for UC volume operations using input stream (#387)
* changes * Minor test changes * add tests * . * test * test * Update changes * Update changes2 * fix imports * Address comments
1 parent e5bc274 commit 4a41ca6

File tree

3 files changed

+137
-20
lines changed

3 files changed

+137
-20
lines changed

‎src/test/java/com/databricks/jdbc/integration/e2e/UCVolumeInputStreamTests.java‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
package com.databricks.jdbc.integration.e2e;
2+
3+
import static com.databricks.jdbc.integration.IntegrationTestUtil.getDogfoodJDBCConnection;
4+
import static org.junit.jupiter.api.Assertions.*;
5+
6+
import com.databricks.jdbc.client.IDatabricksUCVolumeClient;
7+
import com.databricks.jdbc.core.IDatabricksConnection;
8+
import com.databricks.jdbc.driver.DatabricksJdbcConstants;
9+
import java.io.File;
10+
import java.io.FileInputStream;
11+
import java.nio.file.Files;
12+
import java.sql.Connection;
13+
import java.sql.SQLException;
14+
import org.apache.http.entity.InputStreamEntity;
15+
import org.junit.jupiter.api.AfterEach;
16+
import org.junit.jupiter.api.BeforeEach;
17+
import org.junit.jupiter.api.Test;
18+
19+
public class UCVolumeInputStreamTests {
20+
21+
private Connection con;
22+
23+
private static final String LOCAL_FILE = "/tmp/local-e2e.txt";
24+
private static final String VOLUME_FILE = "e2e-stream.csv";
25+
private static final String FILE_CONTENT = "test-data";
26+
private static final String VOL_CATALOG = "samikshya_hackathon";
27+
private static final String VOL_SCHEMA = "default";
28+
private static final String VOL_ROOT = "gopal-psl";
29+
30+
@BeforeEach
31+
void setUp() throws SQLException {
32+
con = getDogfoodJDBCConnection();
33+
System.out.println("Connection established......");
34+
}
35+
36+
@AfterEach
37+
void cleanUp() throws SQLException {
38+
if (con != null) {
39+
con.close();
40+
}
41+
}
42+
43+
@Test
44+
void testUCVolumeOperationsWithInputStream() throws Exception {
45+
IDatabricksUCVolumeClient client = ((IDatabricksConnection) con).getUCVolumeClient();
46+
47+
File file = new File(LOCAL_FILE);
48+
try {
49+
Files.writeString(file.toPath(), FILE_CONTENT);
50+
51+
System.out.println("File created");
52+
System.out.println(
53+
"Object inserted "
54+
+ client.putObject(
55+
VOL_CATALOG,
56+
VOL_SCHEMA,
57+
VOL_ROOT,
58+
VOLUME_FILE,
59+
new FileInputStream(file),
60+
file.length(),
61+
true));
62+
63+
InputStreamEntity inputStream =
64+
client.getObject(VOL_CATALOG, VOL_SCHEMA, VOL_ROOT, VOLUME_FILE);
65+
assertEquals(FILE_CONTENT, new String(inputStream.getContent().readAllBytes()));
66+
inputStream.getContent().close();
67+
68+
assertTrue(client.objectExists(VOL_CATALOG, VOL_SCHEMA, VOL_ROOT, VOLUME_FILE, false));
69+
con.setClientInfo(DatabricksJdbcConstants.ALLOWED_VOLUME_INGESTION_PATHS, "delete");
70+
client.deleteObject(VOL_CATALOG, VOL_SCHEMA, VOL_ROOT, VOLUME_FILE);
71+
assertFalse(client.objectExists(VOL_CATALOG, VOL_SCHEMA, VOL_ROOT, VOLUME_FILE, false));
72+
} finally {
73+
file.delete();
74+
}
75+
}
76+
}

‎src/test/java/com/databricks/jdbc/integration/e2e/UCVolumeTests.java‎

Lines changed: 4 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,8 @@
44
import static com.databricks.jdbc.integration.IntegrationTestUtil.*;
55
import static org.junit.jupiter.api.Assertions.*;
66

7-
import com.databricks.jdbc.client.impl.sdk.DatabricksUCVolumeClient;
8-
import java.io.ByteArrayInputStream;
9-
import java.io.InputStream;
7+
import com.databricks.jdbc.client.IDatabricksUCVolumeClient;
8+
import com.databricks.jdbc.core.IDatabricksConnection;
109
import java.nio.charset.StandardCharsets;
1110
import java.nio.file.Files;
1211
import java.nio.file.Paths;
@@ -18,14 +17,13 @@
1817
import java.util.stream.Stream;
1918
import org.junit.jupiter.api.AfterEach;
2019
import org.junit.jupiter.api.BeforeEach;
21-
import org.junit.jupiter.api.Test;
2220
import org.junit.jupiter.params.ParameterizedTest;
2321
import org.junit.jupiter.params.provider.Arguments;
2422
import org.junit.jupiter.params.provider.MethodSource;
2523

2624
public class UCVolumeTests {
2725

28-
private DatabricksUCVolumeClient client;
26+
private IDatabricksUCVolumeClient client;
2927
private Connection con;
3028

3129
private static final String LOCAL_TEST_DIRECTORY = "/tmp";
@@ -36,7 +34,7 @@ void setUp() throws SQLException {
3634
// test warehouse.
3735
con = getDogfoodJDBCConnection();
3836
System.out.println("Connection established......");
39-
client = new DatabricksUCVolumeClient(con);
37+
client = ((IDatabricksConnection) con).getUCVolumeClient();
4038
con.setClientInfo("allowlistedVolumeOperationLocalFilePaths", LOCAL_TEST_DIRECTORY);
4139
}
4240

@@ -516,18 +514,4 @@ private static Stream<Arguments> provideParametersForPutAndGetOverwriteTest() {
516514
"initialContent",
517515
"overwriteContent"));
518516
}
519-
520-
@Test
521-
public void testPutObjectWithInputStream() throws Exception {
522-
String volume = "test_volume1";
523-
String objectPath = "test_inputstream.txt";
524-
InputStream inputStream = new ByteArrayInputStream("testdata".getBytes(StandardCharsets.UTF_8));
525-
boolean toOverwrite = false;
526-
boolean expected = true;
527-
528-
boolean result =
529-
client.putObject(
530-
UC_VOLUME_CATALOG, UC_VOLUME_SCHEMA, volume, objectPath, inputStream, 10L, toOverwrite);
531-
assertEquals(expected, result);
532-
}
533517
}

‎src/test/java/com/databricks/jdbc/local/DriverTester.java‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,18 @@
11
package com.databricks.jdbc.local;
22

33
import com.databricks.client.jdbc.Driver;
4+
import com.databricks.jdbc.client.IDatabricksUCVolumeClient;
5+
import com.databricks.jdbc.core.IDatabricksConnection;
6+
import com.databricks.jdbc.driver.DatabricksJdbcConstants;
7+
import java.io.File;
8+
import java.io.FileInputStream;
49
import java.math.BigDecimal;
510
import java.math.BigInteger;
11+
import java.nio.file.Files;
612
import java.sql.*;
713
import java.time.LocalDate;
814
import java.util.*;
15+
import org.apache.http.entity.InputStreamEntity;
916
import org.junit.jupiter.api.Test;
1017

1118
public class DriverTester {
@@ -213,6 +220,56 @@ void testGetMetadataTypeBasedOnDBSQLVersion() throws Exception {
213220
con.close();
214221
}
215222

223+
@Test
224+
void testUCVolumeUsingInputStream() throws Exception {
225+
DriverManager.registerDriver(new Driver());
226+
DriverManager.drivers().forEach(driver -> System.out.println(driver.getClass()));
227+
System.out.println("Starting test");
228+
// Getting the connection
229+
String jdbcUrl =
230+
"jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=http;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/791ba2a31c7fd70a;";
231+
Connection con = DriverManager.getConnection(jdbcUrl, "token", "xx");
232+
con.setClientInfo(DatabricksJdbcConstants.ALLOWED_VOLUME_INGESTION_PATHS, "delete");
233+
System.out.println("Connection created");
234+
IDatabricksUCVolumeClient client = ((IDatabricksConnection) con).getUCVolumeClient();
235+
236+
File file = new File("/tmp/put.txt");
237+
try {
238+
Files.writeString(file.toPath(), "test-put");
239+
240+
System.out.println("File created");
241+
242+
System.out.println(
243+
"Object inserted "
244+
+ client.putObject(
245+
"samikshya_hackathon",
246+
"default",
247+
"gopal-psl",
248+
"test-stream.csv",
249+
new FileInputStream(file),
250+
file.length(),
251+
true));
252+
253+
InputStreamEntity inputStream =
254+
client.getObject("samikshya_hackathon", "default", "gopal-psl", "test-stream.csv");
255+
System.out.println("Got data " + new String(inputStream.getContent().readAllBytes()));
256+
inputStream.getContent().close();
257+
258+
System.out.println(
259+
"Object exists "
260+
+ client.objectExists(
261+
"samikshya_hackathon", "default", "gopal-psl", "test-stream.csv", false));
262+
client.deleteObject("samikshya_hackathon", "default", "gopal-psl", "test-stream.csv");
263+
System.out.println(
264+
"Object exists "
265+
+ client.objectExists(
266+
"samikshya_hackathon", "default", "gopal-psl", "test-stream.csv", false));
267+
} finally {
268+
file.delete();
269+
con.close();
270+
}
271+
}
272+
216273
@Test
217274
public void tooManyParameters() throws SQLException {
218275
DriverManager.registerDriver(new Driver());

0 commit comments

Comments
 (0)