fix: Zetasql analysis fragment

We cant discard parts of the string, since we are indexing into it. Store the offset instead. This becomes a problem when we match more than 1 table.
nrkno · Oct 9, 2023 · 7730a89 · 7730a89
1 parent ac41cbe
commit 7730a89
Showing 1 changed file with 13 additions and 10 deletions.
diff --git a/zetasql/src/main/scala/no/nrk/bigquery/ZetaSql.scala b/zetasql/src/main/scala/no/nrk/bigquery/ZetaSql.scala
@@ -6,9 +6,9 @@
 
 package no.nrk.bigquery
 
-import cats.syntax.all._
+import cats.syntax.all.*
 import cats.effect.Sync
-import no.nrk.bigquery.syntax._
+import no.nrk.bigquery.syntax.*
 import com.google.zetasql.{
   AnalyzerOptions,
   ParseLocationRange,
@@ -29,8 +29,8 @@ import com.google.zetasql.parser.{ASTNodes, ParseTreeVisitor}
 import com.google.zetasql.toolkit.{AnalysisException, AnalyzedStatement, ZetaSQLToolkitAnalyzer}
 
 import scala.collection.mutable.ListBuffer
-import scala.jdk.CollectionConverters._
-import scala.jdk.OptionConverters._
+import scala.jdk.CollectionConverters.*
+import scala.jdk.OptionConverters.*
 
 class ZetaSql[F[_]](implicit F: Sync[F]) {
   import ZetaSql._
@@ -56,26 +56,29 @@ class ZetaSql[F[_]](implicit F: Sync[F]) {
     def evalFragments(
         parsedTables: List[(BQTableId, ParseLocationRange)]
     ): BQSqlFrag = {
-      val asString = query
       val found = allTables
         .flatMap(table =>
           parsedTables.flatMap { case (id, range) => if (eqv(table.tableId, id)) List(table -> range) else Nil })
         .distinct
-      val (rest, aggregate) = found.foldLeft((asString, BQSqlFrag.Empty)) { case ((input, agg), (t, loc)) =>
-        val frag = agg ++ BQSqlFrag.Frag(input.substring(0, loc.start() - 1)) ++ toFragment(t)
-        val rest = input.substring(loc.end())
+      val (rest, aggregate) = found.foldLeft((0, Vector.empty[BQSqlFrag])) { case ((offset, agg), (t, loc)) =>
+        val frag = agg ++ List(
+          BQSqlFrag.Frag(query.substring(offset, loc.start() - 1)),
+          BQSqlFrag.Frag(" "),
+          toFragment(t))
+        val rest = loc.end()
         rest -> frag
       }
+      val str = query.substring(rest)
 
-      aggregate ++ BQSqlFrag.Frag(rest)
+      BQSqlFrag.Combined(if (str.nonEmpty) aggregate :+ BQSqlFrag.Frag(str) else aggregate)
     }
 
     parseScript(BQSqlFrag.Frag(query))
       .flatMap(F.fromEither)
       .flatMap { script =>
         val list = script.getStatementListNode.getStatementList
         if (list.size() != 1) {
-          Sync[F].raiseError[List[(no.nrk.bigquery.BQTableId, com.google.zetasql.ParseLocationRange)]](
+          Sync[F].raiseError[List[(BQTableId, ParseLocationRange)]](
             new IllegalArgumentException("Expects only one statement"))
         } else
           Sync[F].delay {