<commit_before>import {defer, partition, merge, of, throwError, asyncScheduler, Observable} from 'rxjs'
import {mergeMap, throttleTime, share, take, filter} from 'rxjs/operators'
import {exhaustMapToWithTrailing} from 'rxjs-exhaustmap-with-trailing'
import {getVersionedClient} from '../../client/versionedClient'
import {ReconnectEvent, WelcomeEvent, MutationEvent} from './types'

type Params = Record<string, string | number | boolean | string[]>

export interface ListenQueryOptions {
  tag?: string
  apiVersion?: string
  transitions?: ('update' | 'appear' | 'disappear')[]
}

const fetch = (query: string, params: Params, options: ListenQueryOptions) =>
  defer(() =>
    getVersionedClient(options.apiVersion).observable.fetch(query, params, {
      tag: options.tag,
      filterResponse: true,
    })
  )

const listen = (query: string, params: Params, options: ListenQueryOptions) =>
  defer(() =>
    getVersionedClient(options.apiVersion).listen(query, params, {
      events: ['welcome', 'mutation', 'reconnect'],
      includeResult: false,
      visibility: 'query',
      tag: options.tag,
    })
  ) as Observable<ReconnectEvent | WelcomeEvent | MutationEvent>

function isWelcomeEvent(
  event: MutationEvent | ReconnectEvent | WelcomeEvent
): event is WelcomeEvent {
  return event.type === 'welcome'
}

// todo: promote as building block for better re-use
// todo: optimize by patching collection in-place
export const listenQuery = (
  query: string | {fetch: string; listen: string},
  params: Params = {},
  options: ListenQueryOptions = {}
) => {
  const fetchQuery = typeof query === 'string' ? query : query.fetch
  const listenerQuery = typeof query === 'string' ? query : query.listen

  const fetchOnce$ = fetch(fetchQuery, params, options)

  const events$ = listen(listenerQuery, params, options).pipe(
    mergeMap((ev, i) => {
      const isFirst = i === 0
      if (isFirst && !isWelcomeEvent(ev)) {
        // if the first event is not welcome, it is most likely a reconnect and
        // if it's not a reconnect something is very wrong
        return throwError(
          new Error(
            ev.type === 'reconnect'
              ? 'Could not establish EventSource connection'
              : `Received unexpected type of first event "${ev.type}"`
          )
        )
      }
      return of(ev)
    }),
    share()
  )

  const [welcome$, mutationAndReconnect$] = partition(events$, isWelcomeEvent)
  const isRelevantEvent = (event: MutationEvent | ReconnectEvent | WelcomeEvent): boolean => {
    if (!options.transitions || event.type !== 'mutation') {
      return true
    }

    return options.transitions.includes(event.transition)
  }

  return merge(
    welcome$.pipe(take(1)),
    mutationAndReconnect$.pipe(
      filter(isRelevantEvent),
      throttleTime(1000, asyncScheduler, {leading: true, trailing: true})
    )
  ).pipe(exhaustMapToWithTrailing(fetchOnce$))
}<commit_msg>feat(base): allow listenQuery to specify custom throttle time<commit_after>import {defer, partition, merge, of, throwError, asyncScheduler, Observable} from 'rxjs'
import {mergeMap, throttleTime, share, take, filter} from 'rxjs/operators'
import {exhaustMapToWithTrailing} from 'rxjs-exhaustmap-with-trailing'
import {getVersionedClient} from '../../client/versionedClient'
import {ReconnectEvent, WelcomeEvent, MutationEvent} from './types'

type Params = Record<string, string | number | boolean | string[]>

export interface ListenQueryOptions {
  tag?: string
  apiVersion?: string
  throttleTime?: number
  transitions?: ('update' | 'appear' | 'disappear')[]
}

const fetch = (query: string, params: Params, options: ListenQueryOptions) =>
  defer(() =>
    getVersionedClient(options.apiVersion).observable.fetch(query, params, {
      tag: options.tag,
      filterResponse: true,
    })
  )

const listen = (query: string, params: Params, options: ListenQueryOptions) =>
  defer(() =>
    getVersionedClient(options.apiVersion).listen(query, params, {
      events: ['welcome', 'mutation', 'reconnect'],
      includeResult: false,
      visibility: 'query',
      tag: options.tag,
    })
  ) as Observable<ReconnectEvent | WelcomeEvent | MutationEvent>

function isWelcomeEvent(
  event: MutationEvent | ReconnectEvent | WelcomeEvent
): event is WelcomeEvent {
  return event.type === 'welcome'
}

// todo: promote as building block for better re-use
// todo: optimize by patching collection in-place
export const listenQuery = (
  query: string | {fetch: string; listen: string},
  params: Params = {},
  options: ListenQueryOptions = {}
) => {
  const fetchQuery = typeof query === 'string' ? query : query.fetch
  const listenerQuery = typeof query === 'string' ? query : query.listen

  const fetchOnce$ = fetch(fetchQuery, params, options)

  const events$ = listen(listenerQuery, params, options).pipe(
    mergeMap((ev, i) => {
      const isFirst = i === 0
      if (isFirst && !isWelcomeEvent(ev)) {
        // if the first event is not welcome, it is most likely a reconnect and
        // if it's not a reconnect something is very wrong
        return throwError(
          new Error(
            ev.type === 'reconnect'
              ? 'Could not establish EventSource connection'
              : `Received unexpected type of first event "${ev.type}"`
          )
        )
      }
      return of(ev)
    }),
    share()
  )

  const [welcome$, mutationAndReconnect$] = partition(events$, isWelcomeEvent)
  const isRelevantEvent = (event: MutationEvent | ReconnectEvent | WelcomeEvent): boolean => {
    if (!options.transitions || event.type !== 'mutation') {
      return true
    }

    return options.transitions.includes(event.transition)
  }

  return merge(
    welcome$.pipe(take(1)),
    mutationAndReconnect$.pipe(
      filter(isRelevantEvent),
      throttleTime(options.throttleTime || 1000, asyncScheduler, {leading: true, trailing: true})
    )
  ).pipe(exhaustMapToWithTrailing(fetchOnce$))
}
{'fraction_non_alphanumeric': 0.12370097217566209, 'fraction_numerical': 0.002011397921555481, 'mean_word_length': 3.6291698991466252, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '4228395', 'n_tokens_mistral': 1742, 'n_tokens_neox': 1701, 'n_words': 486}