From 96d3bf48b9178f9b28606b54a51a1080ac48d836 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Wa=C5=9B?= <jan.was@starburstdata.com>
Date: Thu, 19 Oct 2023 16:01:21 +0200
Subject: [PATCH] Add tests for db.Exec and INSERT

---
 trino/etc/catalog/memory.properties |  1 +
 trino/trino_test.go                 | 49 +++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)
 create mode 100644 trino/etc/catalog/memory.properties

diff --git a/trino/etc/catalog/memory.properties b/trino/etc/catalog/memory.properties
new file mode 100644
index 0000000..833abd3
--- /dev/null
+++ b/trino/etc/catalog/memory.properties
@@ -0,0 +1 @@
+connector.name=memory
diff --git a/trino/trino_test.go b/trino/trino_test.go
index ef8a46a..fc65c5e 100644
--- a/trino/trino_test.go
+++ b/trino/trino_test.go
@@ -1719,3 +1719,52 @@ func BenchmarkQuery(b *testing.B) {
 		rows.Close()
 	}
 }
+
+func TestExec(t *testing.T) {
+	if testing.Short() {
+		t.Skip("Skipping test in short mode.")
+	}
+	c := &Config{
+		ServerURI:         *integrationServerFlag,
+		SessionProperties: map[string]string{"query_priority": "1"},
+	}
+
+	dsn, err := c.FormatDSN()
+	require.NoError(t, err)
+
+	db, err := sql.Open("trino", dsn)
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		assert.NoError(t, db.Close())
+	})
+
+	result, err := db.Exec("CREATE TABLE memory.default.test (id INTEGER, name VARCHAR)")
+	require.NoError(t, err, "Failed executing CREATE TABLE query")
+
+	result, err = db.Exec("INSERT INTO memory.default.test (id, name) VALUES (?, ?), (?, ?), (?, ?)", 123, "abc", 456, "def", 789, "ghi")
+	require.NoError(t, err, "Failed executing INSERT query")
+	_, err = result.LastInsertId()
+	assert.Error(t, err, "trino: operation not supported")
+	numRows, err := result.RowsAffected()
+	require.NoError(t, err, "Failed checking rows affected")
+	assert.Equal(t, numRows, int64(3))
+
+	rows, err := db.Query("SELECT * FROM memory.default.test")
+	require.NoError(t, err, "Failed executing DELETE query")
+
+	expectedIds := []int{123, 456, 789}
+	expectedNames := []string{"abc", "def", "ghi"}
+	actualIds := []int{}
+	actualNames := []string{}
+	for rows.Next() {
+		var id int
+		var name string
+		require.NoError(t, rows.Scan(&id, &name), "Failed scanning query result")
+		actualIds = append(actualIds, id)
+		actualNames = append(actualNames, name)
+
+	}
+	assert.Equal(t, expectedIds, actualIds)
+	assert.Equal(t, expectedNames, actualNames)
+}