@livestore/livestore 0.0.24 → 0.0.27

package/src/react/useRow.ts

@@ -0,0 +1,221 @@
+import { Schema } from '@livestore/utils/effect'
+import * as otel from '@opentelemetry/api'
+import type { SqliteDsl } from 'effect-db-schema'
+import { mapValues } from 'lodash-es'
+import React from 'react'
+
+import type { LiveStoreJSQuery } from '../reactiveQueries/js.js'
+import type { RowQueryArgs, RowResult } from '../row-query.js'
+import { rowQuery } from '../row-query.js'
+import type { DefaultSqliteTableDef, TableDef, TableOptions } from '../schema/table-def.js'
+import { useStore } from './LiveStoreContext.js'
+import { useQueryRef } from './useQuery.js'
+
+export type UseRowResult<TTableDef extends TableDef> = [
+  row: RowResult<TTableDef>,
+  setRow: StateSetters<TTableDef>,
+  query$: LiveStoreJSQuery<RowResult<TTableDef>>,
+]
+
+/**
+ * Similar to `React.useState` but returns a tuple of `[row, setRow, query$]` for a given table where ...
+ *
+ *   - `row` is the current value of the row (fully decoded according to the table schema)
+ *   - `setRow` is a function that can be used to update the row (values will be encoded according to the table schema)
+ *   - `query$` is a `LiveStoreJSQuery` that e.g. can be used to subscribe to changes to the row
+ *
+ * If the table is a singleton table, `useRow` can be called without an `id` argument. Otherwise, the `id` argument is required.
+ */
+export const useRow: {
+  <TTableDef extends TableDef<DefaultSqliteTableDef, boolean, TableOptions & { isSingleton: true }>>(
+    table: TTableDef,
+  ): UseRowResult<TTableDef>
+  <TTableDef extends TableDef<DefaultSqliteTableDef, boolean, TableOptions & { isSingleton: false }>>(
+    table: TTableDef,
+    // TODO adjust so it works with arbitrary primary keys or unique constraints
+    id: string,
+  ): UseRowResult<TTableDef>
+} = <TTableDef extends TableDef>(table: TTableDef, id?: string): UseRowResult<TTableDef> => {
+  const sqliteTableDef = table.schema
+  type TComponentState = SqliteDsl.FromColumns.RowDecoded<TTableDef['schema']['columns']>
+
+  const { store } = useStore()
+
+  const reactId = React.useId()
+
+  const { query$, otelContext } = React.useMemo(() => {
+    const cachedItem = rcCache.get(table, id ?? 'singleton')
+    if (cachedItem !== undefined) {
+      cachedItem.reactIds.add(reactId)
+      cachedItem.span.addEvent('new-subscriber', { reactId })
+
+      return {
+        query$: cachedItem.query$ as LiveStoreJSQuery<RowResult<TTableDef>>,
+        otelContext: cachedItem.otelContext,
+      }
+    }
+
+    const span = store.otel.tracer.startSpan(
+      `LiveStore:useState:${table.schema.name}${id === undefined ? '' : `:${id}`}`,
+      { attributes: { id } },
+      store.otel.queriesSpanContext,
+    )
+
+    const otelContext = otel.trace.setSpan(otel.context.active(), span)
+
+    const query$ = table.options.isSingleton
+      ? rowQuery({ table, store, otelContext } as RowQueryArgs<TTableDef>)
+      : rowQuery({ table, store, id, otelContext } as RowQueryArgs<TTableDef>)
+
+    rcCache.set(table, id ?? 'singleton', query$, reactId, otelContext, span)
+
+    return { query$, otelContext }
+  }, [table, id, reactId, store])
+
+  React.useEffect(
+    () => () => {
+      const cachedItem = rcCache.get(table, id ?? 'singleton')!
+
+      cachedItem.reactIds.delete(reactId)
+      if (cachedItem.reactIds.size === 0) {
+        rcCache.delete(cachedItem.query$)
+        cachedItem.query$.destroy()
+        cachedItem.span.end()
+      }
+    },
+    [table, id, reactId],
+  )
+
+  const query$Ref = useQueryRef(query$, otelContext)
+
+  const setState = React.useMemo<StateSetters<TTableDef>>(() => {
+    if (table.isSingleColumn) {
+      return (newValue: RowResult<TTableDef>) => {
+        if (query$Ref.current === newValue) return
+
+        const encodedValue = Schema.encodeSync(sqliteTableDef.columns['value']!.type.codec)(newValue)
+
+        store.applyEvent('livestore.UpdateComponentState', {
+          tableName: sqliteTableDef.name,
+          columnNames: ['value'],
+          id,
+          bindValues: { ['value']: encodedValue },
+        })
+      }
+    } else {
+      const setState = // TODO: do we have a better type for the values that can go in SQLite?
+        mapValues(sqliteTableDef.columns, (column, columnName) => (newValue: string | number) => {
+          // Don't update the state if it's the same as the value already seen in the component
+          // @ts-expect-error TODO fix typing
+          if (query$Ref.current[columnName] === newValue) return
+
+          const encodedValue = Schema.encodeSync(column.type.codec)(newValue)
+
+          store.applyEvent('livestore.UpdateComponentState', {
+            tableName: sqliteTableDef.name,
+            columnNames: [columnName],
+            id,
+            bindValues: { [columnName]: encodedValue },
+          })
+        })
+
+      // @ts-expect-error TODO fix typing
+      setState.setMany = (columnValues: Partial<TComponentState>) => {
+        // TODO use hashing instead
+        // Don't update the state if it's the same as the value already seen in the component
+        if (
+          // @ts-expect-error TODO fix typing
+          Object.entries(columnValues).every(([columnName, value]) => query$Ref.current[columnName] === value)
+        ) {
+          return
+        }
+
+        const columnNames = Object.keys(columnValues)
+        const bindValues = mapValues(columnValues, (value, columnName) =>
+          Schema.encodeSync(sqliteTableDef.columns[columnName]!.type.codec)(value),
+        )
+
+        store.applyEvent('livestore.UpdateComponentState', {
+          tableName: sqliteTableDef.name,
+          columnNames,
+          id,
+          bindValues,
+        })
+      }
+
+      return setState as any
+    }
+  }, [table.isSingleColumn, id, sqliteTableDef.columns, sqliteTableDef.name, store, query$Ref])
+
+  return [query$Ref.current, setState, query$]
+}
+
+export type Dispatch<A> = (action: A) => void
+export type SetStateAction<S> = S | ((previousValue: S) => S)
+
+export type StateSetters<TTableDef extends TableDef> = TTableDef['isSingleColumn'] extends true
+  ? Dispatch<SetStateAction<RowResult<TTableDef>>>
+  : {
+      [K in keyof RowResult<TTableDef>]: Dispatch<SetStateAction<RowResult<TTableDef>[K]>>
+    } & {
+      setMany: Dispatch<SetStateAction<Partial<RowResult<TTableDef>>>>
+    }
+
+/** Reference counted cache for `query$` and otel context */
+class RCCache {
+  private readonly cache = new Map<
+    TableDef,
+    Map<
+      string,
+      {
+        reactIds: Set<string>
+        span: otel.Span
+        otelContext: otel.Context
+        query$: LiveStoreJSQuery<any>
+      }
+    >
+  >()
+  private reverseCache = new Map<LiveStoreJSQuery<any>, [TableDef, string]>()
+
+  get = (table: TableDef, id: string) => {
+    const queries = this.cache.get(table)
+    if (queries === undefined) return undefined
+    return queries.get(id)
+  }
+
+  set = (
+    table: TableDef,
+    id: string,
+    query$: LiveStoreJSQuery<any>,
+    reactId: string,
+    otelContext: otel.Context,
+    span: otel.Span,
+  ) => {
+    let queries = this.cache.get(table)
+    if (queries === undefined) {
+      queries = new Map()
+      this.cache.set(table, queries)
+    }
+    queries.set(id, { query$, otelContext, span, reactIds: new Set([reactId]) })
+    this.reverseCache.set(query$, [table, id])
+  }
+
+  delete = (query$: LiveStoreJSQuery<any>) => {
+    const item = this.reverseCache.get(query$)
+    if (item === undefined) return
+
+    const [table, id] = item
+    const queries = this.cache.get(table)
+    if (queries === undefined) return
+
+    queries.delete(id)
+
+    if (queries.size === 0) {
+      this.cache.delete(table)
+    }
+
+    this.reverseCache.delete(query$)
+  }
+}
+
+const rcCache = new RCCache()

package/src/react/useTemporaryQuery.ts

@@ -1,23 +1,55 @@
 import React from 'react'
 
 import type { ILiveStoreQuery } from '../reactiveQueries/base-class.js'
-import { useQuery } from './useQuery.js'
+import { useQueryRef } from './useQuery.js'
+
+/**
+ * This is needed because the `React.useMemo` call below, can sometimes be called multiple times 🤷.
+ * The map entry is being removed again in the `React.useEffect` call below.
+ */
+const queryCache = new Map<() => ILiveStoreQuery<any>, { reactIds: Set<string>; query$: ILiveStoreQuery<any> }>()
 
 /**
  * Creates a query, subscribes and destroys it when the component unmounts.
  *
  * Make sure `makeQuery` is a memoized function.
  */
-export const useTemporaryQuery = <TResult>(makeQuery: () => ILiveStoreQuery<TResult>): TResult => {
-  // TODO cache the query outside of the `useMemo` since `useMemo` might be called multiple times
-  // also need to update the `useEffect` below https://stackoverflow.com/questions/66446642/react-usememo-memory-clean/77457605#77457605
-  const query = React.useMemo(() => makeQuery(), [makeQuery])
-
-  React.useEffect(() => {
-    return () => {
-      query.destroy()
+export const useTemporaryQuery = <TResult>(makeQuery: () => ILiveStoreQuery<TResult>): TResult =>
+  useTemporaryQueryRef(makeQuery).current
+
+export const useTemporaryQueryRef = <TResult>(
+  makeQuery: () => ILiveStoreQuery<TResult>,
+): React.MutableRefObject<TResult> => {
+  const reactId = React.useId()
+
+  const query$ = React.useMemo(() => {
+    const cachedItem = queryCache.get(makeQuery)
+    if (cachedItem !== undefined) {
+      cachedItem.reactIds.add(reactId)
+
+      return cachedItem.query$
     }
-  }, [query])
 
-  return useQuery(query)
+    const query$ = makeQuery()
+
+    queryCache.set(makeQuery, { reactIds: new Set([reactId]), query$ })
+
+    return query$
+  }, [reactId, makeQuery])
+
+  React.useEffect(
+    () => () => {
+      const cachedItem = queryCache.get(makeQuery)!
+
+      cachedItem.reactIds.delete(reactId)
+
+      if (cachedItem.reactIds.size === 0) {
+        cachedItem.query$.destroy()
+        queryCache.delete(makeQuery)
+      }
+    },
+    [makeQuery, reactId],
+  )
+
+  return useQueryRef(query$)
 }

package/src/react/utils/stack-info.ts

@@ -38,10 +38,12 @@ export const extractStackInfoFromStackTrace = (stackTrace: string): StackInfo =>
 
   while ((match = namePattern.exec(stackTrace)) !== null) {
     const [, name, filePath] = match as any as [string, string, string]
-    if (name.startsWith('use')) {
+
+    // NOTE No idea where this `Module.` comes from - possibly a Vite thing?
+    if ((name.startsWith('use') || name.startsWith('Module.use')) && name.endsWith('QueryRef') === false) {
       hasReachedStart = true
 
-      frames.unshift({ name, filePath })
+      frames.unshift({ name: name.replace(/^Module\./, ''), filePath })
     } else if (hasReachedStart) {
       // We've reached the end of the `use*` functions, so we're adding the component name and stop
       frames.unshift({ name, filePath })
@@ -57,7 +59,6 @@ export const useStackInfo = (): StackInfo =>
     Error.stackTraceLimit = 10
     // eslint-disable-next-line unicorn/error-message
     const stack = new Error().stack!
-    console.log('stack', stack)
     Error.stackTraceLimit = originalStackLimit
     return extractStackInfoFromStackTrace(stack)
   }, [])

package/src/reactive.ts

@@ -28,7 +28,7 @@ import { pick } from '@livestore/utils'
 import type * as otel from '@opentelemetry/api'
 import { isEqual, uniqueId } from 'lodash-es'
 
-import { BoundArray } from './bounded-collections.js'
+import { BoundArray } from './utils/bounded-collections.js'
 // import { getDurationMsFromSpan } from './otel.js'
 
 export const NOT_REFRESHED_YET = Symbol.for('NOT_REFRESHED_YET')
@@ -36,7 +36,7 @@ export type NOT_REFRESHED_YET = typeof NOT_REFRESHED_YET
 
 export type GetAtom = <T>(atom: Atom<T, any, any>, otelContext?: otel.Context) => T
 
-export type Ref<T, TContext, TDebugRefreshReason extends Taggable> = {
+export type Ref<T, TContext, TDebugRefreshReason extends DebugRefreshReason> = {
   _tag: 'ref'
   id: string
   isDirty: false
@@ -50,14 +50,11 @@ export type Ref<T, TContext, TDebugRefreshReason extends Taggable> = {
   equal: (a: T, b: T) => boolean
 }
 
-export type Thunk<TResult, TContext, TDebugRefreshReason extends Taggable> = {
+export type Thunk<TResult, TContext, TDebugRefreshReason extends DebugRefreshReason> = {
   _tag: 'thunk'
   id: string
   isDirty: boolean
-  computeResult: (
-    otelContext?: otel.Context,
-    debugRefreshReason?: RefreshReasonWithGenericReasons<TDebugRefreshReason>,
-  ) => TResult
+  computeResult: (otelContext?: otel.Context, debugRefreshReason?: TDebugRefreshReason) => TResult
   previousResult: TResult | NOT_REFRESHED_YET
   sub: Set<Atom<any, TContext, TDebugRefreshReason>>
   super: Set<Atom<any, TContext, TDebugRefreshReason> | Effect>
@@ -70,7 +67,7 @@ export type Thunk<TResult, TContext, TDebugRefreshReason extends Taggable> = {
   __getResult: any
 }
 
-export type Atom<T, TContext, TDebugRefreshReason extends Taggable> =
+export type Atom<T, TContext, TDebugRefreshReason extends DebugRefreshReason> =
   | Ref<T, TContext, TDebugRefreshReason>
   | Thunk<T, TContext, TDebugRefreshReason>
 
@@ -82,13 +79,23 @@ export type Effect = {
   label?: string
 }
 
-export type Taggable<T extends string = string> = { _tag: T }
-
 export type DebugThunkInfo<T extends string = string> = {
   _tag: T
   durationMs: number
 }
 
+export type DebugRefreshReasonBase =
+  /** Usually in response to some `applyEvent`/`applyEvents` with `skipRefresh: true` */
+  | {
+      _tag: 'runDeferredEffects'
+      originalRefreshReasons?: ReadonlyArray<DebugRefreshReasonBase>
+      manualRefreshReason?: DebugRefreshReasonBase
+    }
+  | { _tag: 'makeThunk'; label?: string }
+  | { _tag: 'unknown' }
+
+export type DebugRefreshReason<T extends string = string> = DebugRefreshReasonBase | { _tag: T }
+
 export type ReactiveGraphOptions = {
   effectsWrapper?: (runEffects: () => void) => void
 }
@@ -100,7 +107,7 @@ export type AtomDebugInfo<TDebugThunkInfo extends DebugThunkInfo> = {
 }
 
 // TODO possibly find a better name for "refresh"
-export type RefreshDebugInfo<TDebugRefreshReason extends Taggable, TDebugThunkInfo extends DebugThunkInfo> = {
+export type RefreshDebugInfo<TDebugRefreshReason extends DebugRefreshReason, TDebugThunkInfo extends DebugThunkInfo> = {
   /** Currently only used for easier handling in React (e.g. as key) */
   id: string
   reason: TDebugRefreshReason
@@ -112,15 +119,7 @@ export type RefreshDebugInfo<TDebugRefreshReason extends Taggable, TDebugThunkIn
   graphSnapshot: ReactiveGraphSnapshot
 }
 
-export type RefreshReasonWithGenericReasons<T extends Taggable> =
-  | T
-  | {
-      _tag: 'makeThunk'
-      label?: string
-    }
-  | { _tag: 'unknown' }
-
-export const unknownRefreshReason = () => {
+const unknownRefreshReason = () => {
   // debugger
   return { _tag: 'unknown' as const }
 }
@@ -128,34 +127,29 @@ export const unknownRefreshReason = () => {
 export type SerializedAtom = Readonly<
   PrettifyFlat<
     Pick<Atom<unknown, unknown, any>, '_tag' | 'id' | 'label' | 'meta'> & {
-      sub: string[]
-      super: string[]
+      sub: ReadonlyArray<string>
+      super: ReadonlyArray<string>
     }
   >
 >
 
-export type SerializedEffect = Readonly<PrettifyFlat<Pick<Effect, '_tag' | 'id'>>>
-
 type ReactiveGraphSnapshot = {
-  readonly atoms: SerializedAtom[]
-  // readonly effects: SerializedEffect[]
-  /** IDs of atoms and effects that are dirty */
-  // readonly dirtyNodes: string[]
+  readonly atoms: ReadonlyArray<SerializedAtom>
+  /** IDs of deferred effects */
+  readonly deferredEffects: ReadonlyArray<string>
 }
 
 const uniqueNodeId = () => uniqueId('node-')
 const uniqueRefreshInfoId = () => uniqueId('refresh-info-')
 
 const serializeAtom = (atom: Atom<any, unknown, any>): SerializedAtom => ({
-  ...pick(atom, ['_tag', 'id', 'label', 'meta']),
+  ...pick(atom, ['_tag', 'id', 'label', 'meta', 'isDirty']),
   sub: Array.from(atom.sub).map((a) => a.id),
   super: Array.from(atom.super).map((a) => a.id),
 })
 
-// const serializeEffect = (effect: Effect): SerializedEffect => pick(effect, ['_tag', 'id'])
-
 export class ReactiveGraph<
-  TDebugRefreshReason extends Taggable,
+  TDebugRefreshReason extends DebugRefreshReason,
   TDebugThunkInfo extends DebugThunkInfo,
   TContext = {},
 > {
@@ -164,11 +158,13 @@ export class ReactiveGraph<
 
   context: TContext | undefined
 
-  debugRefreshInfos: BoundArray<
-    RefreshDebugInfo<RefreshReasonWithGenericReasons<TDebugRefreshReason>, TDebugThunkInfo>
-  > = new BoundArray(5000)
+  debugRefreshInfos: BoundArray<RefreshDebugInfo<TDebugRefreshReason, TDebugThunkInfo>> = new BoundArray(5000)
+
+  private currentDebugRefresh:
+    | { refreshedAtoms: AtomDebugInfo<TDebugThunkInfo>[]; startMs: DOMHighResTimeStamp }
+    | undefined
 
-  currentDebugRefresh: { refreshedAtoms: AtomDebugInfo<TDebugThunkInfo>[]; startMs: DOMHighResTimeStamp } | undefined
+  private deferredEffects: Map<Effect, Set<TDebugRefreshReason>> = new Map()
 
   constructor(options: ReactiveGraphOptions) {
     this.effectsWrapper = options?.effectsWrapper ?? ((runEffects: () => void) => runEffects())
@@ -208,8 +204,6 @@ export class ReactiveGraph<
           label?: string
           meta?: any
           equal?: (a: T, b: T) => boolean
-          /** Debug info for initializing the thunk (i.e. running it the first time) */
-          // debugRefreshReason?: RefreshReasonWithGenericReasons<TDebugRefreshReason>
         }
       | undefined,
   ): Thunk<T, TContext, TDebugRefreshReason> {
@@ -264,16 +258,15 @@ export class ReactiveGraph<
             const durationMs = performance.now() - this.currentDebugRefresh!.startMs
             this.currentDebugRefresh = undefined
 
-            const refreshDebugInfo = {
+            this.debugRefreshInfos.push({
               id: uniqueRefreshInfoId(),
-              reason: debugRefreshReason ?? { _tag: 'makeThunk', label: options?.label },
+              reason: debugRefreshReason ?? ({ _tag: 'makeThunk', label: options?.label } as TDebugRefreshReason),
               skippedRefresh: false,
               refreshedAtoms,
               durationMs,
               completedTimestamp: Date.now(),
               graphSnapshot: this.getSnapshot(),
-            }
-            this.debugRefreshInfos.push(refreshDebugInfo)
+            })
           }
 
           return result
@@ -308,7 +301,9 @@ export class ReactiveGraph<
       this.removeEdge(node, subComp)
     }
 
-    if (node._tag !== 'effect') {
+    if (node._tag === 'effect') {
+      this.deferredEffects.delete(node)
+    } else {
       this.atoms.delete(node)
     }
   }
@@ -345,23 +340,20 @@ export class ReactiveGraph<
     val: T,
     options?:
       | {
+          skipRefresh?: boolean
           debugRefreshReason?: TDebugRefreshReason
           otelContext?: otel.Context
         }
       | undefined,
   ) {
-    ref.previousResult = val
-
-    const effectsToRefresh = new Set<Effect>()
-    markSuperCompDirtyRec(ref, effectsToRefresh)
-
-    this.runEffects(effectsToRefresh, options)
+    this.setRefs([[ref, val]], options)
   }
 
   setRefs<T>(
     refs: [Ref<T, TContext, TDebugRefreshReason>, T][],
     options?:
       | {
+          skipRefresh?: boolean
           debugRefreshReason?: TDebugRefreshReason
           otelContext?: otel.Context
         }
@@ -374,17 +366,30 @@ export class ReactiveGraph<
       markSuperCompDirtyRec(ref, effectsToRefresh)
     }
 
-    this.runEffects(effectsToRefresh, options)
+    if (options?.skipRefresh) {
+      for (const effect of effectsToRefresh) {
+        if (this.deferredEffects.has(effect) === false) {
+          this.deferredEffects.set(effect, new Set())
+        }
+
+        if (options?.debugRefreshReason !== undefined) {
+          this.deferredEffects.get(effect)!.add(options.debugRefreshReason)
+        }
+      }
+    } else {
+      this.runEffects(effectsToRefresh, {
+        debugRefreshReason: options?.debugRefreshReason ?? (unknownRefreshReason() as TDebugRefreshReason),
+        otelContext: options?.otelContext,
+      })
+    }
   }
 
   private runEffects = (
     effectsToRefresh: Set<Effect>,
-    options?:
-      | {
-          debugRefreshReason?: TDebugRefreshReason
-          otelContext?: otel.Context
-        }
-      | undefined,
+    options: {
+      debugRefreshReason: TDebugRefreshReason
+      otelContext?: otel.Context
+    },
   ) => {
     this.effectsWrapper(() => {
       this.currentDebugRefresh = { refreshedAtoms: [], startMs: performance.now() }
@@ -399,7 +404,7 @@ export class ReactiveGraph<
 
       const refreshDebugInfo: RefreshDebugInfo<TDebugRefreshReason, TDebugThunkInfo> = {
         id: uniqueRefreshInfoId(),
-        reason: options?.debugRefreshReason ?? (unknownRefreshReason() as TDebugRefreshReason),
+        reason: options.debugRefreshReason,
         skippedRefresh: false,
         refreshedAtoms,
         durationMs,
@@ -410,6 +415,22 @@ export class ReactiveGraph<
     })
   }
 
+  runDeferredEffects = (options?: { debugRefreshReason?: TDebugRefreshReason; otelContext?: otel.Context }) => {
+    // TODO improve how refresh reasons are propagated for deferred effect execution
+    // TODO also improve "batching" of running deferred effects (i.e. in a single `this.runEffects` call)
+    // but need to be careful to not overwhelm the main thread
+    for (const [effect, debugRefreshReasons] of this.deferredEffects) {
+      this.runEffects(new Set([effect]), {
+        debugRefreshReason: {
+          _tag: 'runDeferredEffects',
+          originalRefreshReasons: Array.from(debugRefreshReasons) as ReadonlyArray<DebugRefreshReasonBase>,
+          manualRefreshReason: options?.debugRefreshReason,
+        } as TDebugRefreshReason,
+        otelContext: options?.otelContext,
+      })
+    }
+  }
+
   addEdge(
     superComp: Atom<any, TContext, TDebugRefreshReason> | Effect,
     subComp: Atom<any, TContext, TDebugRefreshReason>,
@@ -426,17 +447,12 @@ export class ReactiveGraph<
     subComp.super.delete(superComp)
   }
 
-  private getSnapshot = (): ReactiveGraphSnapshot => ({
+  getSnapshot = (): ReactiveGraphSnapshot => ({
     atoms: Array.from(this.atoms).map(serializeAtom),
-    // effects: Array.from(this.effects).map(serializeEffect),
-    // dirtyNodes: Array.from(this.dirtyNodes).map((a) => a.id),
+    deferredEffects: Array.from(this.deferredEffects.keys()).map((_) => _.id),
   })
 }
 
-// const isAtom = <T, TContext>(a: Atom<T, TContext> | Effect): a is Atom<T, TContext> =>
-//   a._tag === 'ref' || a._tag === 'thunk'
-// const isEffect = <T, TContext>(a: Atom<T, TContext> | Effect): a is Effect => a._tag === 'effect'
-
 const compute = <T>(atom: Atom<T, unknown, any>, otelContext: otel.Context): T => {
   // const __getResult = atom._tag === 'thunk' ? atom.__getResult.toString() : ''
   if (atom.isDirty) {
@@ -462,6 +478,6 @@ const markSuperCompDirtyRec = <T>(atom: Atom<T, unknown, any>, effectsToRefresh:
   }
 }
 
-const throwContextNotSetError = (): never => {
+export const throwContextNotSetError = (): never => {
   throw new Error(`LiveStore Error: \`context\` not set on ReactiveGraph`)
 }

package/src/reactiveQueries/base-class.ts

@@ -1,11 +1,17 @@
 import type * as otel from '@opentelemetry/api'
 
+import { dbGraph } from '../global-state.js'
 import type { StackInfo } from '../react/utils/stack-info.js'
-import type { Atom, GetAtom, RefreshReasonWithGenericReasons, Thunk } from '../reactive.js'
-import type { RefreshReason } from '../store.js'
-import { type DbContext, dbGraph } from './graph.js'
+import { type Atom, type GetAtom, throwContextNotSetError, type Thunk } from '../reactive.js'
+import type { RefreshReason, Store } from '../store.js'
 import type { LiveStoreJSQuery } from './js.js'
 
+export type DbContext = {
+  store: Store
+  otelTracer: otel.Tracer
+  rootOtelContext: otel.Context
+}
+
 export type UnsubscribeQuery = () => void
 
 let queryIdCounter = 0
@@ -18,7 +24,7 @@ export interface ILiveStoreQuery<TResult> {
 
   label: string
 
-  run: (otelContext?: otel.Context, debugRefreshReason?: RefreshReasonWithGenericReasons<RefreshReason>) => TResult
+  run: (otelContext?: otel.Context, debugRefreshReason?: RefreshReason) => TResult
 
   destroy(): void
 
@@ -41,13 +47,10 @@ export abstract class LiveStoreQueryBase<TResult> implements ILiveStoreQuery<TRe
 
   abstract destroy: () => void
 
-  run = (otelContext?: otel.Context, debugRefreshReason?: RefreshReasonWithGenericReasons<RefreshReason>): TResult =>
+  run = (otelContext?: otel.Context, debugRefreshReason?: RefreshReason): TResult =>
     this.results$.computeResult(otelContext, debugRefreshReason)
 
-  runAndDestroy = (
-    otelContext?: otel.Context,
-    debugRefreshReason?: RefreshReasonWithGenericReasons<RefreshReason>,
-  ): TResult => {
+  runAndDestroy = (otelContext?: otel.Context, debugRefreshReason?: RefreshReason): TResult => {
     const result = this.run(otelContext, debugRefreshReason)
     this.destroy()
     return result
@@ -57,7 +60,8 @@ export abstract class LiveStoreQueryBase<TResult> implements ILiveStoreQuery<TRe
     onNewValue: (value: TResult) => void,
     onUnsubsubscribe?: () => void,
     options?: { label?: string; otelContext?: otel.Context } | undefined,
-  ): (() => void) => dbGraph.context!.store.subscribe(this, onNewValue, onUnsubsubscribe, options)
+  ): (() => void) =>
+    dbGraph.context?.store.subscribe(this, onNewValue, onUnsubsubscribe, options) ?? throwContextNotSetError()
 }
 
 export type GetAtomResult = <T>(atom: Atom<T, any, RefreshReason> | LiveStoreJSQuery<T>) => T