feat: add data elements to ES index

querybook/server/datasources/data_element.py

@@ -7,8 +7,8 @@
     "/data_element/keyword/",
     methods=["GET"],
 )
-def get_data_elements_by_keyword(keyword):
-    data_elements = logic.get_data_elements_by_keyword(keyword=keyword)
+def search_data_elements_by_keyword(keyword):
+    data_elements = logic.search_data_elements_by_keyword(keyword=keyword)
     return [data_element.name for data_element in data_elements]
 
 

querybook/server/lib/elasticsearch/search_table.py

@@ -5,7 +5,7 @@
     combine_keyword_and_filter_query,
 )
 
-FILTERS_TO_AND = ["data_elements"]
+FILTERS_TO_AND = ["tags", "data_elements"]
 
 
 def _get_potential_exact_schema_table_name(keywords):

querybook/server/logic/data_element.py

@@ -25,12 +25,11 @@ def get_data_element_by_name(name: str, session=None):
 
 
 @with_session
-def get_data_elements_by_keyword(keyword, limit=10, session=None):
+def search_data_elements_by_keyword(keyword: str, limit=10, session=None):
     return (
         session.query(DataElement)
         .filter(DataElement.name.like("%" + keyword + "%"))
         .order_by(DataElement.name.asc())
-        .offset(0)
         .limit(limit)
         .all()
     )

querybook/server/logic/elasticsearch.py

@@ -47,8 +47,6 @@
 from models.user import User
 from models.datadoc import DataCellType
 from models.board import Board
-from models.data_element import DataElement
-from models.metastore import DataTable
 
 
 LOG = get_logger(__file__)
@@ -540,19 +538,6 @@ def get_table_weight(table_id: int, session=None) -> int:
     return int(math.log2(((num_impressions + num_samples * 10) + 1) + boost_score))
 
 
-def get_unique_data_elements_from_table(table: DataTable) -> list[DataElement]:
-    unique_names = set()
-    data_elements = []
-    for c in table.columns:
-        for data_element in c.data_elements:
-            if data_element.name in unique_names:
-                continue
-            data_elements.append(data_element)
-            unique_names.add(data_element.name)
-
-    return data_elements
-
-
 @with_session
 def table_to_es(table, fields=None, session=None):
     schema = table.data_schema
@@ -561,7 +546,7 @@ def table_to_es(table, fields=None, session=None):
     full_name = "{}.{}".format(schema_name, table_name)
 
     # columns may be associated with the same data element
-    data_elements = get_unique_data_elements_from_table(table)
+    data_elements = {d.name: d for c in table.columns for d in c.data_elements}.values()
 
     def get_table_description():
         return (

querybook/tests/test_lib/test_elasticsearch/test_elasticsearch.py

@@ -339,11 +339,26 @@ def _get_data_schema_mock(self):
         mock_data_schema.name = self.SCHEMA_NAME
         return mock_data_schema
 
+    def _get_data_element_mock(self, name: str, description: str):
+        mock_de = MagicMock()
+        mock_de.name = name
+        mock_de.description = description
+        return mock_de
+
     def _get_columns_mock(self):
         mock_col_a = MagicMock()
         mock_col_a.name = "col_a"
+        mock_col_a.description = "col_a_description"
+        mock_col_a.data_elements = [
+            self._get_data_element_mock("de_a", "de_a_description"),
+            self._get_data_element_mock("de_b", "de_b_description"),
+        ]
         mock_col_b = MagicMock()
         mock_col_b.name = "col_b"
+        mock_col_b.description = "col_b_description"
+        mock_col_b.data_elements = [
+            self._get_data_element_mock("de_a", "de_a_description"),
+        ]
         return [mock_col_a, mock_col_b]
 
     def _get_table_mock(self):
@@ -393,6 +408,9 @@ def test_table_to_es(self):
             "description": self.TABLE_DESCRIPTION,
             "created_at": CREATED_AT_EPOCH,
             "columns": ["col_a", "col_b"],
+            "column_descriptions": ["col_a_description", "col_b_description"],
+            "data_elements": ["de_a", "de_b"],
+            "data_element_descriptions": ["de_a_description", "de_b_description"],
             "golden": False,
             "importance_score": self.TABLE_WEIGHT,
             "tags": ["tag_1", "tag_2"],

querybook/webapp/components/DataElement/DataElement.scss

@@ -1,29 +1,3 @@
 .DataElementCard {
     max-width: 400px;
 }
-
-.DataElementGroupSelect {
-    .data-element-list {
-        display: flex;
-        align-items: center;
-        flex-wrap: wrap;
-        gap: 6px;
-        margin-bottom: 8px;
-
-        .Tag {
-            margin-left: 0 !important;
-            font-size: var(--xxsmall-text-size);
-        }
-    }
-}
-
-.DataElementSelect {
-    min-width: 180px;
-    margin: 0px;
-    padding: 0px;
-    > div {
-        width: 100%;
-    }
-    border: 1px solid transparent;
-    border-radius: var(--border-radius-sm);
-}

querybook/webapp/components/DataTableNavigator/DataTableNavigatorSearch.tsx

@@ -2,7 +2,7 @@ import { startCase } from 'lodash';
 import React, { useCallback, useMemo, useRef } from 'react';
 import { useDispatch, useSelector } from 'react-redux';
 
-import { TableTagGroupSelect } from 'components/DataTableTags/TableTagGroupSelect';
+import { EntitySelect } from 'components/Search/EntitySelect';
 import { ComponentType, ElementType } from 'const/analytics';
 import { useToggleState } from 'hooks/useToggleState';
 import { trackClick } from 'lib/analytics';
@@ -12,6 +12,7 @@ import {
 } from 'redux/dataTableSearch/action';
 import { ITableSearchFilters } from 'redux/dataTableSearch/types';
 import { IStoreState } from 'redux/store/types';
+import { TableTagResource } from 'resource/table';
 import { SoftButton } from 'ui/Button/Button';
 import { IconButton } from 'ui/Button/IconButton';
 import { OrderByButton } from 'ui/OrderByButton/OrderByButton';
@@ -108,9 +109,12 @@ export const DataTableNavigatorSearch: React.FC<{
                         />
                     </SearchFilterRow>
                     <SearchFilterRow title="Tags">
-                        <TableTagGroupSelect
-                            tags={searchFilters?.tags}
-                            updateTags={updateTags}
+                        <EntitySelect
+                            selectedEntities={searchFilters?.tags || []}
+                            loadEntities={TableTagResource.search}
+                            onEntitiesChange={updateTags}
+                            placeholder="Tag name"
+                            mini
                         />
                     </SearchFilterRow>
                 </div>

querybook/webapp/components/DataTableTags/CreateDataTableTag.tsx

@@ -1,12 +1,14 @@
 import * as React from 'react';
 import { useDispatch } from 'react-redux';
 
+import { EntitySelect } from 'components/Search/EntitySelect';
 import { ITag } from 'const/tag';
 import { Dispatch } from 'redux/store/types';
 import { createTableTag } from 'redux/tag/action';
+import { TableTagResource } from 'resource/table';
 import { IconButton } from 'ui/Button/IconButton';
 
-import { TableTagSelect } from './TableTagSelect';
+import { isTagValid } from './utils';
 
 import './CreateDataTableTag.scss';
 
@@ -43,10 +45,15 @@ export const CreateDataTableTag: React.FunctionComponent<IProps> = ({
         <div className="CreateDataTableTag flex-row">
             {showSelect ? (
                 <div className="CreateDataTableTag-input flex-row">
-                    <TableTagSelect
-                        onSelect={handleCreateTag}
-                        existingTags={existingTags}
+                    <EntitySelect
                         creatable
+                        mini
+                        selectedEntities={existingTags || []}
+                        loadEntities={TableTagResource.search}
+                        onSelect={handleCreateTag}
+                        isEntityValid={isTagValid}
+                        placeholder="alphanumeric only"
+                        showSelected={false}
                     />
                 </div>
             ) : (

querybook/webapp/components/DataTableTags/utils.ts

@@ -26,3 +26,8 @@ export function useRankedTags(tags: ITag[]) {
         [tags]
     );
 }
+
+export function isTagValid(val: string) {
+    const regex = /^[a-z0-9_ ]{1,255}$/i;
+    return regex.test(val);
+}

querybook/webapp/components/Search/EntitySelect.scss

@@ -0,0 +1,21 @@
+.EntitySelect {
+    min-width: 180px;
+
+    .invalid {
+        border: 1px solid var(--color-false);
+        border-radius: var(--border-radius-sm);
+    }
+
+    .entity-list {
+        display: flex;
+        align-items: center;
+        flex-wrap: wrap;
+        gap: 6px;
+        margin-bottom: 8px;
+
+        .Tag {
+            margin-left: 0 !important;
+            font-size: var(--xxsmall-text-size);
+        }
+    }
+}