trpc test setup (#6872)

packages/trpc-server/.prettierrc

@@ -0,0 +1,6 @@
+{
+  "trailingComma": "es5",
+  "tabWidth": 2,
+  "semi": false,
+  "singleQuote": true
+}

packages/trpc-server/README.md

@@ -0,0 +1,4 @@
+## Tests
+
+* run once: `bash test.sh run`
+* start watch mode: `bash test.sh`

packages/trpc-server/docker-compose.yml

@@ -0,0 +1,30 @@
+version: '3.1'
+
+services:
+
+  db:
+    image: postgres
+    restart: always
+    environment:
+      POSTGRES_PASSWORD: testing
+    ports:
+      - 35764:5432
+    volumes:
+      - ../discovery-provider/ddl:/ddl
+
+  elasticsearch:
+    image: docker.elastic.co/elasticsearch/elasticsearch:8.10.2
+    environment:
+      - network.host=0.0.0.0
+      - discovery.type=single-node
+      - cluster.name=docker-cluster
+      - node.name=cluster1-node1
+      - xpack.license.self_generated.type=basic
+      - xpack.security.enabled=false
+      - 'ES_JAVA_OPTS=-Xms512m -Xmx512m'
+    ulimits:
+      memlock:
+        soft: -1
+        hard: -1
+    ports:
+      - 35765:9200

packages/trpc-server/package.json

@@ -5,11 +5,13 @@
   "main": "src/index.ts",
   "scripts": {
     "build": "tsc",
-    "dev": "DEBUG=audius:* tsx watch src/index.ts",
+    "dev": "DEBUG=audius:* tsx watch src/server.ts",
     "db:gen": "tsx ./tools/db-gen.ts",
-    "start": "tsx src/index.ts"
+    "start": "tsx src/server.ts",
+    "test": "vitest"
   },
   "dependencies": {
+    "@elastic/elasticsearch": "^8.10.0",
     "@trpc/server": "10.38.4",
     "cors": "2.8.5",
     "dataloader": "2.2.2",
@@ -28,6 +30,8 @@
     "@types/swagger-ui-express": "^4.1.4",
     "eslint": "^8.40.0",
     "tsx": "^3.12.7",
-    "typescript": "^5.1.3"
+    "typescript": "^5.1.3",
+    "vite-node": "^0.34.6",
+    "vitest": "^0.34.6"
   }
-}
+}

packages/trpc-server/src/index.ts

@@ -1,29 +1,18 @@
-import 'dotenv/config'
-
-import { createExpressMiddleware } from '@trpc/server/adapters/express'
-import cors from 'cors'
-import express from 'express'
-import swaggerUi from 'swagger-ui-express'
-import {
-  createOpenApiExpressMiddleware,
-  generateOpenApiDocument
-} from 'trpc-openapi'
-import { userRouter } from './routers/user-router'
-import { createContext, publicProcedure, router } from './trpc'
+import { inferRouterInputs, inferRouterOutputs } from '@trpc/server'
 import { meRouter } from './routers/me-router'
-import { trackRouter } from './routers/track-router'
 import { playlistRouter } from './routers/playlist-router'
-import { inferRouterInputs, inferRouterOutputs } from '@trpc/server'
-
-const app = express()
-app.use(cors())
+import { searchRouter } from './routers/search-router'
+import { trackRouter } from './routers/track-router'
+import { userRouter } from './routers/user-router'
+import { publicProcedure, router } from './trpc'
 
 // AppRouter
-const appRouter = router({
+export const appRouter = router({
   me: meRouter,
   users: userRouter,
   tracks: trackRouter,
   playlists: playlistRouter,
+  search: searchRouter,
 
   version: publicProcedure.query(() => ({
     version: '0.0.2'
@@ -34,28 +23,3 @@ export type AppRouter = typeof appRouter
 
 export type RouterInput = inferRouterInputs<AppRouter>
 export type RouterOutput = inferRouterOutputs<AppRouter>
-
-// endpoints
-app.use('/trpc', createExpressMiddleware({ router: appRouter, createContext }))
-app.use(
-  '/rest',
-  createOpenApiExpressMiddleware({ router: appRouter, createContext })
-)
-
-// OpenAPI schema document
-export const openApiDocument = generateOpenApiDocument(appRouter, {
-  title: 'Example CRUD API',
-  description: 'OpenAPI compliant REST API built using tRPC with Express',
-  version: '1.0.0',
-  baseUrl: 'http://localhost:2022/rest',
-  docsUrl: 'https://docs.audius.co'
-})
-
-// Swagger UI
-app.use('/', swaggerUi.serve)
-app.get('/', swaggerUi.setup(openApiDocument))
-
-const port = 2022
-app.listen(port, () => {
-  console.log('listening on ', port)
-})

packages/trpc-server/src/routers/search-router.ts

@@ -0,0 +1,109 @@
+import { z } from 'zod'
+import { publicProcedure, router } from '../trpc'
+import { TRPCError } from '@trpc/server'
+import { Client as ES } from '@elastic/elasticsearch'
+
+const esc = new ES({ node: process.env.audius_elasticsearch_url })
+
+export const searchRouter = router({
+  users: publicProcedure
+    .input(
+      z.object({
+        q: z.string(),
+        onlyFollowed: z.boolean().default(false),
+        limit: z.number().default(20),
+        cursor: z.string().default('0'),
+      })
+    )
+    .query(async ({ ctx, input }) => {
+      const followedBy = input.onlyFollowed ? ctx.currentUserId : undefined
+      const found = await esc.search({
+        index: 'users',
+        query: userSearchDSL(input.q, followedBy) as any,
+        _source: false,
+      })
+      const ids = found.hits.hits.map((h) => h._id)
+      return ids
+    }),
+
+  tracks: publicProcedure
+    .input(
+      z.object({
+        q: z.string(),
+        onlySaved: z.boolean().default(false),
+        limit: z.number().default(20),
+      })
+    )
+    .query(async ({ ctx, input }) => {
+      const savedBy = input.onlySaved ? ctx.currentUserId : undefined
+      const found = await esc.search({
+        index: 'tracks',
+        query: trackSearchDSL(input.q, savedBy) as any,
+        _source: false,
+      })
+      const ids = found.hits.hits.map((h) => h._id)
+      return ids
+    }),
+})
+
+export function userSearchDSL(q: string, followedBy?: number) {
+  const dsl: any = {
+    bool: {
+      must: [suggestDSL(q), { term: { is_deactivated: { value: false } } }],
+      must_not: [{ exists: { field: 'stem_of' } }],
+      should: [],
+    },
+  }
+
+  if (followedBy) {
+    dsl.bool.must.push({
+      terms: {
+        _id: {
+          index: 'users',
+          id: followedBy.toString(),
+          path: 'following_ids',
+        },
+      },
+    })
+  }
+
+  return dsl
+}
+
+function trackSearchDSL(q: string, savedByUserId?: number) {
+  const dsl: any = {
+    bool: {
+      must: [
+        suggestDSL(q),
+        { term: { is_unlisted: { value: false } } },
+        { term: { is_delete: false } },
+      ],
+      must_not: [],
+      should: [{ term: { is_verified: { value: true } } }],
+    },
+  }
+
+  if (savedByUserId) {
+    dsl.bool.must.push({
+      term: { saved_by: { value: savedByUserId, boost: 1.2 } },
+    })
+  }
+
+  return dsl
+}
+
+function suggestDSL(
+  q: string,
+  operator: string = 'or',
+  extraFields: string[] = []
+) {
+  return {
+    multi_match: {
+      query: q,
+      fields: ['suggest', 'suggest._2gram', 'suggest._3gram', ...extraFields],
+      operator: operator as any,
+      type: 'bool_prefix',
+      fuzziness: 'AUTO',
+    },
+  }
+}

packages/trpc-server/src/server.ts

@@ -0,0 +1,40 @@
+import 'dotenv/config'
+
+import { createExpressMiddleware } from '@trpc/server/adapters/express'
+import cors from 'cors'
+import express from 'express'
+import swaggerUi from 'swagger-ui-express'
+import {
+  createOpenApiExpressMiddleware,
+  generateOpenApiDocument
+} from 'trpc-openapi'
+import { appRouter } from './index'
+import { createContext } from './trpc'
+
+const app = express()
+app.use(cors())
+
+// endpoints
+app.use('/trpc', createExpressMiddleware({ router: appRouter, createContext }))
+app.use(
+  '/rest',
+  createOpenApiExpressMiddleware({ router: appRouter, createContext })
+)
+
+// OpenAPI schema document
+export const openApiDocument = generateOpenApiDocument(appRouter, {
+  title: 'Example CRUD API',
+  description: 'OpenAPI compliant REST API built using tRPC with Express',
+  version: '1.0.0',
+  baseUrl: 'http://localhost:2022/rest',
+  docsUrl: 'https://docs.audius.co'
+})
+
+// Swagger UI
+app.use('/', swaggerUi.serve)
+app.get('/', swaggerUi.setup(openApiDocument))
+
+const port = 2022
+app.listen(port, () => {
+  console.log('listening on ', port)
+})

packages/trpc-server/test.sh

@@ -0,0 +1,23 @@
+#! /bin/bash
+set -e
+
+# start postgres + elasticsearch
+docker compose up -d
+
+# set env variables
+export audius_db_url='postgres://postgres:testing@localhost:35764'
+export audius_elasticsearch_url='http://localhost:35765'
+export DB_URL="$audius_db_url"
+
+# run pg_migrate
+# cd ../discovery-provider/ddl && ./pg_migrate.sh && cd - || exit
+docker exec -w '/ddl' trpc-server-db-1 './pg_migrate.sh'
+
+# populate db fixtures
+npx vite-node test/_fixtures.ts
+
+# run es-indexer
+cd ../es-indexer && npm run catchup:ci && cd - || exit
+
+# run tests
+npx vitest "$1"