[GitHub] [arrow-adbc] lidavidm commented on a diff in pull request #586: feat(go/adbc/driver): Adbc Driver for Snowflake

["lidavidm (via GitHub)" <gi...@apache.org>]

lidavidm commented on code in PR #586:
URL: https://github.com/apache/arrow-adbc/pull/586#discussion_r1174223272


##########
go/adbc/driver/snowflake/statement.go:
##########
@@ -0,0 +1,572 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package snowflake
+
+import (
+	"context"
+	"database/sql/driver"
+	"fmt"
+	"strconv"
+	"strings"
+
+	"github.com/apache/arrow-adbc/go/adbc"
+	"github.com/apache/arrow/go/v12/arrow"
+	"github.com/apache/arrow/go/v12/arrow/array"
+	"github.com/apache/arrow/go/v12/arrow/memory"
+	"github.com/snowflakedb/gosnowflake"
+)
+
+const (
+	OptionStatementQueueSize = "adbc.flight.sql.rpc.queue_size"
+)
+
+type statement struct {
+	cnxn      *cnxn
+	alloc     memory.Allocator
+	queueSize int
+
+	query       string
+	targetTable string
+	append      bool
+
+	bound      arrow.Record
+	streamBind array.RecordReader
+}
+
+// Close releases any relevant resources associated with this statement
+// and closes it (particularly if it is a prepared statement).
+//
+// A statement instance should not be used after Close is called.
+func (st *statement) Close() error {
+	if st.cnxn == nil {
+		return adbc.Error{
+			Msg:  "statement already closed",
+			Code: adbc.StatusInvalidState}
+	}
+
+	if st.bound != nil {
+		st.bound.Release()
+		st.bound = nil
+	} else if st.streamBind != nil {
+		st.streamBind.Release()
+		st.streamBind = nil
+	}
+	st.cnxn = nil
+	return nil
+}
+
+// SetOption sets a string option on this statement
+func (st *statement) SetOption(key string, val string) error {
+	switch key {
+	case adbc.OptionKeyIngestTargetTable:
+		st.query = ""
+		st.targetTable = val
+	case adbc.OptionKeyIngestMode:
+		switch val {
+		case adbc.OptionValueIngestModeAppend:
+			st.append = true
+		case adbc.OptionValueIngestModeCreate:
+			st.append = false
+		default:
+			return adbc.Error{
+				Msg:  fmt.Sprintf("invalid statement option %s=%s", key, val),
+				Code: adbc.StatusInvalidArgument,
+			}
+		}
+	default:
+		return adbc.Error{
+			Msg:  fmt.Sprintf("invalid statement option %s=%s", key, val),
+			Code: adbc.StatusInvalidArgument,
+		}
+	}
+	return nil
+}
+
+// SetSqlQuery sets the query string to be executed.
+//
+// The query can then be executed with any of the Execute methods.
+// For queries expected to be executed repeatedly, Prepare should be
+// called before execution.
+func (st *statement) SetSqlQuery(query string) error {
+	st.query = query
+	st.targetTable = ""
+	return nil
+}
+
+func toSnowflakeType(dt arrow.DataType) string {
+	switch dt.ID() {
+	case arrow.EXTENSION:
+		return toSnowflakeType(dt.(arrow.ExtensionType).StorageType())
+	case arrow.DICTIONARY:
+		return toSnowflakeType(dt.(*arrow.DictionaryType).ValueType)
+	case arrow.RUN_END_ENCODED:
+		return toSnowflakeType(dt.(*arrow.RunEndEncodedType).Encoded())
+	case arrow.INT8, arrow.INT16, arrow.INT32, arrow.INT64,
+		arrow.UINT8, arrow.UINT16, arrow.UINT32, arrow.UINT64:
+		return "integer"
+	case arrow.FLOAT32, arrow.FLOAT16, arrow.FLOAT64:
+		return "double"
+	case arrow.DECIMAL, arrow.DECIMAL256:
+		dec := dt.(arrow.DecimalType)
+		return fmt.Sprintf("NUMERIC(%d,%d)", dec.GetPrecision(), dec.GetScale())
+	case arrow.STRING, arrow.LARGE_STRING:
+		return "text"
+	case arrow.BINARY, arrow.LARGE_BINARY:
+		return "binary"
+	case arrow.FIXED_SIZE_BINARY:
+		fsb := dt.(*arrow.FixedSizeBinaryType)
+		return fmt.Sprintf("binary(%d)", fsb.ByteWidth)
+	case arrow.BOOL:
+		return "boolean"
+	case arrow.TIME32, arrow.TIME64:
+		t := dt.(arrow.TemporalWithUnit)
+		prec := int(t.TimeUnit()) * 3
+		return fmt.Sprintf("time(%d)", prec)
+	case arrow.DATE32, arrow.DATE64:
+		return "date"
+	case arrow.TIMESTAMP:
+		ts := dt.(*arrow.TimestampType)
+		prec := int(ts.Unit) * 3
+		if ts.TimeZone == "" {
+			return fmt.Sprintf("timestamp_tz(%d)", prec)
+		}
+		return fmt.Sprintf("timestamp_ltz(%d)", prec)
+	case arrow.DENSE_UNION, arrow.SPARSE_UNION:
+		return "variant"
+	case arrow.LIST, arrow.LARGE_LIST, arrow.FIXED_SIZE_LIST:
+		return "array"
+	case arrow.STRUCT, arrow.MAP:
+		return "object"
+	}
+
+	return ""
+}
+
+func (st *statement) initIngest(ctx context.Context) (string, error) {
+	var (
+		createBldr, insertBldr strings.Builder
+	)
+
+	createBldr.WriteString("CREATE TABLE ")
+	createBldr.WriteString(st.targetTable)

Review Comment:
   Possibly table names can't be escaped. If it becomes an issue, we could always just restrict it to alphanumeric or similar.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org