fix(go/adbc/driver/internal/driverbase): proper unmarshalling for Con…

…straintColumnNames (#2285) Fixes #2278 The current `UnmarshalJSON` function for `requiredList` ended up unmarshalling into a copy of the underlying slice rather than using the actual slice, so the data wasn't being propagated appropriately. This fixes the `UnmarshalJSON` function to handle the scenario appropriately.
apache · Oct 30, 2024 · f6a8861 · f6a8861
1 parent e424aa2
commit f6a8861
Show file tree

Hide file tree

Showing 3 changed files with 100 additions and 3 deletions.
diff --git a/go/adbc/driver/internal/driverbase/connection.go b/go/adbc/driver/internal/driverbase/connection.go
@@ -583,9 +583,9 @@ func RequiredList[T any](vals []T) requiredList[T] {
 
 type requiredList[T any] []T
 
-func (n requiredList[T]) UnmarshalJSON(data []byte) error {
-	v := []T(n)
-	return json.Unmarshal(data, &v)
+func (n *requiredList[T]) UnmarshalJSON(data []byte) error {
+	v := (*[]T)(n)
+	return json.Unmarshal(data, v)
 }
 
 func (n requiredList[T]) MarshalJSON() ([]byte, error) {

diff --git a/go/adbc/driver/internal/driverbase/driver_test.go b/go/adbc/driver/internal/driverbase/driver_test.go
@@ -19,6 +19,7 @@ package driverbase_test
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"log/slog"
 	"slices"
@@ -30,6 +31,7 @@ import (
 	"github.com/apache/arrow/go/v18/arrow"
 	"github.com/apache/arrow/go/v18/arrow/array"
 	"github.com/apache/arrow/go/v18/arrow/memory"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 )
@@ -708,3 +710,13 @@ func tableFromRecordReader(rdr array.RecordReader) arrow.Table {
 	}
 	return array.NewTableFromRecords(rdr.Schema(), recs)
 }
+
+func TestRequiredList(t *testing.T) {
+	v := driverbase.RequiredList([]string{"a", "b", "c"})
+	result, err := json.Marshal(v)
+	require.NoError(t, err)
+	assert.JSONEq(t, `["a", "b", "c"]`, string(result))
+
+	require.NoError(t, json.Unmarshal([]byte(`["d", "e", "f"]`), &v))
+	assert.Equal(t, driverbase.RequiredList([]string{"d", "e", "f"}), v)
+}
diff --git a/go/adbc/driver/snowflake/driver_test.go b/go/adbc/driver/snowflake/driver_test.go
@@ -25,6 +25,7 @@ import (
 	"crypto/x509"
 	"database/sql"
 	"encoding/base64"
+	"encoding/json"
 	"encoding/pem"
 	"fmt"
 	"math"
@@ -35,6 +36,7 @@ import (
 	"testing"
 
 	"github.com/apache/arrow-adbc/go/adbc"
+	"github.com/apache/arrow-adbc/go/adbc/driver/internal/driverbase"
 	driver "github.com/apache/arrow-adbc/go/adbc/driver/snowflake"
 	"github.com/apache/arrow-adbc/go/adbc/validation"
 	"github.com/apache/arrow/go/v18/arrow"
@@ -43,6 +45,7 @@ import (
 	"github.com/apache/arrow/go/v18/arrow/memory"
 	"github.com/google/uuid"
 	"github.com/snowflakedb/gosnowflake"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 )
@@ -2167,3 +2170,85 @@ func (suite *SnowflakeTests) TestGetObjectsWithNilCatalog() {
 	// release on the result reader
 	rdr.Release()
 }
+
+func TestJSONUnmarshal(t *testing.T) {
+	jsonData := `{
+  "catalog_db_schemas": [
+    {
+      "db_schema_name": "PUBLIC",
+      "db_schema_tables": [
+        {
+          "table_columns": [
+            {
+              "column_name": "PRODUCT_ID",
+              "ordinal_position": 1,
+              "xdbc_char_octet_length": 16777216,
+              "xdbc_column_size": 16777216,
+              "xdbc_is_nullable": "NO",
+              "xdbc_nullable": 0,
+              "xdbc_type_name": "TEXT"
+            },
+            {
+              "column_name": "PRODUCT_NAME",
+              "ordinal_position": 2,
+              "xdbc_char_octet_length": 16777216,
+              "xdbc_column_size": 16777216,
+              "xdbc_is_nullable": "YES",
+              "xdbc_nullable": 1,
+              "xdbc_type_name": "TEXT"
+            },
+            {
+              "column_name": "ORDER_ID",
+              "ordinal_position": 3,
+              "xdbc_char_octet_length": 16777216,
+              "xdbc_column_size": 16777216,
+              "xdbc_is_nullable": "YES",
+              "xdbc_nullable": 1,
+              "xdbc_type_name": "TEXT"
+            }
+          ],
+          "table_constraints": [
+            {
+              "constraint_column_names": [
+                "PRODUCT_ID"
+              ],
+              "constraint_column_usage": [],
+              "constraint_name": "SYS_CONSTRAINT_386a9022-ad6e-47ed-94b6-f48501938f5f",
+              "constraint_type": "PRIMARY KEY"
+            },
+            {
+              "constraint_column_names": [
+                "ORDER_ID"
+              ],
+              "constraint_column_usage": [
+                {
+                  "fk_catalog": "DEMODB",
+                  "fk_column_name": "ORDER_ID",
+                  "fk_db_schema": "PUBLIC",
+                  "fk_table": "ORDERS2"
+                }
+              ],
+              "constraint_name": "SYS_CONSTRAINT_cfb3b763-4a97-4c0b-8356-a493c03b7e66",
+              "constraint_type": "FOREIGN KEY"
+            }
+          ],
+          "table_name": "PRODUCT2",
+          "table_type": "BASE TABLE"
+        }
+      ]
+    }
+  ],
+  "catalog_name": "DEMODB"
+}`
+
+	var getObjectsCatalog driverbase.GetObjectsInfo
+	require.NoError(t, json.Unmarshal([]byte(jsonData), &getObjectsCatalog))
+
+	for _, sch := range getObjectsCatalog.CatalogDbSchemas {
+		for _, tab := range sch.DbSchemaTables {
+			for _, con := range tab.TableConstraints {
+				assert.NotEmpty(t, con.ConstraintColumnNames)
+			}
+		}
+	}
+}