[Maps] ESQL geo_shape support (elastic#175156)

elastic/elasticsearch#104269 adds geo_shape
support to ESQL

This PR updates maps ESQL source to support geo_shape column type
<img width="800" alt="Screenshot 2024-01-18 at 1 15 31 PM"
src="https://github.com/elastic/kibana/assets/373691/844251a6-4409-45f4-89f5-513f4537d833">

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>

x-pack/plugins/maps/public/classes/sources/esql_source/create_source_editor.test.tsx

@@ -0,0 +1,93 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import React from 'react';
+import { render, waitFor } from '@testing-library/react';
+import { CreateSourceEditor } from './create_source_editor';
+
+jest.mock('../../../kibana_services', () => {
+  const mockDefaultDataView = {
+    fields: [
+      {
+        name: 'location',
+        type: 'geo_point',
+      },
+      {
+        name: '@timestamp',
+        type: 'date',
+      },
+    ],
+    timeFieldName: '@timestamp',
+    getIndexPattern: () => {
+      return 'logs';
+    },
+  };
+  const mockDataView = {
+    fields: [
+      {
+        name: 'geometry',
+        type: 'geo_shape',
+      },
+    ],
+    getIndexPattern: () => {
+      return 'world_countries';
+    },
+  };
+  return {
+    getIndexPatternService() {
+      return {
+        get: async () => {
+          return mockDataView;
+        },
+        getDefaultDataView: async () => {
+          return mockDefaultDataView;
+        },
+      };
+    },
+  };
+});
+
+describe('CreateSourceEditor', () => {
+  test('should preview default data view on load', async () => {
+    const onSourceConfigChange = jest.fn();
+    render(<CreateSourceEditor onSourceConfigChange={onSourceConfigChange} />);
+    await waitFor(() =>
+      expect(onSourceConfigChange).toBeCalledWith({
+        columns: [
+          {
+            name: 'location',
+            type: 'geo_point',
+          },
+        ],
+        dateField: '@timestamp',
+        esql: 'from logs | keep location | limit 10000',
+      })
+    );
+  });
+
+  test('should preview requested data view on load when mostCommonDataViewId prop provided', async () => {
+    const onSourceConfigChange = jest.fn();
+    render(
+      <CreateSourceEditor
+        onSourceConfigChange={onSourceConfigChange}
+        mostCommonDataViewId="123abc"
+      />
+    );
+    await waitFor(() =>
+      expect(onSourceConfigChange).toBeCalledWith({
+        columns: [
+          {
+            name: 'geometry',
+            type: 'geo_shape',
+          },
+        ],
+        dateField: undefined,
+        esql: 'from world_countries | keep geometry | limit 10000',
+      })
+    );
+  });
+});

x-pack/plugins/maps/public/classes/sources/esql_source/create_source_editor.tsx

@@ -7,11 +7,12 @@
 
 import React, { useEffect, useState } from 'react';
 import { EuiSkeletonText } from '@elastic/eui';
+import { DataViewField } from '@kbn/data-views-plugin/public';
 import { ES_GEO_FIELD_TYPE } from '../../../../common/constants';
 import type { ESQLSourceDescriptor } from '../../../../common/descriptor_types';
 import { getIndexPatternService } from '../../../kibana_services';
 import { ESQLEditor } from './esql_editor';
-import { ESQL_GEO_POINT_TYPE } from './esql_utils';
+import { ESQL_GEO_POINT_TYPE, ESQL_GEO_SHAPE_TYPE } from './esql_utils';
 
 interface Props {
   mostCommonDataViewId?: string;
@@ -39,12 +40,17 @@ export function CreateSourceEditor(props: Props) {
         }
 
         if (dataView) {
-          let geoField: string | undefined;
+          let geoField: DataViewField | undefined;
           const initialDateFields: string[] = [];
           for (let i = 0; i < dataView.fields.length; i++) {
             const field = dataView.fields[i];
-            if (!geoField && ES_GEO_FIELD_TYPE.GEO_POINT === field.type) {
-              geoField = field.name;
+            if (
+              !geoField &&
+              [ES_GEO_FIELD_TYPE.GEO_POINT, ES_GEO_FIELD_TYPE.GEO_SHAPE].includes(
+                field.type as ES_GEO_FIELD_TYPE
+              )
+            ) {
+              geoField = field;
             } else if ('date' === field.type) {
               initialDateFields.push(field.name);
             }
@@ -57,14 +63,19 @@ export function CreateSourceEditor(props: Props) {
             } else if (initialDateFields.length) {
               initialDateField = initialDateFields[0];
             }
-            const initialEsql = `from ${dataView.getIndexPattern()} | keep ${geoField} | limit 10000`;
+            const initialEsql = `from ${dataView.getIndexPattern()} | keep ${
+              geoField.name
+            } | limit 10000`;
             setDateField(initialDateField);
             setEsql(initialEsql);
             props.onSourceConfigChange({
               columns: [
                 {
-                  name: geoField,
-                  type: ESQL_GEO_POINT_TYPE,
+                  name: geoField.name,
+                  type:
+                    geoField.type === ES_GEO_FIELD_TYPE.GEO_SHAPE
+                      ? ESQL_GEO_SHAPE_TYPE
+                      : ESQL_GEO_POINT_TYPE,
                 },
               ],
               dateField: initialDateField,

x-pack/plugins/maps/public/classes/sources/esql_source/esql_source.test.ts

@@ -0,0 +1,51 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { ESQLSource } from './esql_source';
+import { VECTOR_SHAPE_TYPE } from '../../../../common/constants';
+
+describe('getSupportedShapeTypes', () => {
+  test('should return point for geo_point column', async () => {
+    const descriptor = ESQLSource.createDescriptor({
+      esql: 'from kibana_sample_data_logs | keep geo.coordinates | limit 10000',
+      columns: [
+        {
+          name: 'geo.coordinates',
+          type: 'geo_point',
+        },
+      ],
+    });
+    const esqlSource = new ESQLSource(descriptor);
+    expect(await esqlSource.getSupportedShapeTypes()).toEqual([VECTOR_SHAPE_TYPE.POINT]);
+  });
+
+  test('should return all geometry types for geo_shape column', async () => {
+    const descriptor = ESQLSource.createDescriptor({
+      esql: 'from world_countries | keep geometry | limit 10000',
+      columns: [
+        {
+          name: 'geometry',
+          type: 'geo_shape',
+        },
+      ],
+    });
+    const esqlSource = new ESQLSource(descriptor);
+    expect(await esqlSource.getSupportedShapeTypes()).toEqual([
+      VECTOR_SHAPE_TYPE.POINT,
+      VECTOR_SHAPE_TYPE.LINE,
+      VECTOR_SHAPE_TYPE.POLYGON,
+    ]);
+  });
+
+  test('should fallback to point when geometry column can not be found', async () => {
+    const descriptor = ESQLSource.createDescriptor({
+      esql: 'from world_countries | keep geometry | limit 10000',
+    });
+    const esqlSource = new ESQLSource(descriptor);
+    expect(await esqlSource.getSupportedShapeTypes()).toEqual([VECTOR_SHAPE_TYPE.POINT]);
+  });
+});

x-pack/plugins/maps/public/classes/sources/esql_source/esql_source.tsx

@@ -31,7 +31,12 @@ import type { IField } from '../../fields/field';
 import { InlineField } from '../../fields/inline_field';
 import { getData, getUiSettings } from '../../../kibana_services';
 import { convertToGeoJson } from './convert_to_geojson';
-import { getFieldType, getGeometryColumnIndex } from './esql_utils';
+import {
+  getFieldType,
+  getGeometryColumnIndex,
+  ESQL_GEO_POINT_TYPE,
+  ESQL_GEO_SHAPE_TYPE,
+} from './esql_utils';
 import { UpdateSourceEditor } from './update_source_editor';
 
 type ESQLSourceSyncMeta = Pick<
@@ -114,7 +119,18 @@ export class ESQLSource extends AbstractVectorSource implements IVectorSource {
   }
 
   async getSupportedShapeTypes() {
-    return [VECTOR_SHAPE_TYPE.POINT];
+    let geomtryColumnType = ESQL_GEO_POINT_TYPE;
+    try {
+      const index = getGeometryColumnIndex(this._descriptor.columns);
+      if (index > -1) {
+        geomtryColumnType = this._descriptor.columns[index].type;
+      }
+    } catch (error) {
+      // errors for missing geometry columns surfaced in UI by data loading
+    }
+    return geomtryColumnType === ESQL_GEO_SHAPE_TYPE
+      ? [VECTOR_SHAPE_TYPE.POINT, VECTOR_SHAPE_TYPE.LINE, VECTOR_SHAPE_TYPE.POLYGON]
+      : [VECTOR_SHAPE_TYPE.POINT];
   }
 
   supportsJoins() {

x-pack/plugins/maps/public/classes/sources/esql_source/esql_utils.test.ts

@@ -0,0 +1,22 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { isGeometryColumn } from './esql_utils';
+
+describe('isGeometryColumn', () => {
+  test('should return true for geo_point columns', () => {
+    expect(isGeometryColumn({ name: 'myColumn', type: 'geo_point' })).toBe(true);
+  });
+
+  test('should return true for geo_shape columns', () => {
+    expect(isGeometryColumn({ name: 'myColumn', type: 'geo_shape' })).toBe(true);
+  });
+
+  test('should return false for non-geometry columns', () => {
+    expect(isGeometryColumn({ name: 'myColumn', type: 'string' })).toBe(false);
+  });
+});

x-pack/plugins/maps/public/classes/sources/esql_source/esql_utils.ts

@@ -12,6 +12,7 @@ import type { ESQLColumn } from '@kbn/es-types';
 import { getData, getIndexPatternService } from '../../../kibana_services';
 
 export const ESQL_GEO_POINT_TYPE = 'geo_point';
+export const ESQL_GEO_SHAPE_TYPE = 'geo_shape';
 
 const NO_GEOMETRY_COLUMN_ERROR_MSG = i18n.translate(
   'xpack.maps.source.esql.noGeometryColumnErrorMsg',
@@ -20,8 +21,8 @@ const NO_GEOMETRY_COLUMN_ERROR_MSG = i18n.translate(
   }
 );
 
-function isGeometryColumn(column: ESQLColumn) {
-  return column.type === ESQL_GEO_POINT_TYPE;
+export function isGeometryColumn(column: ESQLColumn) {
+  return [ESQL_GEO_POINT_TYPE, ESQL_GEO_SHAPE_TYPE].includes(column.type);
 }
 
 export function verifyGeometryColumn(columns: ESQLColumn[]) {