SchemaRegistry.ts

/* eslint-disable prettier/prettier */
import { Response } from 'mappersmith'
import { types } from 'avsc'

import { encode, MAGIC_BYTE } from './encoder'
import decode from './decoder'
import { COMPATIBILITY, DEFAULT_SEPERATOR } from './constants'
import API, {
  SchemaRegistryAPIClientArgs,
  SchemaRegistryAPIClientOptions,
  SchemaRegistryAPIClient,
} from './api'
import Cache from './cache'
import {
  ConfluentSchemaRegistryError,
  ConfluentSchemaRegistryArgumentError,
  ConfluentSchemaRegistryCompatibilityError,
} from './errors'
import { RawSchema, Schema, SchemaReference } from './@types'

interface RegisteredSchema {
  id: number
}

interface Opts {
  compatibility?: COMPATIBILITY
  separator?: string
  subject?: string
}

const DEFAULT_OPTS = {
  compatibility: COMPATIBILITY.BACKWARD,
  separator: DEFAULT_SEPERATOR,
}

export default class SchemaRegistry {
  private api: SchemaRegistryAPIClient
  private cacheMissRequests: { [key: number]: Promise<Response> } = {}

  public cache: Cache

  constructor(
    { auth, clientId, host, retry }: SchemaRegistryAPIClientArgs,
    options?: SchemaRegistryAPIClientOptions,
  ) {
    this.api = API({ auth, clientId, host, retry })
    this.cache = new Cache(options?.forSchemaOptions)
  }

  public async register(schema: RawSchema, userOpts?: Opts): Promise<RegisteredSchema> {
    const { compatibility, separator } = { ...DEFAULT_OPTS, ...userOpts }

    if (!schema.name) {
      throw new ConfluentSchemaRegistryArgumentError(`Invalid name: ${schema.name}`)
    }

    let subject: string
    if (userOpts && userOpts.subject) {
      subject = userOpts.subject
    } else {
      if (!schema.namespace) {
        throw new ConfluentSchemaRegistryArgumentError(`Invalid namespace: ${schema.namespace}`)
      }

      subject = [schema.namespace, schema.name].join(separator)
    }

    try {
      const response = await this.api.Subject.config({ subject })
      const { compatibilityLevel }: { compatibilityLevel: COMPATIBILITY } = response.data()

      if (compatibilityLevel.toUpperCase() !== compatibility) {
        throw new ConfluentSchemaRegistryCompatibilityError(
          `Compatibility does not match the configuration (${compatibility} != ${compatibilityLevel.toUpperCase()})`,
        )
      }
    } catch (error) {
      if (error.status !== 404) {
        throw error
      }

      if (compatibility) {
        await this.api.Subject.updateConfig({ subject, body: { compatibility } })
      }
    }

    const response = await this.api.Subject.register({
      subject,
      body: { schema: JSON.stringify(schema) },
    })

    const registeredSchema: RegisteredSchema = response.data()
    this.cache.setLatestRegistryId(subject, registeredSchema.id)
    this.cache.setSchema(registeredSchema.id, schema)

    return registeredSchema
  }

  public async getSchema(registryId: number): Promise<Schema> {
    const schema = this.cache.getSchema(registryId)

    if (schema) {
      return schema
    }

    const response = await this.getSchemaOriginRequest(registryId)
    const foundSchema: { schema: string; references?: Array<SchemaReference> } = response.data()
    const rawSchema: RawSchema = JSON.parse(foundSchema.schema)
    let logicalTypes: Record<string, new () => types.LogicalType> | undefined
    if (foundSchema.references) {
      logicalTypes = Object.fromEntries(
        await Promise.all(
          foundSchema.references.map(async ({ name, subject, version }) => {
            const schemaType = await this.getSchema(await this.getRegistryId(subject, version))

            return [ name, schemaType ];
          }),
        ),
      )
    }
    return this.cache.setSchema(registryId, rawSchema, logicalTypes)
  }

  public async encode(registryId: number, jsonPayload: any): Promise<Buffer> {
    if (!registryId) {
      throw new ConfluentSchemaRegistryArgumentError(
        `Invalid registryId: ${JSON.stringify(registryId)}`,
      )
    }

    const schema = await this.getSchema(registryId)

    return encode(schema, registryId, jsonPayload)
  }

  public async decode(buffer: Buffer): Promise<any> {
    if (!Buffer.isBuffer(buffer)) {
      throw new ConfluentSchemaRegistryArgumentError('Invalid buffer')
    }

    const { magicByte, registryId, payload } = decode(buffer)
    if (Buffer.compare(MAGIC_BYTE, magicByte) !== 0) {
      throw new ConfluentSchemaRegistryArgumentError(
        `Message encoded with magic byte ${JSON.stringify(magicByte)}, expected ${JSON.stringify(
          MAGIC_BYTE,
        )}`,
      )
    }

    const schema = await this.getSchema(registryId)

    return schema.fromBuffer(payload)
  }

  public async getRegistryId(subject: string, version: number | string): Promise<number> {
    const cached = this.cache.getRegistryIdBySchemaRef({subject, version});
    if (cached) {
        return cached;
    }
    const response = await this.api.Subject.version({ subject, version })
    const { id }: { id: number } = response.data()
    this.cache.setRegistryIdBySchemaRef({subject, version}, id);

    return id
  }

  public async getRegistryIdBySchema(subject: string, schema: Schema): Promise<number> {
    try {
      const response = await this.api.Subject.registered({
        subject,
        body: { schema: JSON.stringify(schema) },
      })
      const { id }: { id: number } = response.data()

      return id
    } catch (error) {
      if (error.status && error.status === 404) {
        throw new ConfluentSchemaRegistryError(error)
      }

      throw error
    }
  }

  public async getLatestSchemaId(subject: string): Promise<number> {
    const response = await this.api.Subject.latestVersion({ subject })
    const { id }: { id: number } = response.data()

    return id
  }

  private getSchemaOriginRequest(registryId: number) {
    // ensure that cache-misses result in a single origin request
    if (this.cacheMissRequests[registryId]) {
      return this.cacheMissRequests[registryId]
    } else {
      const request = this.api.Schema.find({ id: registryId }).finally(() => {
        delete this.cacheMissRequests[registryId]
      })

      this.cacheMissRequests[registryId] = request

      return request
    }
  }
}