IntelliJ/DataGrip Data Extractors

SQL-IN.groovy

/*
 * Available context bindings:
 *   COLUMNS     List<DataColumn>
 *   ROWS        Iterable<DataRow>
 *   OUT         { append() }
 *   FORMATTER   { format(row, col); formatValue(Object, col) }
 *   TRANSPOSED  Boolean
 * plus ALL_COLUMNS, TABLE, DIALECT
 *
 * where:
 *   DataRow     { rowNumber(); first(); last(); data(): List<Object>; value(column): Object }
 *   DataColumn  { columnNumber(), name() }
 */

SEPARATOR = ","
QUOTE     = "\""
NEWLINE   = System.getProperty("line.separator")


KEYWORDS_LOWERCASE = com.intellij.database.util.DbSqlUtil.areKeywordsLowerCase(PROJECT)
KW_NULL = KEYWORDS_LOWERCASE ? "null" : "NULL"


def printRow = { values, valueToString ->
  values.eachWithIndex { value, idx ->
    def str = valueToString(value)
    def q = str.contains(SEPARATOR) || str.contains(QUOTE) || str.contains(NEWLINE)
    OUT.append(q ? QUOTE : "")
      .append(str.replace(QUOTE, QUOTE + QUOTE))
      .append(q ? QUOTE : "")
      .append(idx != values.size() - 1 ? SEPARATOR : NEWLINE)
  }
}


def formatSingleValue = {dataRow, column ->
  def RET = new StringBuilder()
  def value = dataRow.value(column)
  def skipQuote = value.toString().isNumber() || value == null
  def stringValue = value != null ? FORMATTER.format(dataRow, column) : KW_NULL
  if (DIALECT.getFamilyId().isMysql()) stringValue = stringValue.replace("\\", "\\\\")
  RET.append(skipQuote ? "": QUOTE).append(stringValue.replace(QUOTE, QUOTE + QUOTE))
          .append(skipQuote ? "": QUOTE)
  return RET.toString()
}

def printCol = { col, colIdx, values ->

  OUT
    .append(col.name())
    .append(" IN (")
  def isLastCol = colIdx == COLUMNS.size() - 1;

  values.eachWithIndex { value, idx ->
    def isLastRow = idx == values.size() - 1
    OUT.append(value)
    if (!isLastRow) {
      OUT.append(", ")
    }
  }

  OUT.append(")")
  if (!isLastCol) {
    OUT.append(" AND ")
  }
}

if (!TRANSPOSED) {
  // ROWS.each { row -> printRow(COLUMNS, { FORMATTER.format(row, it) }) }
  def byColumn = COLUMNS.collect { new ArrayList<String>()}
  // byColumn[0].add(byColumn[0])
  ROWS.eachWithIndex { row, rowIdx ->
    COLUMNS.eachWithIndex {col, i ->
      def value = formatSingleValue(row, col)
      print(value)
      byColumn[i].add(value.toString())
    }
  }
  COLUMNS.eachWithIndex { column, colNumber -> printCol(column, colNumber, byColumn[colNumber]) }
}
else {
  def values = COLUMNS.collect { new ArrayList<String>() }
  ROWS.each { row -> COLUMNS.eachWithIndex { col, i -> values[i].add(FORMATTER.format(row, col)) } }
  values.each { printRow(it, { it }) }
}