Skip to content

Commit 2122d53

Browse files
krinartkrinart
committed
Add tests
1 parent 63c666c commit 2122d53

4 files changed

Lines changed: 543 additions & 0 deletions

File tree

‎go/bulk_ingest_test.go‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
// Copyright (c) 2026 ADBC Drivers Contributors
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package databricks
16+
17+
import (
18+
"testing"
19+
20+
"github.com/apache/arrow-go/v18/arrow"
21+
"github.com/stretchr/testify/assert"
22+
)
23+
24+
func TestBuildTableName(t *testing.T) {
25+
tests := []struct {
26+
name string
27+
catalog string
28+
schema string
29+
table string
30+
want string
31+
}{
32+
{
33+
name: "table only",
34+
table: "my_table",
35+
want: "`my_table`",
36+
},
37+
{
38+
name: "schema and table",
39+
schema: "my_schema",
40+
table: "my_table",
41+
want: "`my_schema`.`my_table`",
42+
},
43+
{
44+
name: "fully qualified",
45+
catalog: "my_catalog",
46+
schema: "my_schema",
47+
table: "my_table",
48+
want: "`my_catalog`.`my_schema`.`my_table`",
49+
},
50+
{
51+
name: "with backticks in name",
52+
catalog: "my`catalog",
53+
schema: "my`schema",
54+
table: "my`table",
55+
want: "`my``catalog`.`my``schema`.`my``table`",
56+
},
57+
{
58+
name: "catalog without schema",
59+
catalog: "my_catalog",
60+
table: "my_table",
61+
want: "`my_catalog`.`my_table`",
62+
},
63+
}
64+
65+
for _, tt := range tests {
66+
t.Run(tt.name, func(t *testing.T) {
67+
got := buildTableName(tt.catalog, tt.schema, tt.table)
68+
assert.Equal(t, tt.want, got)
69+
})
70+
}
71+
}
72+
73+
func TestQuoteIdentifier(t *testing.T) {
74+
tests := []struct {
75+
input string
76+
want string
77+
}{
78+
{"simple", "`simple`"},
79+
{"with space", "`with space`"},
80+
{"with`backtick", "`with``backtick`"},
81+
{"multiple``backticks", "`multiple````backticks`"},
82+
{"", "``"},
83+
}
84+
85+
for _, tt := range tests {
86+
t.Run(tt.input, func(t *testing.T) {
87+
assert.Equal(t, tt.want, quoteIdentifier(tt.input))
88+
})
89+
}
90+
}
91+
92+
func TestArrowTypeToDatabricksType(t *testing.T) {
93+
tests := []struct {
94+
name string
95+
dataType arrow.DataType
96+
want string
97+
}{
98+
{"bool", arrow.FixedWidthTypes.Boolean, "BOOLEAN"},
99+
{"int8", arrow.PrimitiveTypes.Int8, "TINYINT"},
100+
{"int16", arrow.PrimitiveTypes.Int16, "SMALLINT"},
101+
{"int32", arrow.PrimitiveTypes.Int32, "INT"},
102+
{"int64", arrow.PrimitiveTypes.Int64, "BIGINT"},
103+
{"uint8", arrow.PrimitiveTypes.Uint8, "SMALLINT"},
104+
{"uint16", arrow.PrimitiveTypes.Uint16, "INT"},
105+
{"uint32", arrow.PrimitiveTypes.Uint32, "BIGINT"},
106+
{"uint64", arrow.PrimitiveTypes.Uint64, "BIGINT"},
107+
{"float32", arrow.PrimitiveTypes.Float32, "FLOAT"},
108+
{"float64", arrow.PrimitiveTypes.Float64, "DOUBLE"},
109+
{"string", arrow.BinaryTypes.String, "STRING"},
110+
{"large_string", arrow.BinaryTypes.LargeString, "STRING"},
111+
{"binary", arrow.BinaryTypes.Binary, "BINARY"},
112+
{"large_binary", arrow.BinaryTypes.LargeBinary, "BINARY"},
113+
{"date32", arrow.PrimitiveTypes.Date32, "DATE"},
114+
{"date64", arrow.PrimitiveTypes.Date64, "DATE"},
115+
{"timestamp_with_tz", &arrow.TimestampType{Unit: arrow.Microsecond, TimeZone: "UTC"}, "TIMESTAMP"},
116+
{"timestamp_without_tz", &arrow.TimestampType{Unit: arrow.Microsecond}, "TIMESTAMP_NTZ"},
117+
{"decimal128", &arrow.Decimal128Type{Precision: 10, Scale: 2}, "DECIMAL(10, 2)"},
118+
{"decimal128_large", &arrow.Decimal128Type{Precision: 38, Scale: 18}, "DECIMAL(38, 18)"},
119+
}
120+
121+
for _, tt := range tests {
122+
t.Run(tt.name, func(t *testing.T) {
123+
assert.Equal(t, tt.want, arrowTypeToDatabricksType(tt.dataType))
124+
})
125+
}
126+
}
127+
128+
func TestPendingCopy(t *testing.T) {
129+
p := &pendingCopy{
130+
path: "Volumes/cat/sch/vol/staging/abc123.parquet",
131+
rows: 42,
132+
}
133+
134+
assert.Equal(t, "Volumes/cat/sch/vol/staging/abc123.parquet", p.String())
135+
assert.Equal(t, int64(42), p.Rows())
136+
}

‎go/database_test.go‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
// Copyright (c) 2026 ADBC Drivers Contributors
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package databricks
16+
17+
import (
18+
"testing"
19+
20+
"github.com/apache/arrow-adbc/go/adbc"
21+
"github.com/stretchr/testify/assert"
22+
"github.com/stretchr/testify/require"
23+
)
24+
25+
func TestParseURIForStaging(t *testing.T) {
26+
tests := []struct {
27+
name string
28+
uri string
29+
wantHost string
30+
wantPort int
31+
wantToken string
32+
wantErr bool
33+
}{
34+
{
35+
name: "standard URI with token",
36+
uri: "databricks://token:dapi1234abcd@workspace.cloud.databricks.com:443/sql/1.0/warehouses/abc",
37+
wantHost: "workspace.cloud.databricks.com",
38+
wantPort: 443,
39+
wantToken: "dapi1234abcd",
40+
},
41+
{
42+
name: "URI with custom port",
43+
uri: "databricks://token:mytoken@host.example.com:8443/path",
44+
wantHost: "host.example.com",
45+
wantPort: 8443,
46+
wantToken: "mytoken",
47+
},
48+
{
49+
name: "URI without port",
50+
uri: "databricks://token:mytoken@host.example.com/path",
51+
wantHost: "host.example.com",
52+
wantPort: 0,
53+
wantToken: "mytoken",
54+
},
55+
{
56+
name: "URI without password",
57+
uri: "databricks://user@host.example.com:443/path",
58+
wantHost: "host.example.com",
59+
wantPort: 443,
60+
wantToken: "",
61+
},
62+
{
63+
name: "URI with special characters in token",
64+
uri: "databricks://token:dapi%2Fabc123@host.com:443/path",
65+
wantHost: "host.com",
66+
wantPort: 443,
67+
wantToken: "dapi/abc123",
68+
},
69+
}
70+
71+
for _, tt := range tests {
72+
t.Run(tt.name, func(t *testing.T) {
73+
d := &databaseImpl{}
74+
err := d.parseURIForStaging(tt.uri)
75+
if tt.wantErr {
76+
assert.Error(t, err)
77+
return
78+
}
79+
require.NoError(t, err)
80+
assert.Equal(t, tt.wantHost, d.serverHostname)
81+
assert.Equal(t, tt.wantPort, d.port)
82+
assert.Equal(t, tt.wantToken, d.accessToken)
83+
})
84+
}
85+
}
86+
87+
func TestStagingOptionsWithURI(t *testing.T) {
88+
d := &databaseImpl{}
89+
90+
// Staging options should be allowed alongside URI
91+
err := d.SetOptions(map[string]string{
92+
adbc.OptionKeyURI: "databricks://token:mytoken@host.com:443/path",
93+
OptionStagingVolumePath: "cat.sch.vol",
94+
OptionStagingPrefix: "my_prefix",
95+
})
96+
require.NoError(t, err)
97+
98+
assert.Equal(t, "cat.sch.vol", d.stagingVolumePath)
99+
assert.Equal(t, "my_prefix", d.stagingPrefix)
100+
assert.Equal(t, "host.com", d.serverHostname)
101+
assert.Equal(t, 443, d.port)
102+
assert.Equal(t, "mytoken", d.accessToken)
103+
}
104+
105+
func TestStagingOptionsGetSet(t *testing.T) {
106+
d := &databaseImpl{}
107+
108+
require.NoError(t, d.SetOption(OptionStagingVolumePath, "cat.sch.vol"))
109+
require.NoError(t, d.SetOption(OptionStagingPrefix, "my_prefix"))
110+
111+
val, err := d.GetOption(OptionStagingVolumePath)
112+
require.NoError(t, err)
113+
assert.Equal(t, "cat.sch.vol", val)
114+
115+
val, err = d.GetOption(OptionStagingPrefix)
116+
require.NoError(t, err)
117+
assert.Equal(t, "my_prefix", val)
118+
}
119+
120+
func TestNewStagingClient(t *testing.T) {
121+
t.Run("returns nil when no volume path", func(t *testing.T) {
122+
d := &databaseImpl{}
123+
assert.Nil(t, d.newStagingClient())
124+
})
125+
126+
t.Run("creates client with volume path", func(t *testing.T) {
127+
d := &databaseImpl{
128+
serverHostname: "host.com",
129+
port: 443,
130+
accessToken: "token",
131+
stagingVolumePath: "cat.sch.vol",
132+
stagingPrefix: "prefix",
133+
}
134+
client := d.newStagingClient()
135+
require.NotNil(t, client)
136+
assert.Equal(t, "host.com", client.serverHostname)
137+
assert.Equal(t, 443, client.port)
138+
assert.Equal(t, "token", client.accessToken)
139+
assert.Equal(t, "cat.sch.vol", client.volumePath)
140+
assert.Equal(t, "prefix", client.prefix)
141+
assert.NotNil(t, client.httpClient)
142+
})
143+
}

‎go/staging_client.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,11 @@ func (c *stagingClient) generateFileName() (string, error) {
7676
// Upload uploads data to the staging area via the Databricks Files API.
7777
func (c *stagingClient) Upload(ctx context.Context, path string, data io.Reader) error {
7878
url := fmt.Sprintf("%s/%s", c.baseURL(), path)
79+
return c.uploadToURL(ctx, url, data)
80+
}
81+
82+
// uploadToURL performs the actual PUT request to the given URL.
83+
func (c *stagingClient) uploadToURL(ctx context.Context, url string, data io.Reader) error {
7984
req, err := http.NewRequestWithContext(ctx, http.MethodPut, url, data)
8085
if err != nil {
8186
return fmt.Errorf("failed to create upload request: %w", err)
@@ -99,6 +104,11 @@ func (c *stagingClient) Upload(ctx context.Context, path string, data io.Reader)
99104
// Delete removes a file from the staging area via the Databricks Files API.
100105
func (c *stagingClient) Delete(ctx context.Context, path string) error {
101106
url := fmt.Sprintf("%s/%s", c.baseURL(), path)
107+
return c.deleteFromURL(ctx, url)
108+
}
109+
110+
// deleteFromURL performs the actual DELETE request to the given URL.
111+
func (c *stagingClient) deleteFromURL(ctx context.Context, url string) error {
102112
req, err := http.NewRequestWithContext(ctx, http.MethodDelete, url, nil)
103113
if err != nil {
104114
return fmt.Errorf("failed to create delete request: %w", err)

0 commit comments

Comments
 (0)