Add validation for SQL queries in DataFrame.readSqlQuery

dataframe-jdbc/src/main/kotlin/org/jetbrains/kotlinx/dataframe/io/readJdbc.kt

@@ -26,6 +26,21 @@ private val logger = KotlinLogging.logger {}
  */
 private const val DEFAULT_LIMIT = Int.MIN_VALUE
 
+/**
+ * Constant variable indicating the start of an SQL read query.
+ * The value of this variable is "SELECT".
+ */
+private const val START_OF_READ_SQL_QUERY = "SELECT"
+
+/**
+ * Constant representing the separator used to separate multiple SQL queries.
+ *
+ * This separator is used when multiple SQL queries need to be executed together.
+ * Each query should be separated by this separator to indicate the end of one query
+ * and the start of the next query.
+ */
+private const val MULTIPLE_SQL_QUERY_SEPARATOR = ";"
+
 /**
  * Represents a column in a database table to keep all required meta-information.
  *
@@ -179,6 +194,8 @@ public fun DataFrame.Companion.readSqlQuery(connection: Connection, sqlQuery: St
  * @see DriverManager.getConnection
  */
 public fun DataFrame.Companion.readSqlQuery(connection: Connection, sqlQuery: String, limit: Int): AnyFrame {
+    require (isValid(sqlQuery)) { "SQL query should start from SELECT and contain one query for reading data without any manipulation. "}
+
     val url = connection.metaData.url
     val dbType = extractDBTypeFromUrl(url)
 
@@ -199,6 +216,14 @@ public fun DataFrame.Companion.readSqlQuery(connection: Connection, sqlQuery: St
     }
 }
 
+/** SQL-query is accepted only if it starts from SELECT */
+private fun isValid(sqlQuery: String): Boolean {
+    val normalizedSqlQuery = sqlQuery.trim().uppercase()
+
+    return normalizedSqlQuery.startsWith(START_OF_READ_SQL_QUERY) &&
+        !normalizedSqlQuery.contains(MULTIPLE_SQL_QUERY_SEPARATOR)
+}
+
 /**
  * Reads the data from a [ResultSet] and converts it into a DataFrame.
  *

dataframe-jdbc/src/test/kotlin/org/jetbrains/kotlinx/dataframe/io/h2Test.kt

@@ -327,6 +327,7 @@ class JdbcTest {
         }
     }
 
+    // to cover a reported case from https://github.com/Kotlin/dataframe/issues/494
     @Test
     fun `repeated read from ResultSet with limit`() {
         connection.createStatement(ResultSet.TYPE_SCROLL_SENSITIVE, ResultSet.CONCUR_UPDATABLE).use { st ->
@@ -362,6 +363,88 @@ class JdbcTest {
         }
     }
 
+    // to cover a reported case from https://github.com/Kotlin/dataframe/issues/498
+    @Test
+    fun `read from incorrect SQL query`() {
+        @Language("SQL")
+        val createSQL = """
+            CREATE TABLE Orders (
+            order_id INT PRIMARY KEY,
+            customer_id INT,
+            order_date DATE,
+            total_amount DECIMAL(10, 2))
+            """
+
+
+        @Language("SQL")
+        val dropSQL = """
+            DROP TABLE Customer
+            """
+
+        @Language("SQL")
+        val alterSQL = """
+            ALTER TABLE Customer
+            ADD COLUMN email VARCHAR(100)
+            """
+
+        @Language("SQL")
+        val deleteSQL = """
+            DELETE FROM Customer
+            WHERE id = 1
+            """
+
+        @Language("SQL")
+        val repeatedSQL = """
+            SELECT * FROM Customer
+            WHERE id = 1;
+            SELECT * FROM Customer
+            WHERE id = 1;
+            """
+
+        shouldThrow<IllegalArgumentException> {
+            DataFrame.readSqlQuery(connection, createSQL)
+        }
+
+        shouldThrow<IllegalArgumentException> {
+            DataFrame.readSqlQuery(connection, dropSQL)
+        }
+
+        shouldThrow<IllegalArgumentException> {
+            DataFrame.readSqlQuery(connection, alterSQL)
+        }
+
+        shouldThrow<IllegalArgumentException> {
+            DataFrame.readSqlQuery(connection, deleteSQL)
+        }
+
+        shouldThrow<IllegalArgumentException> {
+            DataFrame.readSqlQuery(connection, repeatedSQL)
+        }
+    }
+
+    @Test
+    fun `read from table with name from reserved SQL keywords`() {
+        // Create table Sale
+        @Language("SQL")
+        val createAlterTableQuery = """
+                CREATE TABLE "ALTER" (
+                id INT PRIMARY KEY,
+                description TEXT
+                )
+            """
+
+        connection.createStatement().execute(
+            createAlterTableQuery
+        )
+
+        @Language("SQL")
+        val selectFromWeirdTableSQL = """
+            SELECT * from "ALTER"
+            """
+
+        DataFrame.readSqlQuery(connection, selectFromWeirdTableSQL).rowsCount() shouldBe 0
+    }
+
     @Test
     fun `read from non-existing jdbc url`() {
         shouldThrow<SQLException> {