fix: fix table post/put api bug (#172)

* fix: fix table post/put api bug

Signed-off-by: tianru zhou <tianru.zhou@databricks.com>

* fix syntax issue v1

Signed-off-by: tianru zhou <tianru.zhou@databricks.com>

* address PR comments

Signed-off-by: tianru zhou <tianru.zhou@databricks.com>

* fix lint test failure

Signed-off-by: tianru zhou <tianru.zhou@databricks.com>

search_service/api/document.py

@@ -2,10 +2,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
+import json
 
 from http import HTTPStatus
 from typing import Tuple, Any
-
+from ast import literal_eval
 from flasgger import swag_from
 from flask_restful import Resource, reqparse
 from search_service.proxy import get_proxy_client
@@ -58,7 +59,7 @@ def post(self) -> Tuple[Any, int]:
          :param data: list of data objects to be indexed in Elasticsearch
          :return: name of new index
          """
-        self.parser.add_argument('data', required=True)
+        self.parser.add_argument('data', required=True, action='append')
         args = self.parser.parse_args()
 
         try:
@@ -78,11 +79,13 @@ def put(self) -> Tuple[Any, int]:
         :param data: list of data objects to be indexed in Elasticsearch
         :return: name of index
         """
-        self.parser.add_argument('data', required=True)
+        self.parser.add_argument('data', required=True, action='append')
         args = self.parser.parse_args()
 
         try:
-            data = self.schema(many=True, strict=False).loads(args.get('data')).data
+            table_dict_list = [literal_eval(table_str) for table_str in args.get('data')]
+            table_list_json = json.dumps(table_dict_list)
+            data = self.schema(many=True, strict=False).loads(table_list_json).data
             results = self.proxy.update_document(data=data, index=args.get('index'))
             return results, HTTPStatus.OK
         except RuntimeError as e:

search_service/api/swagger_doc/document/table_post.yml

@@ -11,19 +11,18 @@ parameters:
       type: string
       default: 'table_search_index'
     required: false
-  - name: body
-    in: body
-    schema:
-      type: object
-      name: data
-      properties:
-        data:
-          type: array
-          description: 'List of tables'
-          items:
-            $ref: '#/components/schemas/TableFields'
-    description: 'Tables to create'
-    required: true
+requestBody:
+  content:
+    'application/json':
+      schema:
+        type: object
+        properties:
+          data:
+            type: array
+            items:
+              $ref: '#/components/schemas/TableFields'
+  description: 'Tables to create'
+  required: true
 responses:
   200:
     description: Empty json response

search_service/api/swagger_doc/document/table_put.yml

@@ -11,19 +11,18 @@ parameters:
       type: string
       default: 'table_search_index'
     required: false
-  - name: body
-    in: body
-    schema:
-      type: object
-      name: data
-      properties:
-        data:
-          type: array
-          description: 'List of tables'
-          items:
-            $ref: '#/components/schemas/TableFields'
-    description: 'Tables to update'
-    required: true
+requestBody:
+  content:
+    'application/json':
+      schema:
+        type: object
+        properties:
+          data:
+            type: array
+            items:
+              $ref: '#/components/schemas/TableFields'
+  description: 'Tables to update'
+  required: true
 responses:
   200:
     description: Empty json response

search_service/api/swagger_doc/document/user_post.yml

@@ -11,19 +11,18 @@ parameters:
       type: string
       default: user_search_index
     required: false
-  - name: body
-    in: body
-    schema:
-      type: object
-      name: data
-      properties:
-        data:
-          type: array
-          description: 'List of users'
-          items:
-            $ref: '#/components/schemas/UserFields'
-    description: 'Users to create'
-    required: true
+requestBody:
+  content:
+    'application/json':
+      schema:
+        type: object
+        properties:
+          data:
+            type: array
+            items:
+              $ref: '#/components/schemas/UserFields'
+  description: 'Users to create'
+  required: true
 responses:
   200:
     description: Empty json response

search_service/api/swagger_doc/document/user_put.yml

@@ -11,19 +11,18 @@ parameters:
       type: string
       default: user_search_index
     required: false
-  - name: body
-    in: body
-    schema:
-      type: object
-      name: data
-      properties:
-        data:
-          type: array
-          description: 'List of users'
-          items:
-            $ref: '#/components/schemas/UserFields'
-    description: 'Users to update'
-    required: true
+requestBody:
+  content:
+    'application/json':
+      schema:
+        type: object
+        properties:
+          data:
+            type: array
+            items:
+              $ref: '#/components/schemas/UserFields'
+  description: 'Users to update'
+  required: true
 responses:
   200:
     description: Empty json response

search_service/api/swagger_doc/template.yml

@@ -52,6 +52,10 @@ components:
     TableFields:
       type: object
       properties:
+        id:
+          type: string
+          description: 'elasticsearch doc id'
+          example: 'M81jD3cBdULZTSY96PSh'
         name:
           type: string
           description: 'name of table'
@@ -82,16 +86,54 @@ components:
               type: string
           description: 'list of column names'
           example: ['col1', 'col2']
+        column_descriptions:
+          type: array
+          items:
+            type: string
+          description: 'list of column descriptions'
+          example: ['column description1', 'column description2']
+        programmatic_descriptions:
+          type: array
+          items:
+            type: string
+          description: 'list of programmatic descriptions'
+          example: ['programmatic description1', 'programmatic description2']
         tags:
           type: array
           items:
-              type: string
+            type: object
+            properties:
+              tag_name:
+                type: string
           description: 'list of table tags'
-          example: ['tag2', 'tag1']
+          example: [{'tag_name': 'tag1'}, {'tag_name': 'tag2'}]
+        badges:
+          type: array
+          items:
+            type: object
+            properties:
+              tag_name:
+                type: string
+          description: 'list of table badges'
+          example: [{'tag_name': 'badge1'}, {'tag_name': 'badge2'}]
         last_updated_timestamp:
           type: integer
           description: 'table last updated time'
           example: 1568814420
+        display_name:
+          type: string
+          description: 'table display name'
+          example: 'display_name'
+        total_usage:
+          type: int
+          description: 'total usage'
+          example: 0
+        schema_description:
+          type: array
+          items:
+            type: string
+          description: 'list of schema descriptions'
+          example: ['schema description1', 'schema description2']
     DashboardFields:
       type: object
       properties:
@@ -130,6 +172,10 @@ components:
     UserFields:
       type: object
       properties:
+        id:
+          type: string
+          description: 'elasticsearch doc id'
+          example: 'M81jD3cBdULZTSY96PSh'
         name:
           type: string
           description: 'user name'

search_service/models/table.py

@@ -8,38 +8,52 @@
 
 from .base import Base
 from search_service.models.tag import Tag
+import time
 
 
 @attr.s(auto_attribs=True, kw_only=True)
 class Table(Base):
     """
     This represents the part of a table stored in the search proxy
     """
+    id: str
     database: str
     cluster: str
     schema: str
     name: str
     key: str
     display_name: Optional[str] = None
-    tags: List[Tag]
-    badges: List[Tag]
+    tags: Optional[List[Tag]] = None
+    badges: Optional[List[Tag]] = None
     description: Optional[str] = None
-    last_updated_timestamp: int
+    last_updated_timestamp: int = int(time.time())
     # The following properties are lightly-transformed properties from the normal table object:
-    column_names: List[str]
-    column_descriptions: List[str] = []
-    programmatic_descriptions: List[str] = []
+    column_names: Optional[List[str]] = None
+    column_descriptions: Optional[List[str]] = None
+    programmatic_descriptions: Optional[List[str]] = None
     # The following are search-only properties:
     total_usage: int = 0
     schema_description: Optional[str] = attr.ib(default=None)
 
     def get_id(self) -> str:
-        # uses the table key as the document id in ES
-        return self.key
+        return self.id
+
+    def get_attrs_dict(self) -> dict:
+        attrs_dict = self.__dict__.copy()
+        if self.tags is not None:
+            attrs_dict['tags'] = [str(tag) for tag in self.tags]
+        else:
+            attrs_dict['tags'] = None
+        if self.badges is not None:
+            attrs_dict['badges'] = [str(badge) for badge in self.badges]
+        else:
+            attrs_dict['badges'] = None
+        return attrs_dict
 
     @classmethod
     def get_attrs(cls) -> Set:
         return {
+            'id',
             'name',
             'key',
             'description',

search_service/models/tag.py

@@ -13,6 +13,9 @@ class Tag:
     def __init__(self, tag_name: str):
         self.tag_name = tag_name
 
+    def __str__(self) -> str:
+        return self.tag_name
+
 
 class TagSchema(AttrsSchema):
     class Meta:

search_service/models/user.py

@@ -17,14 +17,16 @@ class User(Base, CommonUser):
     This represents the part of a user stored in the search proxy
     """
     manager_email: Optional[str] = None
+    id: str
 
     def get_id(self) -> str:
         # uses the user email as the document id in ES
-        return self.email if self.email else ''
+        return self.id
 
     @classmethod
     def get_attrs(cls) -> Set:
         return {
+            'id',
             'full_name',
             'first_name',
             'last_name',

search_service/proxy/atlas.py

@@ -92,7 +92,8 @@ def _prepare_tables(self, response: EntityCollection, enhance_metadata: bool = F
 
             badges: List[Tag] = tags
 
-            table = Table(name=entity_name,
+            table = Table(id=f"{entity.typeName}://{db_cluster}.{db_name}/{entity_name}",
+                          name=entity_name,
                           key=f"{entity.typeName}://{db_cluster}.{db_name}/{entity_name}",
                           description=entity_attrs.get('description'),
                           cluster=db_cluster,

search_service/proxy/elasticsearch.py

@@ -116,14 +116,42 @@ def _get_search_result(self, page_index: int,
 
         for hit in response:
             try:
-                # ES hit: {'_d_': {'key': xxx...}
+                es_metadata = hit.__dict__.get('meta', {})
+                """
+                ES hit example:
+                {
+                    '_d_': {
+                        'name': 'name',
+                        'database': 'database',
+                        'schema': 'schema',
+                        'key': 'database://cluster.schema/name',
+                        'cluster': 'cluster',
+                        'column_descriptions': ['description1', 'description2'],
+                        'column_names': ['colname1', 'colname2'],
+                        'description': None,
+                        'display_name': 'display name',
+                        'last_updated_timestamp': 12345678,
+                        'programmatic_descriptions': [],
+                        'schema_description': None,
+                        'tags': ['tag1', 'tag2'],
+                        'badges': [],
+                        'total_usage': 0
+                    },
+                    'mata': {
+                        'index': 'table index',
+                        'id': 'table id',
+                        'type': 'type'
+                    }
+                }
+                """
                 es_payload = hit.__dict__.get('_d_', {})
                 if not es_payload:
                     raise Exception('The ES doc not contain required field')
                 result = {}
                 for attr, val in es_payload.items():
                     if attr in model.get_attrs():
                         result[attr] = self._get_instance(attr=attr, val=val)
+                result['id'] = self._get_instance(attr='id', val=es_metadata['id'])
 
                 results.append(model(**result))
             except Exception:
@@ -590,7 +618,7 @@ def _build_update_actions(self, data: List[Table], index_key: str) -> List[Dict[
 
         for item in data:
             actions.append({'update': {'_index': index_key, '_type': item.get_type(), '_id': item.get_id()}})
-            actions.append({'doc': item.__dict__})
+            actions.append({'doc': item.get_attrs_dict()})
         return actions
 
     def _build_delete_actions(self, data: List[str], index_key: str, type: str) -> List[Dict[str, Any]]: