<commit_before>/*
 * Copyright 2019 Spotify AB.
 *
 * Licensed 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 com.spotify.scio.sql

import com.spotify.scio.values.SCollection
import com.spotify.scio.schemas.Schema
import org.apache.beam.sdk.values.TupleTag

import scala.language.experimental.macros
import scala.reflect.macros.whitebox
import com.spotify.scio.schemas.SchemaMacroHelpers

trait SQLBuilder {
  def as[B: Schema]: SCollection[B]
}

sealed trait SqlParam

final case class SCollectionRef[A: Schema](coll: SCollection[A]) extends SqlParam {
  type _A = A
  val schema = Schema[A]
}

final case class UdfRef(udf: Udf) extends SqlParam

final class SqlInterpolator(private val sc: StringContext) extends AnyVal {

  private def paramToString(
    tags: Map[String, (SCollectionRef[_], TupleTag[_])]
  )(p: SqlParam): String =
    p match {
      case SCollectionRef(scoll) =>
        tags(scoll.name)._2.getId
      case UdfRef(udf) =>
        udf.fnName
    }

  def sql(p0: SqlParam, ps: SqlParam*): SQLBuilder = {
    val params = p0 :: ps.toList

    val tags =
      params.zipWithIndex.collect {
        case (ref @ SCollectionRef(scoll), i) =>
          (scoll.name, (ref, new TupleTag[ref._A](s"SCOLLECTION_$i")))
      }.toMap

    val udfs = params.collect { case UdfRef(u) => u }
    val strings = sc.parts.iterator
    val expressions = params.iterator
    val toString = paramToString(tags) _

    val expr = expressions.map(toString)
    val q =
      strings.zipAll(expr, "", "").foldLeft("") { case (a, (x, y)) => s"${a}${x} ${y}" }

    SQLBuilders.from(q, tags.values.toList, udfs)
  }

  def tsql(ps: Any*): SQLBuilder =
    macro SqlInterpolatorMacro.builder
}

private trait SqlInterpolatorMacroHelpers {
  val ctx: whitebox.Context
  import ctx.universe._

  def partsFromContext: List[Tree] = {
    ctx.prefix.tree match {
      case Apply(_, Apply(_, xs: List[_]) :: Nil) => xs<commit_msg>Fix sql interpolator tag index (#1961)<commit_after>/*
 * Copyright 2019 Spotify AB.
 *
 * Licensed 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 com.spotify.scio.sql

import com.spotify.scio.values.SCollection
import com.spotify.scio.schemas.Schema
import org.apache.beam.sdk.values.TupleTag

import scala.language.experimental.macros
import scala.reflect.macros.whitebox
import com.spotify.scio.schemas.SchemaMacroHelpers

trait SQLBuilder {
  def as[B: Schema]: SCollection[B]
}

sealed trait SqlParam

final case class SCollectionRef[A: Schema](coll: SCollection[A]) extends SqlParam {
  type _A = A
  val schema = Schema[A]
}

final case class UdfRef(udf: Udf) extends SqlParam

final class SqlInterpolator(private val sc: StringContext) extends AnyVal {

  private def paramToString(
    tags: Map[String, (SCollectionRef[_], TupleTag[_])]
  )(p: SqlParam): String =
    p match {
      case SCollectionRef(scoll) =>
        tags(scoll.name)._2.getId
      case UdfRef(udf) =>
        udf.fnName
    }

  def sql(p0: SqlParam, ps: SqlParam*): SQLBuilder = {
    val params = p0 :: ps.toList

    val tags = params
      .collect { case ref @ SCollectionRef(scoll) => (scoll, ref) }
      .distinct
      .zipWithIndex
      .map {
        case ((scoll, ref), i) =>
          (scoll.name, (ref, new TupleTag[ref._A](s"SCOLLECTION_$i")))
      }
      .toMap

    val udfs = params.collect { case UdfRef(u) => u }
    val strings = sc.parts.iterator
    val expressions = params.iterator
    val toString = paramToString(tags) _

    val expr = expressions.map(toString)
    val q =
      strings.zipAll(expr, "", "").foldLeft("") { case (a, (x, y)) => s"${a}${x} ${y}" }

    SQLBuilders.from(q, tags.values.toList, udfs)
  }

  def tsql(ps: Any*): SQLBuilder =
    macro SqlInterpolatorMacro.builder
}

private trait SqlInterpolatorMacroHelpers {
  val ctx: whitebox.Context
  import ctx.universe._

  def partsFromContext: List[Tree] = {
    ctx.prefix.tree match {
      case Apply(_, Apply(_, xs: List[_]) :: Nil) => xs
{'fraction_non_alphanumeric': 0.11352220212979706, 'fraction_numerical': 0.005224030540486237, 'mean_word_length': 3.424888888888889, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 16, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6079186', 'n_tokens_mistral': 1618, 'n_tokens_neox': 1578, 'n_words': 477}