@livestore/livestore 0.0.19 → 0.0.22

package/src/store.ts

@@ -1,34 +1,29 @@
-import type { TypedDocumentNode as DocumentNode } from '@graphql-typed-document-node/core'
 import { assertNever, makeNoopSpan, makeNoopTracer, shouldNeverHappen } from '@livestore/utils'
 import { identity } from '@livestore/utils/effect'
 import * as otel from '@opentelemetry/api'
 import type { GraphQLSchema } from 'graphql'
-import * as graphql from 'graphql'
-import { uniqueId } from 'lodash-es'
-import * as ReactDOM from 'react-dom'
 import type * as Sqlite from 'sqlite-esm'
 import { v4 as uuid } from 'uuid'
 
 import type { ComponentKey } from './componentKey.js'
 import { tableNameForComponentKey } from './componentKey.js'
-import type { QueryDefinition } from './effect/LiveStore.js'
 import type { LiveStoreEvent } from './events.js'
 import { InMemoryDatabase } from './inMemoryDatabase.js'
 import { migrateDb } from './migrations.js'
 import { getDurationMsFromSpan } from './otel.js'
-import type { Atom, Ref } from './reactive.js'
-import { ReactiveGraph } from './reactive.js'
-import { LiveStoreGraphQLQuery } from './reactiveQueries/graphql.js'
-import { LiveStoreJSQuery } from './reactiveQueries/js.js'
-import { LiveStoreSQLQuery } from './reactiveQueries/sql.js'
+import type { StackInfo } from './react/utils/stack-info.js'
+import type { ReactiveGraph, Ref } from './reactive.js'
+import type { ILiveStoreQuery } from './reactiveQueries/base-class.js'
+import { type DbContext, dbGraph } from './reactiveQueries/graph.js'
+import type { LiveStoreGraphQLQuery } from './reactiveQueries/graphql.js'
+import type { LiveStoreJSQuery } from './reactiveQueries/js.js'
+import type { LiveStoreSQLQuery } from './reactiveQueries/sql.js'
 import type { ActionDefinition, GetActionArgs, Schema, SQLWriteStatement } from './schema.js'
 import { componentStateTables } from './schema.js'
 import type { Storage, StorageInit } from './storage/index.js'
-import type { Bindable, ParamsObject } from './util.js'
+import type { ParamsObject } from './util.js'
 import { isPromise, prepareBindValues, sql } from './util.js'
 
-export type GetAtomResult = <T>(atom: Atom<T> | LiveStoreJSQuery<T>) => T
-
 export type LiveStoreQuery<TResult extends Record<string, any> = any> =
   | LiveStoreSQLQuery<TResult>
   | LiveStoreJSQuery<TResult>
@@ -48,8 +43,6 @@ export type QueryResult<TQuery> = TQuery extends LiveStoreSQLQuery<infer R>
   ? Readonly<Result>
   : never
 
-const globalComponentKey: ComponentKey = { _tag: 'singleton', componentName: '__global', id: 'singleton' }
-
 export type GraphQLOptions<TContext> = {
   schema: GraphQLSchema
   makeContext: (db: InMemoryDatabase, tracer: otel.Tracer) => TContext
@@ -93,9 +86,21 @@ export type RefreshReason =
       _tag: 'makeThunk'
       label?: string
     }
+  | {
+      _tag: 'react'
+      api: string
+      label?: string
+      stackInfo?: StackInfo
+    }
+  | { _tag: 'manual'; label?: string }
   | { _tag: 'unknown' }
 
-export type QueryDebugInfo = { _tag: 'graphql' | 'sql' | 'js' | 'unknown'; label: string; query: string }
+export type QueryDebugInfo = {
+  _tag: 'graphql' | 'sql' | 'js' | 'unknown'
+  label: string
+  query: string
+  durationMs: number
+}
 
 export type StoreOtel = {
   tracer: otel.Tracer
@@ -104,7 +109,7 @@ export type StoreOtel = {
 }
 
 export class Store<TGraphQLContext extends BaseGraphQLContext = BaseGraphQLContext> {
-  graph: ReactiveGraph<RefreshReason, QueryDebugInfo>
+  graph: ReactiveGraph<RefreshReason, QueryDebugInfo, DbContext>
   inMemoryDB: InMemoryDatabase
   // TODO refactor
   _proxyDb: InMemoryDatabase
@@ -116,10 +121,11 @@ export class Store<TGraphQLContext extends BaseGraphQLContext = BaseGraphQLConte
    * Note we're using `Ref<null>` here as we don't care about the value but only about *that* something has changed.
    * This only works in combination with `equal: () => false` which will always trigger a refresh.
    */
-  tableRefs: { [key: string]: Ref<null> }
-  activeQueries: Set<LiveStoreQuery>
+  tableRefs: { [key: string]: Ref<null, DbContext, RefreshReason> }
+
+  /** RC-based set to see which queries are currently subscribed to */
+  activeQueries: ReferenceCountedSet<LiveStoreQuery>
   storage?: Storage
-  temporaryQueries: Set<LiveStoreQuery> | undefined
 
   private constructor({
     db,
@@ -132,16 +138,10 @@ export class Store<TGraphQLContext extends BaseGraphQLContext = BaseGraphQLConte
   }: StoreOptions<TGraphQLContext>) {
     this.inMemoryDB = db
     this._proxyDb = dbProxy
-    this.graph = new ReactiveGraph({
-      // TODO move this into React module
-      // Do all our updates inside a single React setState batch to avoid multiple UI re-renders
-      effectsWrapper: (run) => ReactDOM.unstable_batchedUpdates(() => run()),
-      otelTracer,
-    })
     this.schema = schema
     // TODO generalize the `tableRefs` concept to allow finer-grained refs
     this.tableRefs = {}
-    this.activeQueries = new Set()
+    this.activeQueries = new ReferenceCountedSet()
     this.storage = storage
 
     const applyEventsSpan = otelTracer.startSpan('LiveStore:applyEvents', {}, otelRootSpanContext)
@@ -150,6 +150,10 @@ export class Store<TGraphQLContext extends BaseGraphQLContext = BaseGraphQLConte
     const queriesSpan = otelTracer.startSpan('LiveStore:queries', {}, otelRootSpanContext)
     const otelQueriesSpanContext = otel.trace.setSpan(otel.context.active(), queriesSpan)
 
+    // TODO allow passing in a custom graph
+    this.graph = dbGraph
+    this.graph.context = { store: this, otelTracer, rootOtelContext: otelQueriesSpanContext }
+
     this.otel = {
       tracer: otelTracer,
       applyEventsSpanContext: otelApplyEventsSpanContext,
@@ -189,403 +193,61 @@ export class Store<TGraphQLContext extends BaseGraphQLContext = BaseGraphQLConte
     })
   }
 
-  /**
-   * Creates a reactive LiveStore SQL query
-   *
-   * NOTE The query is actually running (even if no one has subscribed to it yet) and will be kept up to date.
-   */
-  querySQL = <TResult>(
-    genQueryString: string | ((get: GetAtomResult) => string),
-    {
-      queriedTables,
-      bindValues,
-      componentKey,
-      label,
-      otelContext = otel.context.active(),
-    }: {
-      /**
-       * List of tables that are queried in this query;
-       * used to determine reactive dependencies.
-       *
-       * NOTE In the future we want to auto-generate this via parsing the query
-       */
-      queriedTables: string[]
-      bindValues?: Bindable | undefined
-      componentKey?: ComponentKey | undefined
-      label?: string | undefined
-      otelContext?: otel.Context
-    },
-  ): LiveStoreSQLQuery<TResult> =>
-    this.otel.tracer.startActiveSpan(
-      'querySQL', // NOTE span name will be overridden further down
-      { attributes: { label } },
-      otelContext,
-      (span) => {
-        const otelContext = otel.trace.setSpan(otel.context.active(), span)
-
-        const queryString$ = this.graph.makeThunk(
-          (get, addDebugInfo) => {
-            if (typeof genQueryString === 'function') {
-              const getAtom: GetAtomResult = (atom) => {
-                if (atom._tag === 'thunk' || atom._tag === 'ref') return get(atom)
-                return get(atom.results$)
-              }
-              const queryString = genQueryString(getAtom)
-              addDebugInfo({ _tag: 'js', label: `${label}:queryString`, query: queryString })
-              return queryString
-            } else {
-              return genQueryString
-            }
-          },
-          { label: `${label}:queryString`, meta: { liveStoreThunkType: 'sqlQueryString' } },
-          otelContext,
-        )
-
-        label = label ?? queryString$.result
-        span.updateName(`querySQL:${label}`)
-
-        const queryLabel = `${label}:results` + (this.temporaryQueries ? ':temp' : '')
-
-        const results$ = this.graph.makeThunk<ReadonlyArray<TResult>>(
-          (get, addDebugInfo) =>
-            this.otel.tracer.startActiveSpan(
-              'sql:', // NOTE span name will be overridden further down
-              {},
-              otelContext,
-              (span) => {
-                try {
-                  const otelContext = otel.trace.setSpan(otel.context.active(), span)
-
-                  // Establish a reactive dependency on the tables used in the query
-                  for (const tableName of queriedTables) {
-                    const tableRef =
-                      this.tableRefs[tableName] ?? shouldNeverHappen(`No table ref found for ${tableName}`)
-                    get(tableRef)
-                  }
-                  const sqlString = get(queryString$)
-
-                  span.setAttribute('sql.query', sqlString)
-                  span.updateName(`sql:${sqlString.slice(0, 50)}`)
-
-                  const results = this.inMemoryDB.select<TResult>(sqlString, {
-                    queriedTables,
-                    bindValues: bindValues ? prepareBindValues(bindValues, sqlString) : undefined,
-                    otelContext,
-                  })
-
-                  span.setAttribute('sql.rowsCount', results.length)
-                  addDebugInfo({ _tag: 'sql', label: label ?? '', query: sqlString })
-
-                  return results
-                } finally {
-                  span.end()
-                }
-              },
-            ),
-          { label: queryLabel },
-          otelContext,
-        )
-
-        const query = new LiveStoreSQLQuery<TResult>({
-          label,
-          queryString$,
-          results$,
-          componentKey: componentKey ?? globalComponentKey,
-          store: this,
-          otelContext,
-        })
-
-        this.activeQueries.add(query)
-
-        // TODO get rid of temporary query workaround
-        if (this.temporaryQueries !== undefined) {
-          this.temporaryQueries.add(query)
-        }
-
-        // NOTE we are not ending the span here but in the query `destroy` method
-        return query
-      },
-    )
-
-  queryJS = <TResult>(
-    genResults: (get: GetAtomResult) => TResult,
-    {
-      componentKey = globalComponentKey,
-      label = `js${uniqueId()}`,
-      otelContext = otel.context.active(),
-    }: { componentKey?: ComponentKey; label?: string; otelContext?: otel.Context },
-  ): LiveStoreJSQuery<TResult> =>
-    this.otel.tracer.startActiveSpan(`queryJS:${label}`, { attributes: { label } }, otelContext, (span) => {
-      const otelContext = otel.trace.setSpan(otel.context.active(), span)
-      const queryLabel = `${label}:results` + (this.temporaryQueries ? ':temp' : '')
-      const results$ = this.graph.makeThunk(
-        (get, addDebugInfo) => {
-          const getAtom: GetAtomResult = (atom) => {
-            if (atom._tag === 'thunk' || atom._tag === 'ref') return get(atom)
-            return get(atom.results$)
-          }
-          addDebugInfo({ _tag: 'js', label, query: genResults.toString() })
-          return genResults(getAtom)
-        },
-        { label: queryLabel, meta: { liveStoreThunkType: 'jsResults' } },
-        otelContext,
-      )
-
-      const query = new LiveStoreJSQuery<TResult>({
-        label,
-        results$,
-        componentKey,
-        store: this,
-        otelContext,
-      })
-
-      this.activeQueries.add(query)
-
-      // TODO get rid of temporary query workaround
-      if (this.temporaryQueries !== undefined) {
-        this.temporaryQueries.add(query)
-      }
-
-      // NOTE we are not ending the span here but in the query `destroy` method
-      return query
-    })
-
-  queryGraphQL = <TResult extends Record<string, any>, TVariableValues extends Record<string, any>>(
-    document: DocumentNode<TResult, TVariableValues>,
-    genVariableValues: TVariableValues | ((get: GetAtomResult) => TVariableValues),
-    {
-      componentKey,
-      label,
-      otelContext = otel.context.active(),
-    }: {
-      componentKey: ComponentKey
-      label?: string
-      otelContext?: otel.Context
-    },
-  ): LiveStoreGraphQLQuery<TResult, TVariableValues, TGraphQLContext> =>
-    this.otel.tracer.startActiveSpan(
-      `queryGraphQL:`, // NOTE span name will be overridden further down
-      {},
-      otelContext,
-      (span) => {
-        const otelContext = otel.trace.setSpan(otel.context.active(), span)
-
-        if (this.graphQLContext === undefined) {
-          return shouldNeverHappen("Can't run a GraphQL query on a store without GraphQL context")
-        }
-
-        const labelWithDefault = label ?? graphql.getOperationAST(document)?.name?.value ?? 'graphql'
-
-        span.updateName(`queryGraphQL:${labelWithDefault}`)
-
-        const variableValues$ = this.graph.makeThunk(
-          (get) => {
-            if (typeof genVariableValues === 'function') {
-              const getAtom: GetAtomResult = (atom) => {
-                if (atom._tag === 'thunk' || atom._tag === 'ref') return get(atom)
-                return get(atom.results$)
-              }
-              return genVariableValues(getAtom)
-            } else {
-              return genVariableValues
-            }
-          },
-          { label: `${labelWithDefault}:variableValues`, meta: { liveStoreThunkType: 'graphqlVariableValues' } },
-          otelContext,
-        )
-
-        const resultsLabel = `${labelWithDefault}:results` + (this.temporaryQueries ? ':temp' : '')
-        const results$ = this.graph.makeThunk<TResult>(
-          (get, addDebugInfo) => {
-            const variableValues = get(variableValues$)
-            const { result, queriedTables } = this.queryGraphQLOnce(document, variableValues, otelContext)
-
-            // Add dependencies on any tables that were used
-            for (const tableName of queriedTables) {
-              const tableRef = this.tableRefs[tableName]
-              assertNever(tableRef !== undefined, `No table ref found for ${tableName}`)
-              get(tableRef!)
-            }
-
-            addDebugInfo({ _tag: 'graphql', label: resultsLabel, query: graphql.print(document) })
-
-            return result
-          },
-          { label: resultsLabel, meta: { liveStoreThunkType: 'graphqlResults' } },
-          otelContext,
-        )
-
-        const query = new LiveStoreGraphQLQuery({
-          document,
-          context: this.graphQLContext,
-          results$,
-          componentKey,
-          label: labelWithDefault,
-          store: this,
-          otelContext,
-        })
-
-        this.activeQueries.add(query)
-
-        // TODO get rid of temporary query workaround
-        if (this.temporaryQueries !== undefined) {
-          this.temporaryQueries.add(query)
-        }
-
-        // NOTE we are not ending the span here but in the query `destroy` method
-        return query
-      },
-    )
-
-  queryGraphQLOnce = <TResult extends Record<string, any>, TVariableValues extends Record<string, any>>(
-    document: DocumentNode<TResult, TVariableValues>,
-    variableValues: TVariableValues,
-    otelContext: otel.Context = this.otel.queriesSpanContext,
-  ): { result: TResult; queriedTables: string[] } => {
-    const schema =
-      this.graphQLSchema ?? shouldNeverHappen("Can't run a GraphQL query on a store without GraphQL schema")
-    const context =
-      this.graphQLContext ?? shouldNeverHappen("Can't run a GraphQL query on a store without GraphQL context")
-    const tracer = this.otel.tracer
-
-    const operationName = graphql.getOperationAST(document)?.name?.value
-
-    return tracer.startActiveSpan(`executeGraphQLQuery: ${operationName}`, {}, otelContext, (span) => {
-      try {
-        span.setAttribute('graphql.variables', JSON.stringify(variableValues))
-        span.setAttribute('graphql.query', graphql.print(document))
-
-        context.queriedTables.clear()
-
-        context.otelContext = otel.trace.setSpan(otel.context.active(), span)
-
-        const res = graphql.executeSync({
-          document,
-          contextValue: context,
-          schema: schema,
-          variableValues,
-        })
-
-        // TODO track number of nested SQL queries via Otel + debug info
-
-        if (res.errors) {
-          span.setStatus({ code: otel.SpanStatusCode.ERROR, message: 'GraphQL error' })
-          span.setAttribute('graphql.error', res.errors.join('\n'))
-          span.setAttribute('graphql.error-detail', JSON.stringify(res.errors))
-          console.error(`graphql error (${operationName})`, res.errors)
-        }
-
-        return { result: res.data as unknown as TResult, queriedTables: Array.from(context.queriedTables.values()) }
-      } finally {
-        span.end()
-      }
-    })
-  }
-
   /**
    * Subscribe to the results of a query
    * Returns a function to cancel the subscription.
    */
-  subscribe = <TQuery extends LiveStoreQuery>(
-    query: TQuery,
-    onNewValue: (value: QueryResult<TQuery>) => void,
-    onSubsubscribe?: () => void,
-    options?: { label?: string } | undefined,
+  subscribe = <TResult>(
+    query: ILiveStoreQuery<TResult>,
+    onNewValue: (value: TResult) => void,
+    onUnsubsubscribe?: () => void,
+    options?: { label?: string; otelContext?: otel.Context; skipInitialRun?: boolean } | undefined,
   ): (() => void) =>
     this.otel.tracer.startActiveSpan(
       `LiveStore.subscribe`,
       { attributes: { label: options?.label } },
-      query.otelContext,
+      options?.otelContext ?? this.otel.queriesSpanContext,
       (span) => {
         const otelContext = otel.trace.setSpan(otel.context.active(), span)
 
-        const effect = this.graph.makeEffect(
-          (get) => {
-            const result = get(query.results$) as QueryResult<TQuery>
-            onNewValue(result)
-          },
-          { label: `subscribe:${options?.label}` },
-          otelContext,
-        )
+        const label = `subscribe:${options?.label}`
+        const effect = this.graph.makeEffect((get) => onNewValue(get(query.results$)), { label })
+
+        this.activeQueries.add(query as LiveStoreQuery)
 
-        const subscriptionKey = uuid()
+        // Running effect right away to get initial value (unless `skipInitialRun` is set)
+        if (options?.skipInitialRun !== true) {
+          effect.doEffect(otelContext)
+        }
 
         const unsubscribe = () => {
           try {
             this.graph.destroy(effect)
-            query.activeSubscriptions.delete(subscriptionKey)
-            onSubsubscribe?.()
+            this.activeQueries.remove(query as LiveStoreQuery)
+            onUnsubsubscribe?.()
           } finally {
             span.end()
           }
         }
 
-        query.activeSubscriptions.set(subscriptionKey, unsubscribe)
-
         return unsubscribe
       },
     )
 
-  /**
-   * Any queries created in the callback will be destroyed when the callback is complete.
-   * Useful for temporarily creating reactive queries, which is an idempotent operation
-   * that can be safely called inside a React useMemo hook.
-   */
-  inTempQueryContext = <TResult>(callback: () => TResult): TResult => {
-    this.temporaryQueries = new Set()
-    // TODO: consider errors / try/finally here?
-    const result = callback()
-    for (const query of this.temporaryQueries) {
-      this.destroyQuery(query)
-    }
-    this.temporaryQueries = undefined
-    return result
-  }
-
   /**
    * Destroys the entire store, including all queries and subscriptions.
    *
    * Currently only used when shutting down the app for debugging purposes (e.g. to close Otel spans).
    */
   destroy = () => {
-    for (const query of this.activeQueries) {
-      this.destroyQuery(query)
-    }
-
     Object.values(this.tableRefs).forEach((tableRef) => this.graph.destroy(tableRef))
 
-    const applyEventsSpan = otel.trace.getSpan(this.otel.applyEventsSpanContext)!
-    applyEventsSpan.end()
-
-    const queriesSpan = otel.trace.getSpan(this.otel.queriesSpanContext)!
-    queriesSpan.end()
+    otel.trace.getSpan(this.otel.applyEventsSpanContext)!.end()
+    otel.trace.getSpan(this.otel.queriesSpanContext)!.end()
 
     // TODO destroy active subscriptions
   }
 
-  private destroyQuery = (query: LiveStoreQuery) => {
-    if (query._tag === 'sql') {
-      // results are downstream of query string, so will automatically be destroyed together
-      this.graph.destroy(query.queryString$)
-    } else {
-      this.graph.destroy(query.results$)
-    }
-    this.activeQueries.delete(query)
-    query.destroy()
-  }
-
-  /**
-   * Clean up queries and downstream subscriptions associated with a component.
-   * This is critical to avoid memory leaks.
-   */
-  unmountComponent = (componentKey: ComponentKey) => {
-    for (const query of this.activeQueries) {
-      if (query.componentKey === componentKey) {
-        this.destroyQuery(query)
-      }
-    }
-  }
-
   /* Apply a single write event to the store, and refresh all queries in response */
   applyEvent = <TEventType extends string & keyof LiveStoreActionDefinitionsTypes>(
     eventType: TEventType,
@@ -607,27 +269,32 @@ export class Store<TGraphQLContext extends BaseGraphQLContext = BaseGraphQLConte
           const otelContext = otel.trace.setSpan(otel.context.active(), span)
           const writeTables = this.applyEventWithoutRefresh(eventType, args, otelContext).writeTables
 
-          const tablesToUpdate = [] as [Ref<null>, null][]
+          const tablesToUpdate = [] as [Ref<null, DbContext, RefreshReason>, null][]
           for (const tableName of writeTables) {
             const tableRef = this.tableRefs[tableName]
             assertNever(tableRef !== undefined, `No table ref found for ${tableName}`)
             tablesToUpdate.push([tableRef!, null])
           }
 
+          const debugRefreshReason = {
+            _tag: 'applyEvent' as const,
+            event: { type: eventType, args },
+            writeTables: [...writeTables],
+          }
+
           // Update all table refs together in a batch, to only trigger one reactive update
-          this.graph.setRefs(
-            tablesToUpdate,
-            {
-              otelHint: 'applyEvents',
-              skipRefresh,
-              debugRefreshReason: {
-                _tag: 'applyEvent',
-                event: { type: eventType, args },
-                writeTables: [...writeTables],
-              },
-            },
-            otelContext,
-          )
+          this.graph.setRefs(tablesToUpdate, { debugRefreshReason, otelContext })
+
+          if (skipRefresh === false) {
+            // TODO update the graph
+            // this.graph.refresh(
+            //   {
+            //     otelHint: 'applyEvents',
+            //     debugRefreshReason,
+            //   },
+            //   otelContext,
+            // )
+          }
         } catch (e: any) {
           span.setStatus({ code: otel.SpanStatusCode.ERROR, message: e.toString() })
 
@@ -703,27 +370,25 @@ export class Store<TGraphQLContext extends BaseGraphQLContext = BaseGraphQLConte
             },
           )
 
-          const tablesToUpdate = [] as [Ref<null>, null][]
+          const tablesToUpdate = [] as [Ref<null, DbContext, RefreshReason>, null][]
           for (const tableName of writeTables) {
             const tableRef = this.tableRefs[tableName]
             assertNever(tableRef !== undefined, `No table ref found for ${tableName}`)
             tablesToUpdate.push([tableRef!, null])
           }
 
+          const debugRefreshReason = {
+            _tag: 'applyEvents' as const,
+            events: [...events].map((e) => ({ type: e.eventType, args: e.args })),
+            writeTables: [...writeTables],
+          }
           // Update all table refs together in a batch, to only trigger one reactive update
-          this.graph.setRefs(
-            tablesToUpdate,
-            {
-              otelHint: 'applyEvents',
-              skipRefresh,
-              debugRefreshReason: {
-                _tag: 'applyEvents',
-                events: [...events].map((e) => ({ type: e.eventType, args: e.args })),
-                writeTables: [...writeTables],
-              },
-            },
-            otelContext,
-          )
+          this.graph.setRefs(tablesToUpdate, { debugRefreshReason, otelContext })
+
+          if (skipRefresh === false) {
+            // TODO update the graph
+            // this.graph.refresh({ debugRefreshReason, otelHint: 'applyEvents' }, otelContext)
+          }
         } catch (e: any) {
           span.setStatus({ code: otel.SpanStatusCode.ERROR, message: e.toString() })
         } finally {
@@ -746,20 +411,14 @@ export class Store<TGraphQLContext extends BaseGraphQLContext = BaseGraphQLConte
       { attributes: { 'livestore.manualRefreshLabel': label } },
       this.otel.applyEventsSpanContext,
       (span) => {
-        const otelContext = otel.trace.setSpan(otel.context.active(), span)
-        this.graph.refresh({ otelHint: 'manualRefresh', debugRefreshReason: { _tag: 'manualRefresh' } }, otelContext)
+        // const otelContext = otel.trace.setSpan(otel.context.active(), span)
+        // TODO update the graph
+        // this.graph.refresh({ otelHint: 'manualRefresh', debugRefreshReason: { _tag: 'manualRefresh' } }, otelContext)
         span.end()
       },
     )
   }
 
-  // TODO get rid of this as part of new query definition approach https://www.notion.so/schickling/New-query-definition-approach-1097a78ef0e9495bac25f90417374756?pvs=4
-  runOnce = <TQueryDef extends QueryDefinition>(queryDef: TQueryDef): QueryResult<ReturnType<TQueryDef>> => {
-    return this.inTempQueryContext(() => {
-      return queryDef(this).results$.result
-    })
-  }
-
   /**
    * Apply an event to the store.
    * Returns the tables that were affected by the event.
@@ -853,8 +512,8 @@ export class Store<TGraphQLContext extends BaseGraphQLContext = BaseGraphQLConte
    * This should only be used for framework-internal purposes;
    * all app writes should go through applyEvent.
    */
-  execute = async (query: string, params: ParamsObject = {}, writeTables?: string[]) => {
-    this.inMemoryDB.execute(query, prepareBindValues(params, query), writeTables)
+  execute = (query: string, params: ParamsObject = {}, writeTables?: string[], otelContext?: otel.Context) => {
+    this.inMemoryDB.execute(query, prepareBindValues(params, query), writeTables, { otelContext })
 
     if (this.storage !== undefined) {
       const parentSpan = otel.trace.getSpan(otel.context.active())
@@ -977,3 +636,39 @@ const eventToSql = (
 
   return { statement, bindValues }
 }
+
+class ReferenceCountedSet<T> {
+  private map: Map<T, number>
+
+  constructor() {
+    this.map = new Map<T, number>()
+  }
+
+  add = (key: T) => {
+    const count = this.map.get(key) ?? 0
+    this.map.set(key, count + 1)
+  }
+
+  remove = (key: T) => {
+    const count = this.map.get(key) ?? 0
+    if (count === 1) {
+      this.map.delete(key)
+    } else {
+      this.map.set(key, count - 1)
+    }
+  }
+
+  has = (key: T) => {
+    return this.map.has(key)
+  }
+
+  get size() {
+    return this.map.size
+  }
+
+  *[Symbol.iterator]() {
+    for (const key of this.map.keys()) {
+      yield key
+    }
+  }
+}