Refactor Elasticsearch client calls into client wrapper and break out types

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
2022-06-15 11:07:07 +02:00
parent ca9ff1eaac
commit 68b7739aab
10 changed files with 659 additions and 160 deletions
@@ -5,60 +5,27 @@
 ```ts
 /// <reference types="node" />

+import { ApiResponse } from '@opensearch-project/opensearch';
 import { BatchSearchEngineIndexer } from '@backstage/plugin-search-backend-node';
-import { Client } from '@elastic/elasticsearch';
+import { BulkHelper } from '@opensearch-project/opensearch/lib/Helpers';
+import { BulkStats } from '@opensearch-project/opensearch/lib/Helpers';
 import { Config } from '@backstage/config';
 import type { ConnectionOptions } from 'tls';
 import { IndexableDocument } from '@backstage/plugin-search-common';
 import { IndexableResultSet } from '@backstage/plugin-search-common';
 import { Logger } from 'winston';
+import { Readable } from 'stream';
 import { SearchEngine } from '@backstage/plugin-search-common';
 import { SearchQuery } from '@backstage/plugin-search-common';
-
-// @public (undocumented)
-export interface ElasticSearchAgentOptions {
-  // (undocumented)
-  keepAlive?: boolean;
-  // (undocumented)
-  keepAliveMsecs?: number;
-  // (undocumented)
-  maxFreeSockets?: number;
-  // (undocumented)
-  maxSockets?: number;
-}
-
-// @public (undocumented)
-export type ElasticSearchAuth =
-  | {
-      username: string;
-      password: string;
-    }
-  | {
-      apiKey:
-        | string
-        | {
-            id: string;
-            api_key: string;
-          };
-    };
+import { TransportRequestPromise } from '@opensearch-project/opensearch/lib/Transport';

 // @public
-export interface ElasticSearchClientOptions {
+export interface BaseElasticSearchClientOptions {
  // (undocumented)
  agent?: ElasticSearchAgentOptions | ((opts?: any) => unknown) | false;
  // (undocumented)
-  auth?: ElasticSearchAuth;
-  // (undocumented)
-  cloud?: {
-    id: string;
-    username?: string;
-    password?: string;
-  };
-  // (undocumented)
  compression?: 'gzip';
  // (undocumented)
-  Connection?: ElasticSearchConnectionConstructor;
-  // (undocumented)
  disablePrototypePoisoningProtection?: boolean | 'proto' | 'constructor';
  // (undocumented)
  enableMetaHeader?: boolean;
@@ -69,28 +36,14 @@ export interface ElasticSearchClientOptions {
  // (undocumented)
  name?: string | symbol;
  // (undocumented)
-  node?:
-    | string
-    | string[]
-    | ElasticSearchNodeOptions
-    | ElasticSearchNodeOptions[];
-  // (undocumented)
  nodeFilter?: (connection: any) => boolean;
  // (undocumented)
-  nodes?:
-    | string
-    | string[]
-    | ElasticSearchNodeOptions
-    | ElasticSearchNodeOptions[];
-  // (undocumented)
  nodeSelector?: ((connections: any[]) => any) | string;
  // (undocumented)
  opaqueIdPrefix?: string;
  // (undocumented)
  pingTimeout?: number;
  // (undocumented)
-  provider?: 'aws' | 'elastic';
-  // (undocumented)
  proxy?: string | URL;
  // (undocumented)
  requestTimeout?: number;
@@ -112,6 +65,118 @@ export interface ElasticSearchClientOptions {
  Transport?: ElasticSearchTransportConstructor;
 }

+// @public (undocumented)
+export interface ElasticSearchAgentOptions {
+  // (undocumented)
+  keepAlive?: boolean;
+  // (undocumented)
+  keepAliveMsecs?: number;
+  // (undocumented)
+  maxFreeSockets?: number;
+  // (undocumented)
+  maxSockets?: number;
+}
+
+// @public (undocumented)
+export type ElasticSearchAliasAction =
+  | {
+      remove: {
+        index: any;
+        alias: any;
+      };
+      add?: undefined;
+    }
+  | {
+      add: {
+        indices: any;
+        alias: any;
+        index?: undefined;
+      };
+      remove?: undefined;
+    }
+  | {
+      add: {
+        index: any;
+        alias: any;
+        indices?: undefined;
+      };
+      remove?: undefined;
+    }
+  | undefined;
+
+// @public (undocumented)
+export type ElasticSearchAuth =
+  | OpenSearchAuth
+  | {
+      apiKey:
+        | string
+        | {
+            id: string;
+            api_key: string;
+          };
+    };
+
+// @public
+export type ElasticSearchClientOptions =
+  | ElasticSearchElasticSearchClientOptions
+  | OpenSearchElasticSearchClientOptions;
+
+// @public
+export class ElasticSearchClientWrapper {
+  // (undocumented)
+  bulk(bulkOptions: {
+    datasource: Readable;
+    onDocument: () => ElasticSearchIndexAction;
+    refreshOnCompletion?: string | boolean;
+  }): BulkHelper<BulkStats>;
+  // (undocumented)
+  createIndex({
+    index,
+  }: {
+    index: string;
+  }): TransportRequestPromise<ApiResponse<Record<string, any>, unknown>>;
+  // (undocumented)
+  deleteIndex({
+    index,
+  }: {
+    index: string | string[];
+  }): TransportRequestPromise<ApiResponse<Record<string, any>, unknown>>;
+  // (undocumented)
+  static fromClientOptions(
+    options: ElasticSearchClientOptions,
+  ): ElasticSearchClientWrapper;
+  // (undocumented)
+  getAliases({
+    aliases,
+  }: {
+    aliases: string[];
+  }): TransportRequestPromise<ApiResponse<Record<string, any>, unknown>>;
+  // (undocumented)
+  indexExists({
+    index,
+  }: {
+    index: string | string[];
+  }): TransportRequestPromise<ApiResponse<boolean, unknown>>;
+  // (undocumented)
+  putIndexTemplate(
+    template: ElasticSearchCustomIndexTemplate,
+  ): TransportRequestPromise<ApiResponse<Record<string, any>, unknown>>;
+  // (undocumented)
+  search({
+    index,
+    body,
+  }: {
+    index: string | string[];
+    body: Object;
+  }): TransportRequestPromise<ApiResponse<Record<string, any>, unknown>>;
+  // (undocumented)
+  updateAliases({
+    actions,
+  }: {
+    actions: ElasticSearchAliasAction[];
+  }): TransportRequestPromise<ApiResponse<Record<string, any>, unknown>>;
+}
+
 // @public
 export type ElasticSearchConcreteQuery = {
  documentTypes?: string[];
@@ -150,6 +215,35 @@ export type ElasticSearchCustomIndexTemplateBody = {
  template?: Record<string, any>;
 };

+// @public
+export interface ElasticSearchElasticSearchClientOptions
+  extends BaseElasticSearchClientOptions {
+  // (undocumented)
+  auth?: ElasticSearchAuth;
+  // (undocumented)
+  cloud?: {
+    id: string;
+    username?: string;
+    password?: string;
+  };
+  // (undocumented)
+  Connection?: ElasticSearchConnectionConstructor;
+  // (undocumented)
+  node?:
+    | string
+    | string[]
+    | ElasticSearchNodeOptions
+    | ElasticSearchNodeOptions[];
+  // (undocumented)
+  nodes?:
+    | string
+    | string[]
+    | ElasticSearchNodeOptions
+    | ElasticSearchNodeOptions[];
+  // (undocumented)
+  provider?: 'elastic';
+}
+
 // @public (undocumented)
 export type ElasticSearchHighlightConfig = {
  fragmentDelimiter: string;
@@ -166,6 +260,14 @@ export type ElasticSearchHighlightOptions = {
  numFragments?: number;
 };

+// @public (undocumented)
+export type ElasticSearchIndexAction = {
+  index: {
+    _index: string;
+    [key: string]: any;
+  };
+};
+
 // @public (undocumented)
 export interface ElasticSearchNodeOptions {
  // (undocumented)
@@ -258,7 +360,7 @@ export type ElasticSearchSearchEngineIndexerOptions = {
  indexSeparator: string;
  alias: string;
  logger: Logger;
-  elasticSearchClient: Client;
+  elasticSearchClientWrapper: ElasticSearchClientWrapper;
 };

 // @public (undocumented)
@@ -273,4 +375,67 @@ export interface ElasticSearchTransportConstructor {
    DEFAULT: string;
  };
 }
+
+// @public
+export const isOpenSearchCompatible: (
+  opts: ElasticSearchClientOptions,
+) => opts is OpenSearchElasticSearchClientOptions;
+
+// @public (undocumented)
+export type OpenSearchAuth = {
+  username: string;
+  password: string;
+};
+
+// @public (undocumented)
+export interface OpenSearchConnectionConstructor {
+  // (undocumented)
+  new (opts?: any): any;
+  // (undocumented)
+  roles: {
+    MASTER: string;
+    DATA: string;
+    INGEST: string;
+  };
+  // (undocumented)
+  statuses: {
+    ALIVE: string;
+    DEAD: string;
+  };
+}
+
+// @public
+export interface OpenSearchElasticSearchClientOptions
+  extends BaseElasticSearchClientOptions {
+  // (undocumented)
+  auth?: OpenSearchAuth;
+  // (undocumented)
+  connection?: OpenSearchConnectionConstructor;
+  // (undocumented)
+  node?: string | string[] | OpenSearchNodeOptions | OpenSearchNodeOptions[];
+  // (undocumented)
+  nodes?: string | string[] | OpenSearchNodeOptions | OpenSearchNodeOptions[];
+  // (undocumented)
+  provider?: 'aws';
+}
+
+// @public (undocumented)
+export interface OpenSearchNodeOptions {
+  // (undocumented)
+  agent?: ElasticSearchAgentOptions;
+  // (undocumented)
+  headers?: Record<string, any>;
+  // (undocumented)
+  id?: string;
+  // (undocumented)
+  roles?: {
+    master: boolean;
+    data: boolean;
+    ingest: boolean;
+  };
+  // (undocumented)
+  ssl?: ConnectionOptions;
+  // (undocumented)
+  url: URL;
+}
 ```
@@ -16,17 +16,60 @@
 import type { ConnectionOptions as TLSConnectionOptions } from 'tls';

 /**
- * Options used to configure the `@elastic/elasticsearch` client and
- * are what will be passed as an argument to the
- * {@link ElasticSearchSearchEngine.newClient} method
- *
- * They are drawn from the `ClientOptions` class of `@elastic/elasticsearch`,
- * but are maintained separately so that this interface is not coupled to
+ * Typeguard to differentiate ElasticSearch client options which are compatible
+ * with OpenSearch vs. ElasticSearch clients. Useful when calling the
+ * {@link ElasticSearchSearchEngine.newClient} method.
 *
 * @public
 */
-export interface ElasticSearchClientOptions {
-  provider?: 'aws' | 'elastic';
+export const isOpenSearchCompatible = (
+  opts: ElasticSearchClientOptions,
+): opts is OpenSearchElasticSearchClientOptions => {
+  return opts?.provider === 'aws';
+};
+
+/**
+ * Options used to configure the `@elastic/elasticsearch` client or the
+ * `@opensearch-project/opensearch` client, depending on the given config. It
+ * will be passed as an argument to the
+ * {@link ElasticSearchSearchEngine.newClient} method.
+ *
+ * @public
+ */
+export type ElasticSearchClientOptions =
+  | ElasticSearchElasticSearchClientOptions
+  | OpenSearchElasticSearchClientOptions;
+
+/**
+ * Options used to configure the `@opensearch-project/opensearch` client.
+ *
+ * They are drawn from the `ClientOptions` class of `@opensearch-project/opensearch`,
+ * but are maintained separately so that this interface is not coupled to it.
+ *
+ * @public
+ */
+export interface OpenSearchElasticSearchClientOptions
+  extends BaseElasticSearchClientOptions {
+  provider?: 'aws';
+  auth?: OpenSearchAuth;
+  connection?: OpenSearchConnectionConstructor;
+  node?: string | string[] | OpenSearchNodeOptions | OpenSearchNodeOptions[];
+  nodes?: string | string[] | OpenSearchNodeOptions | OpenSearchNodeOptions[];
+}
+
+/**
+ * Options used to configure the `@elastic/elasticsearch` client.
+ *
+ * They are drawn from the `ClientOptions` class of `@elastic/elasticsearch`,
+ * but are maintained separately so that this interface is not coupled to it.
+ *
+ * @public
+ */
+export interface ElasticSearchElasticSearchClientOptions
+  extends BaseElasticSearchClientOptions {
+  provider?: 'elastic';
+  auth?: ElasticSearchAuth;
+  Connection?: ElasticSearchConnectionConstructor;
  node?:
    | string
    | string[]
@@ -37,8 +80,21 @@ export interface ElasticSearchClientOptions {
    | string[]
    | ElasticSearchNodeOptions
    | ElasticSearchNodeOptions[];
+  cloud?: {
+    id: string;
+    username?: string;
+    password?: string;
+  };
+}
+
+/**
+ * Base client options that are shared across `@opensearch-project/opensearch`
+ * and `@elastic/elasticsearch` clients.
+ *
+ * @public
+ */
+export interface BaseElasticSearchClientOptions {
  Transport?: ElasticSearchTransportConstructor;
-  Connection?: ElasticSearchConnectionConstructor;
  maxRetries?: number;
  requestTimeout?: number;
  pingTimeout?: number;
@@ -56,25 +112,24 @@ export interface ElasticSearchClientOptions {
  headers?: Record<string, any>;
  opaqueIdPrefix?: string;
  name?: string | symbol;
-  auth?: ElasticSearchAuth;
  proxy?: string | URL;
  enableMetaHeader?: boolean;
-  cloud?: {
-    id: string;
-    username?: string;
-    password?: string;
-  };
  disablePrototypePoisoningProtection?: boolean | 'proto' | 'constructor';
 }

+/**
+ * @public
+ */
+export type OpenSearchAuth = {
+  username: string;
+  password: string;
+};
+
 /**
 * @public
 */
 export type ElasticSearchAuth =
-  | {
-      username: string;
-      password: string;
-    }
+  | OpenSearchAuth
  | {
      apiKey:
        | string
@@ -101,6 +156,22 @@ export interface ElasticSearchNodeOptions {
  };
 }

+/**
+ * @public
+ */
+export interface OpenSearchNodeOptions {
+  url: URL;
+  id?: string;
+  agent?: ElasticSearchAgentOptions;
+  ssl?: TLSConnectionOptions;
+  headers?: Record<string, any>;
+  roles?: {
+    master: boolean;
+    data: boolean;
+    ingest: boolean;
+  };
+}
+
 /**
 * @public
 */
@@ -127,6 +198,23 @@ export interface ElasticSearchConnectionConstructor {
    ML: string;
  };
 }
+
+/**
+ * @public
+ */
+export interface OpenSearchConnectionConstructor {
+  new (opts?: any): any;
+  statuses: {
+    ALIVE: string;
+    DEAD: string;
+  };
+  roles: {
+    MASTER: string;
+    DATA: string;
+    INGEST: string;
+  };
+}
+
 /**
 * @public
 */
@@ -139,3 +227,5 @@ export interface ElasticSearchTransportConstructor {
    DEFAULT: string;
  };
 }
+
+// todo(iamEAP) implement canary types to ensure we remain compatible through upgrades.
@@ -0,0 +1,209 @@
+/*
+ * Copyright 2021 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import { Client as ElasticSearchClient } from '@elastic/elasticsearch';
+import { Client as OpenSearchClient } from '@opensearch-project/opensearch';
+import { Readable } from 'stream';
+import {
+  ElasticSearchClientOptions,
+  isOpenSearchCompatible,
+} from './ElasticSearchClientOptions';
+import { ElasticSearchCustomIndexTemplate } from './types';
+
+/**
+ * @public
+ */
+export type ElasticSearchAliasAction =
+  | {
+      remove: { index: any; alias: any };
+      add?: undefined;
+    }
+  | {
+      add: { indices: any; alias: any; index?: undefined };
+      remove?: undefined;
+    }
+  | {
+      add: { index: any; alias: any; indices?: undefined };
+      remove?: undefined;
+    }
+  | undefined;
+
+/**
+ * @public
+ */
+export type ElasticSearchIndexAction = {
+  index: {
+    _index: string;
+    [key: string]: any;
+  };
+};
+
+/**
+ * A wrapper class that exposes logical methods that are conditionally fired
+ * against either a configured Elasticsearch client or a configured Opensearch
+ * client.
+ *
+ * This is necessary because, despite its intention to be API-compatible, the
+ * opensearch client does not support API key-based authentication. This is
+ * also the sanest way to accomplish this while making typescript happy.
+ *
+ * In the future, if the differences between implementations become
+ * unmaintainably divergent, we should split out the Opensearch and
+ * Elasticsearch search engine implementations.
+ *
+ * @public
+ */
+export class ElasticSearchClientWrapper {
+  private readonly elasticSearchClient: ElasticSearchClient | undefined;
+  private readonly openSearchClient: OpenSearchClient | undefined;
+
+  private constructor({
+    openSearchClient,
+    elasticSearchClient,
+  }: {
+    openSearchClient?: OpenSearchClient;
+    elasticSearchClient?: ElasticSearchClient;
+  }) {
+    this.openSearchClient = openSearchClient;
+    this.elasticSearchClient = elasticSearchClient;
+  }
+
+  static fromClientOptions(options: ElasticSearchClientOptions) {
+    if (isOpenSearchCompatible(options)) {
+      return new ElasticSearchClientWrapper({
+        openSearchClient: new OpenSearchClient(options),
+      });
+    }
+
+    return new ElasticSearchClientWrapper({
+      elasticSearchClient: new ElasticSearchClient(options),
+    });
+  }
+
+  search({ index, body }: { index: string | string[]; body: Object }) {
+    if (this.openSearchClient) {
+      return this.openSearchClient.search({ index, body });
+    }
+
+    if (this.elasticSearchClient) {
+      return this.elasticSearchClient.search({ index, body });
+    }
+
+    throw new Error('No client defined');
+  }
+
+  bulk(bulkOptions: {
+    datasource: Readable;
+    onDocument: () => ElasticSearchIndexAction;
+    refreshOnCompletion?: string | boolean;
+  }) {
+    if (this.openSearchClient) {
+      return this.openSearchClient.helpers.bulk(bulkOptions);
+    }
+
+    if (this.elasticSearchClient) {
+      return this.elasticSearchClient.helpers.bulk(bulkOptions);
+    }
+
+    throw new Error('No client defined');
+  }
+
+  putIndexTemplate(template: ElasticSearchCustomIndexTemplate) {
+    if (this.openSearchClient) {
+      return this.openSearchClient.indices.putIndexTemplate(template);
+    }
+
+    if (this.elasticSearchClient) {
+      return this.elasticSearchClient.indices.putIndexTemplate(template);
+    }
+
+    throw new Error('No client defined');
+  }
+
+  indexExists({ index }: { index: string | string[] }) {
+    if (this.openSearchClient) {
+      return this.openSearchClient.indices.exists({ index });
+    }
+
+    if (this.elasticSearchClient) {
+      return this.elasticSearchClient.indices.exists({ index });
+    }
+
+    throw new Error('No client defined');
+  }
+
+  deleteIndex({ index }: { index: string | string[] }) {
+    if (this.openSearchClient) {
+      return this.openSearchClient.indices.delete({ index });
+    }
+
+    if (this.elasticSearchClient) {
+      return this.elasticSearchClient.indices.delete({ index });
+    }
+
+    throw new Error('No client defined');
+  }
+
+  createIndex({ index }: { index: string }) {
+    if (this.openSearchClient) {
+      return this.openSearchClient.indices.create({ index });
+    }
+
+    if (this.elasticSearchClient) {
+      return this.elasticSearchClient.indices.create({ index });
+    }
+
+    throw new Error('No client defined');
+  }
+
+  getAliases({ aliases }: { aliases: string[] }) {
+    if (this.openSearchClient) {
+      return this.openSearchClient.cat.aliases({
+        format: 'json',
+        name: aliases,
+      });
+    }
+
+    if (this.elasticSearchClient) {
+      return this.elasticSearchClient.cat.aliases({
+        format: 'json',
+        name: aliases,
+      });
+    }
+
+    throw new Error('No client defined');
+  }
+
+  updateAliases({ actions }: { actions: ElasticSearchAliasAction[] }) {
+    const filteredActions = actions.filter(Boolean);
+    if (this.openSearchClient) {
+      return this.openSearchClient.indices.updateAliases({
+        body: {
+          actions: filteredActions,
+        },
+      });
+    }
+
+    if (this.elasticSearchClient) {
+      return this.elasticSearchClient.indices.updateAliases({
+        body: {
+          actions: filteredActions,
+        },
+      });
+    }
+
+    throw new Error('No client defined');
+  }
+}
@@ -16,8 +16,9 @@

 import { getVoidLogger } from '@backstage/backend-common';
 import { ConfigReader } from '@backstage/config';
-import { Client, errors } from '@elastic/elasticsearch';
+import { errors } from '@elastic/elasticsearch';
 import Mock from '@elastic/elasticsearch-mock';
+import { ElasticSearchClientWrapper } from './ElasticSearchClientWrapper';
 import {
  ElasticSearchConcreteQuery,
  decodePageCursor,
@@ -70,7 +71,7 @@ const customIndexTemplate = {
 describe('ElasticSearchSearchEngine', () => {
  let testSearchEngine: ElasticSearchSearchEngine;
  let inspectableSearchEngine: ElasticSearchSearchEngineForTranslatorTests;
-  let client: Client;
+  let clientWrapper: ElasticSearchClientWrapper;

  beforeEach(() => {
    testSearchEngine = new ElasticSearchSearchEngine(
@@ -86,7 +87,7 @@ describe('ElasticSearchSearchEngine', () => {
      getVoidLogger(),
    );
    // eslint-disable-next-line dot-notation
-    client = testSearchEngine['elasticSearchClient'];
+    clientWrapper = testSearchEngine['elasticSearchClientWrapper'];
  });

  describe('custom index template', () => {
@@ -686,7 +687,7 @@ describe('ElasticSearchSearchEngine', () => {
    });

    it('should handle index/search type filtering correctly', async () => {
-      const elasticSearchQuerySpy = jest.spyOn(client, 'search');
+      const elasticSearchQuerySpy = jest.spyOn(clientWrapper, 'search');
      await testSearchEngine.query({
        term: 'testTerm',
        filters: {},
@@ -725,7 +726,7 @@ describe('ElasticSearchSearchEngine', () => {
    });

    it('should create matchAll query if no term defined', async () => {
-      const elasticSearchQuerySpy = jest.spyOn(client, 'search');
+      const elasticSearchQuerySpy = jest.spyOn(clientWrapper, 'search');
      await testSearchEngine.query({
        term: '',
        filters: {},
@@ -759,7 +760,7 @@ describe('ElasticSearchSearchEngine', () => {
    });

    it('should query only specified indices if defined', async () => {
-      const elasticSearchQuerySpy = jest.spyOn(client, 'search');
+      const elasticSearchQuerySpy = jest.spyOn(clientWrapper, 'search');
      await testSearchEngine.query({
        term: '',
        filters: {},
@@ -809,7 +810,7 @@ describe('ElasticSearchSearchEngine', () => {
          type: 'test-index',
          indexPrefix: '',
          indexSeparator: '-index__',
-          elasticSearchClient: client,
+          elasticSearchClientWrapper: clientWrapper,
        }),
      );
      expect(indexerMock.on).toHaveBeenCalledWith(
@@ -14,10 +14,7 @@
 * limitations under the License.
 */

-import {
-  awsGetCredentials,
-  createAWSConnection,
-} from '@acuris/aws-es-connection';
+import { awsGetCredentials, createAWSConnection } from 'aws-os-connection';
 import { Config } from '@backstage/config';
 import {
  IndexableDocument,
@@ -26,47 +23,17 @@ import {
  SearchEngine,
  SearchQuery,
 } from '@backstage/plugin-search-common';
-import { Client } from '@elastic/elasticsearch';
 import esb from 'elastic-builder';
 import { isEmpty, isNaN as nan, isNumber } from 'lodash';
 import { v4 as uuid } from 'uuid';
 import { Logger } from 'winston';
-import type { ElasticSearchClientOptions } from './ElasticSearchClientOptions';
+import { ElasticSearchClientOptions } from './ElasticSearchClientOptions';
 import { ElasticSearchSearchEngineIndexer } from './ElasticSearchSearchEngineIndexer';
+import { ElasticSearchCustomIndexTemplate } from './types';
+import { ElasticSearchClientWrapper } from './ElasticSearchClientWrapper';

 export type { ElasticSearchClientOptions };

-/**
- * Elasticsearch specific index template
- * @public
- */
-export type ElasticSearchCustomIndexTemplate = {
-  name: string;
-  body: ElasticSearchCustomIndexTemplateBody;
-};
-
-/**
- * Elasticsearch specific index template body
- * @public
- */
-export type ElasticSearchCustomIndexTemplateBody = {
-  /**
-   * Array of wildcard (*) expressions used to match the names of data streams and indices during creation.
-   */
-  index_patterns: string[];
-  /**
-   * An ordered list of component template names.
-   * Component templates are merged in the order specified,
-   * meaning that the last component template specified has the highest precedence.
-   */
-  composed_of?: string[];
-  /**
-   * See available properties of template
-   * https://www.elastic.co/guide/en/elasticsearch/reference/7.15/indices-put-template.html#put-index-template-api-request-body
-   */
-  template?: Record<string, any>;
-};
-
 /**
 * Search query that the elasticsearch engine understands.
 * @public
@@ -143,7 +110,7 @@ function isBlank(str: string) {
 * @public
 */
 export class ElasticSearchSearchEngine implements SearchEngine {
-  private readonly elasticSearchClient: Client;
+  private readonly elasticSearchClientWrapper: ElasticSearchClientWrapper;
  private readonly highlightOptions: ElasticSearchHighlightConfig;

  constructor(
@@ -153,7 +120,8 @@ export class ElasticSearchSearchEngine implements SearchEngine {
    private readonly logger: Logger,
    highlightOptions?: ElasticSearchHighlightOptions,
  ) {
-    this.elasticSearchClient = this.newClient(options => new Client(options));
+    this.elasticSearchClientWrapper =
+      ElasticSearchClientWrapper.fromClientOptions(elasticSearchClientOptions);
    const uuidTag = uuid();
    this.highlightOptions = {
      preTag: `<${uuidTag}>`,
@@ -177,7 +145,7 @@ export class ElasticSearchSearchEngine implements SearchEngine {
    if (options.provider === 'elastic') {
      logger.info('Initializing Elastic.co ElasticSearch search engine.');
    } else if (options.provider === 'aws') {
-      logger.info('Initializing AWS ElasticSearch search engine.');
+      logger.info('Initializing AWS OpenSearch search engine.');
    } else {
      logger.info('Initializing ElasticSearch search engine.');
    }
@@ -269,7 +237,7 @@ export class ElasticSearchSearchEngine implements SearchEngine {

  async setIndexTemplate(template: ElasticSearchCustomIndexTemplate) {
    try {
-      await this.elasticSearchClient.indices.putIndexTemplate(template);
+      await this.elasticSearchClientWrapper.putIndexTemplate(template);
      this.logger.info('Custom index template set');
    } catch (error) {
      this.logger.error(`Unable to set custom index template: ${error}`);
@@ -284,7 +252,7 @@ export class ElasticSearchSearchEngine implements SearchEngine {
      indexPrefix: this.indexPrefix,
      indexSeparator: this.indexSeparator,
      alias,
-      elasticSearchClient: this.elasticSearchClient,
+      elasticSearchClientWrapper: this.elasticSearchClientWrapper,
      logger: this.logger,
    });

@@ -292,13 +260,13 @@ export class ElasticSearchSearchEngine implements SearchEngine {
    indexer.on('error', async e => {
      this.logger.error(`Failed to index documents for type ${type}`, e);
      try {
-        const response = await this.elasticSearchClient.indices.exists({
+        const response = await this.elasticSearchClientWrapper.indexExists({
          index: indexer.indexName,
        });
        const indexCreated = response.body;
        if (indexCreated) {
          this.logger.info(`Removing created index ${indexer.indexName}`);
-          await this.elasticSearchClient.indices.delete({
+          await this.elasticSearchClientWrapper.deleteIndex({
            index: indexer.indexName,
          });
        }
@@ -319,7 +287,7 @@ export class ElasticSearchSearchEngine implements SearchEngine {
      ? documentTypes.map(it => this.constructSearchAlias(it))
      : this.constructSearchAlias('*');
    try {
-      const result = await this.elasticSearchClient.search({
+      const result = await this.elasticSearchClientWrapper.search({
        index: queryIndices,
        body: elasticSearchQuery,
      });
@@ -16,13 +16,13 @@

 import { getVoidLogger } from '@backstage/backend-common';
 import { TestPipeline } from '@backstage/plugin-search-backend-node';
-import { Client } from '@elastic/elasticsearch';
 import Mock from '@elastic/elasticsearch-mock';
 import { range } from 'lodash';
+import { ElasticSearchClientWrapper } from './ElasticSearchClientWrapper';
 import { ElasticSearchSearchEngineIndexer } from './ElasticSearchSearchEngineIndexer';

 const mock = new Mock();
-const client = new Client({
+const clientWrapper = ElasticSearchClientWrapper.fromClientOptions({
  node: 'http://localhost:9200',
  Connection: mock.getConnection(),
 });
@@ -43,7 +43,7 @@ describe('ElasticSearchSearchEngineIndexer', () => {
      indexSeparator: '-index__',
      alias: 'some-type-index__search',
      logger: getVoidLogger(),
-      elasticSearchClient: client,
+      elasticSearchClientWrapper: clientWrapper,
    });

    // Set up all requisite Elastic mocks.
@@ -16,9 +16,9 @@

 import { BatchSearchEngineIndexer } from '@backstage/plugin-search-backend-node';
 import { IndexableDocument } from '@backstage/plugin-search-common';
-import { Client } from '@elastic/elasticsearch';
 import { Readable } from 'stream';
 import { Logger } from 'winston';
+import { ElasticSearchClientWrapper } from './ElasticSearchClientWrapper';

 /**
 * Options for instansiate ElasticSearchSearchEngineIndexer
@@ -30,7 +30,7 @@ export type ElasticSearchSearchEngineIndexerOptions = {
  indexSeparator: string;
  alias: string;
  logger: Logger;
-  elasticSearchClient: Client;
+  elasticSearchClientWrapper: ElasticSearchClientWrapper;
 };

 function duration(startTimestamp: [number, number]): string {
@@ -57,7 +57,7 @@ export class ElasticSearchSearchEngineIndexer extends BatchSearchEngineIndexer {
  private readonly removableAlias: string;
  private readonly logger: Logger;
  private readonly sourceStream: Readable;
-  private readonly elasticSearchClient: Client;
+  private readonly elasticSearchClientWrapper: ElasticSearchClientWrapper;
  private bulkResult: Promise<any>;

  constructor(options: ElasticSearchSearchEngineIndexerOptions) {
@@ -70,7 +70,7 @@ export class ElasticSearchSearchEngineIndexer extends BatchSearchEngineIndexer {
    this.indexName = this.constructIndexName(`${Date.now()}`);
    this.alias = options.alias;
    this.removableAlias = `${this.alias}_removable`;
-    this.elasticSearchClient = options.elasticSearchClient;
+    this.elasticSearchClientWrapper = options.elasticSearchClientWrapper;

    // The ES client bulk helper supports stream-based indexing, but we have to
    // supply the stream directly to it at instantiation-time. We can't supply
@@ -83,7 +83,7 @@ export class ElasticSearchSearchEngineIndexer extends BatchSearchEngineIndexer {

    // Keep a reference to the ES Bulk helper so that we can know when all
    // documents have been successfully written to ES.
-    this.bulkResult = this.elasticSearchClient.helpers.bulk({
+    this.bulkResult = this.elasticSearchClientWrapper.bulk({
      datasource: this.sourceStream,
      onDocument() {
        that.processed++;
@@ -98,16 +98,15 @@ export class ElasticSearchSearchEngineIndexer extends BatchSearchEngineIndexer {
  async initialize(): Promise<void> {
    this.logger.info(`Started indexing documents for index ${this.type}`);

-    const aliases = await this.elasticSearchClient.cat.aliases({
-      format: 'json',
-      name: [this.alias, this.removableAlias],
+    const aliases = await this.elasticSearchClientWrapper.getAliases({
+      aliases: [this.alias, this.removableAlias],
    });

    this.removableIndices = [
      ...new Set(aliases.body.map((r: Record<string, any>) => r.index)),
    ] as string[];

-    await this.elasticSearchClient.indices.create({
+    await this.elasticSearchClientWrapper.createIndex({
      index: this.indexName,
    });
  }
@@ -140,25 +139,23 @@ export class ElasticSearchSearchEngineIndexer extends BatchSearchEngineIndexer {
      )}`,
      result,
    );
-    await this.elasticSearchClient.indices.updateAliases({
-      body: {
-        actions: [
-          {
-            remove: { index: this.constructIndexName('*'), alias: this.alias },
-          },
-          this.removableIndices.length
-            ? {
-                add: {
-                  indices: this.removableIndices,
-                  alias: this.removableAlias,
-                },
-              }
-            : undefined,
-          {
-            add: { index: this.indexName, alias: this.alias },
-          },
-        ].filter(Boolean),
-      },
+    await this.elasticSearchClientWrapper.updateAliases({
+      actions: [
+        {
+          remove: { index: this.constructIndexName('*'), alias: this.alias },
+        },
+        this.removableIndices.length
+          ? {
+              add: {
+                indices: this.removableIndices,
+                alias: this.removableAlias,
+              },
+            }
+          : undefined,
+        {
+          add: { index: this.indexName, alias: this.alias },
+        },
+      ].filter(Boolean),
    });

    // If any indices are removable, remove them. Do not bubble up this error,
@@ -166,7 +163,7 @@ export class ElasticSearchSearchEngineIndexer extends BatchSearchEngineIndexer {
    if (this.removableIndices.length) {
      this.logger.info('Removing stale search indices', this.removableIndices);
      try {
-        await this.elasticSearchClient.indices.delete({
+        await this.elasticSearchClientWrapper.deleteIndex({
          index: this.removableIndices,
        });
      } catch (e) {
@@ -15,13 +15,25 @@
 */

 export { ElasticSearchSearchEngine } from './ElasticSearchSearchEngine';
+export { isOpenSearchCompatible } from './ElasticSearchClientOptions';
 export type {
+  BaseElasticSearchClientOptions,
  ElasticSearchAgentOptions,
  ElasticSearchConnectionConstructor,
+  ElasticSearchElasticSearchClientOptions,
  ElasticSearchTransportConstructor,
  ElasticSearchNodeOptions,
  ElasticSearchAuth,
+  OpenSearchAuth,
+  OpenSearchConnectionConstructor,
+  OpenSearchElasticSearchClientOptions,
+  OpenSearchNodeOptions,
 } from './ElasticSearchClientOptions';
+export type {
+  ElasticSearchAliasAction,
+  ElasticSearchClientWrapper,
+  ElasticSearchIndexAction,
+} from './ElasticSearchClientWrapper';
 export type {
  ElasticSearchConcreteQuery,
  ElasticSearchClientOptions,
@@ -30,10 +42,12 @@ export type {
  ElasticSearchQueryTranslator,
  ElasticSearchQueryTranslatorOptions,
  ElasticSearchOptions,
-  ElasticSearchCustomIndexTemplate,
-  ElasticSearchCustomIndexTemplateBody,
 } from './ElasticSearchSearchEngine';
 export type {
  ElasticSearchSearchEngineIndexer,
  ElasticSearchSearchEngineIndexerOptions,
 } from './ElasticSearchSearchEngineIndexer';
+export type {
+  ElasticSearchCustomIndexTemplate,
+  ElasticSearchCustomIndexTemplateBody,
+} from './types';
@@ -0,0 +1,46 @@
+/*
+ * Copyright 2022 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * Elasticsearch specific index template
+ * @public
+ */
+export type ElasticSearchCustomIndexTemplate = {
+  name: string;
+  body: ElasticSearchCustomIndexTemplateBody;
+};
+
+/**
+ * Elasticsearch specific index template body
+ * @public
+ */
+export type ElasticSearchCustomIndexTemplateBody = {
+  /**
+   * Array of wildcard (*) expressions used to match the names of data streams and indices during creation.
+   */
+  index_patterns: string[];
+  /**
+   * An ordered list of component template names.
+   * Component templates are merged in the order specified,
+   * meaning that the last component template specified has the highest precedence.
+   */
+  composed_of?: string[];
+  /**
+   * See available properties of template
+   * https://www.elastic.co/guide/en/elasticsearch/reference/7.15/indices-put-template.html#put-index-template-api-request-body
+   */
+  template?: Record<string, any>;
+};
@@ -20,13 +20,18 @@
 * @packageDocumentation
 */

-export { ElasticSearchSearchEngine } from './engines';
+export { ElasticSearchSearchEngine, isOpenSearchCompatible } from './engines';
 export type {
+  BaseElasticSearchClientOptions,
  ElasticSearchAgentOptions,
+  ElasticSearchAliasAction,
+  ElasticSearchClientWrapper,
  ElasticSearchConcreteQuery,
  ElasticSearchClientOptions,
+  ElasticSearchElasticSearchClientOptions,
  ElasticSearchHighlightConfig,
  ElasticSearchHighlightOptions,
+  ElasticSearchIndexAction,
  ElasticSearchQueryTranslator,
  ElasticSearchQueryTranslatorOptions,
  ElasticSearchConnectionConstructor,
@@ -38,4 +43,8 @@ export type {
  ElasticSearchSearchEngineIndexerOptions,
  ElasticSearchCustomIndexTemplate,
  ElasticSearchCustomIndexTemplateBody,
+  OpenSearchAuth,
+  OpenSearchConnectionConstructor,
+  OpenSearchElasticSearchClientOptions,
+  OpenSearchNodeOptions,
 } from './engines';