search_v3.py 10.9 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# Copyright 2021 Schlumberger
#
# 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.

from fastapi import APIRouter, Depends
from odes_search.models import (
    QueryRequest,
fabian serin's avatar
fabian serin committed
18
19
20
21
22
    CursorQueryResponse,
    CursorQueryRequest,
    BaseModel,
    Field,
    Optional)
23
from app.clients.search_service_client import get_search_service
fabian serin's avatar
fabian serin committed
24
from app.routers.common_parameters import REQUIRED_ROLES_READ
25
26
27
from app.utils import Context
import app.routers.search.search_wrapper as search_wrapper
from .search import (
fabian serin's avatar
fabian serin committed
28
    LIMIT,
29
30
31
32
33
34
    query_type,
    SearchQuery,
    get_ctx,
    query_type_returned_fields,
    basic_query_request,
    basic_query_request_with_cursor)
fabian serin's avatar
fabian serin committed
35

36
37
router = APIRouter()

fabian serin's avatar
fabian serin committed
38
# osdu kind
39
40
41
OSDU_WELLBORE_KIND = '*:wks:master-data--Wellbore:*'
OSDU_WELLLOG_KIND = '*:wks:work-product-component--WellLog:*'
OSDU_WELLBOREMARKERSET_KIND = '*:wks:work-product-component--WellboreMarkerSet:*'
fabian serin's avatar
fabian serin committed
42
OSDU_WELLBORETRAJECTORY_KIND = '*:wks:work-product-component--WellboreTrajectory:*'
43
44
WELLBORE_RELATIONSHIP = "WellboreID"

fabian serin's avatar
fabian serin committed
45

fabian serin's avatar
fabian serin committed
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
class SearchQueryRequest(BaseModel):
    # Used by as input, w/o kind, etc...
    limit: "Optional[int]" = Field(None, alias="limit")
    query: "Optional[str]" = Field(None, alias="query")
    cursor: "Optional[str]" = Field(None, alias="cursor")
    offset: "Optional[int]" = Field(None, alias="offset")


SearchQueryRequest.update_forward_refs()
DEFAULT_SEARCHQUERYREQUEST = SearchQueryRequest(limit=None, query=None, cursor=None, offset=None)


class SimpleCursorQueryRequest(BaseModel):
    # Used by as input, w/o kind, etc...
    limit: "Optional[int]" = Field(None, alias="limit")
    query: "Optional[str]" = Field(None, alias="query")
    cursor: "Optional[str]" = Field(None, alias="cursor")


SimpleCursorQueryRequest.update_forward_refs()
DEFAULT_CURSORQUERYREQUEST = SimpleCursorQueryRequest(limit=None, query=None, cursor=None)


class SimpleOffsetQueryRequest(BaseModel):
    limit: "Optional[int]" = Field(None, alias="limit")
    query: "Optional[str]" = Field(None, alias="query")
    offset: "Optional[int]" = Field(None, alias="offset")


SimpleOffsetQueryRequest.update_forward_refs()
DEFAULT_QUERYREQUEST = SimpleOffsetQueryRequest(limit=None, query=None, offset=None)


79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def create_relationships_id_str(data_type: str, id: str):
    return f'data.{data_type}:\"{id}\"'


def added_query(generated_query: str, user_query: str = None):
    if user_query:
        query = f'{generated_query} AND ({user_query})'
    else:
        query = generated_query
    return query


def added_relationships_query(id: str, data_type: str, query: str = None):
    relationships_id = create_relationships_id_str(data_type, id)
    return added_query(relationships_id, query)


async def query_request_with_specific_attribute(query_type: str, attribute: str, attribute_kind: str, kind: str,
                                                data_type: str,
                                                ctx: Context, query: str = None):
    query_request = QueryRequest(kind=attribute_kind,
                                 query=attribute,
                                 returnedFields=['id'])

    client = await get_search_service(ctx)
    query_result = await search_wrapper.SearchWrapper.query_cursorless(
        search_service=client,
        data_partition_id=ctx.partition_id,
        query_request=query_request)

    response = CursorQueryResponse.parse_obj(query_result.dict())

    if not response.results:
        return query_result

    relationships_ids = [create_relationships_id_str(data_type, r["id"]) for r in response.results]
fabian serin's avatar
fabian serin committed
115
    id_list = ' OR '.join(relationships_ids)  # [a, b, c] => 'a OR b OR c'
116
117
118
119
120
121
122
123
124
125
126
127
128

    query = added_query(id_list, query)

    returned_fields = query_type_returned_fields(query_type)
    query_request = QueryRequest(kind=kind,
                                 query=query,
                                 returnedFields=[returned_fields])
    return await search_wrapper.SearchWrapper.query_cursorless(
        search_service=client,
        data_partition_id=ctx.partition_id,
        query_request=query_request)


129
130
def update_query_with_names_based_search(names: str = None, user_query: str = None,
                                         name_field: str = "data.FacilityName") -> str:
131
132
    if names is None:
        return user_query
fabian serin's avatar
fabian serin committed
133
    generated_query = f"{name_field}:{names}"
fabian serin's avatar
fabian serin committed
134
135
136
    return added_query(generated_query, user_query)


137
138
139
140
141
142
143
144
def update_query_with_nested_names_based_search(array_field: str, nested_field: str, names: str = None,
                                                user_query: str = None) -> str:
    if names is None:
        return user_query
    generated_query = f"(nested({array_field}, ({nested_field}:({names}))))"
    return added_query(generated_query, user_query)


fabian serin's avatar
fabian serin committed
145
146
147
def escape_forbidden_characters_for_search(input_str: str) -> str:
    # Reserved character are listed here https://community.opengroup.org/osdu/documentation/-/blob/master/platform/tutorials/core-services/SearchService.md
    # ? and * are allowed for wildcard search
148
149
    if input_str is None:
        return None
fabian serin's avatar
fabian serin committed
150
151
152
153
154
155
156
157
158
159
    reserved_char_list = ['+', '-', '=', '>', '<', '!', '(', ')', '{', '}', '[', ']', '^', '"', '~',
                          ':', '\\', '/']

    def escape_char(input_char: str, reserved_char_list: [str]) -> str:
        return input_char if input_char not in reserved_char_list else f"\\{input_char}"

    result_str = ''.join([escape_char(char, reserved_char_list) for char in input_str])
    return result_str


fabian serin's avatar
cleanup    
fabian serin committed
160
161
async def query_request_with_cursor(query_type: str, kind: str, ctx: Context, query: SimpleCursorQueryRequest = None):
    returned_fields = query_type_returned_fields(query_type)
fabian serin's avatar
fabian serin committed
162
163
164
165
166
167
168
169
170
171
172
    query_request = CursorQueryRequest(kind=kind,
                                       limit=query.limit or LIMIT,
                                       query=query.query,
                                       returnedFields=[returned_fields],
                                       cursor=query.cursor)
    client = await get_search_service(ctx)
    return await client.query_with_cursor(
        data_partition_id=ctx.partition_id,
        cursor_query_request=query_request)


173
174
async def query_request_with_offset(query_type: str, kind: str, ctx: Context, query: SimpleOffsetQueryRequest = None,
                                    custom_returned_fields: [str] = None):
fabian serin's avatar
cleanup    
fabian serin committed
175
    returned_fields = query_type_returned_fields(query_type)
fabian serin's avatar
fabian serin committed
176
177
178
179
180
181
182
183
184
185
186
187

    query_request = QueryRequest(kind=kind,
                                 limit=query.limit or LIMIT,
                                 query=query.query,
                                 returnedFields=[returned_fields],
                                 offset=query.offset)
    client = await get_search_service(ctx)
    return await client.query(
        data_partition_id=ctx.partition_id,
        query_request=query_request)


fabian serin's avatar
cleanup    
fabian serin committed
188
async def query_request(query_type: str, kind: str, ctx: Context, query: SearchQueryRequest = None):
fabian serin's avatar
fabian serin committed
189
190
191
192
    # use offset if not not none else use cursor
    query_as_dict = query.dict(exclude_none=True, exclude_unset=True)
    if query.offset is not None:
        cursor_query = SimpleOffsetQueryRequest(**query_as_dict)
fabian serin's avatar
cleanup    
fabian serin committed
193
        return await query_request_with_offset(query_type, kind, ctx, cursor_query)
fabian serin's avatar
fabian serin committed
194
195

    cursor_query = SimpleCursorQueryRequest(**query_as_dict)
fabian serin's avatar
cleanup    
fabian serin committed
196
    return await query_request_with_cursor(query_type, kind, ctx, cursor_query)
fabian serin's avatar
fabian serin committed
197
198


199
200
201
202
@router.post('/query/wellbores', summary='Query with cursor, get wellbores',
             description=f"""Get all Wellbores object.  <p>The wellbore kind is {OSDU_WELLBORE_KIND}
        returns all records directly based on existing schemas</p>{REQUIRED_ROLES_READ}""",
             response_model=CursorQueryResponse)
fabian serin's avatar
fabian serin committed
203
204
async def query_wellbores(body: SearchQueryRequest = DEFAULT_QUERYREQUEST, ctx: Context = Depends(get_ctx)):
    return await query_request(query_type, OSDU_WELLBORE_KIND, ctx, body)
205

fabian serin's avatar
fabian serin committed
206

207
208
209
210
211
@router.post('/query/wellbores/{wellboreId}/welllogs', summary='Query with cursor, search WellLogs by wellbore ID',
             description=f"""Get all WellLogs object using its relationship Wellbore ID.  <p>All WellLogs linked to this
            specific ID will be returned</p>
            <p>The WellLogs kind is {OSDU_WELLLOG_KIND} returns all records directly based on existing schemas</p>{REQUIRED_ROLES_READ}""",
             response_model=CursorQueryResponse)
fabian serin's avatar
fabian serin committed
212
async def query_welllogs_bywellbore(wellboreId: str, body: SearchQueryRequest = DEFAULT_QUERYREQUEST,
fabian serin's avatar
fabian serin committed
213
                                    ctx: Context = Depends(get_ctx)):
fabian serin's avatar
fabian serin committed
214
215
    body.query = added_relationships_query(wellboreId, WELLBORE_RELATIONSHIP, body.query)
    return await query_request(query_type, OSDU_WELLLOG_KIND, ctx, body)
216
217
218
219
220
221
222
223


@router.post('/query/wellbore/{wellboreAttribute}/welllogs',
             summary='Query with cursor, search WellLogs by wellbore attribute',
             description=f"""Get all WellLogs object using a specific attribute of Wellbores.  <p>All WellLogs linked to Wellbores
            with this specific attribute will be returned</p>
            <p>The WellLogs kind is {OSDU_WELLLOG_KIND} returns all records directly based on existing schemas</p>{REQUIRED_ROLES_READ}""",
             response_model=CursorQueryResponse)
224
async def query_welllogs_bywellboreattribute(wellboreAttribute: str, body: SearchQuery = SearchQuery(query=None),
fabian serin's avatar
fabian serin committed
225
226
227
                                             ctx: Context = Depends(get_ctx)):
    return await query_request_with_specific_attribute(query_type, wellboreAttribute, OSDU_WELLBORE_KIND,
                                                       OSDU_WELLLOG_KIND,
228
229
230
231
                                                       WELLBORE_RELATIONSHIP, ctx,
                                                       body.query)


fabian serin's avatar
fabian serin committed
232
233
@router.post('/query/wellbores/{wellboreId}/wellboremarkersets',
             summary='Query with cursor, search wellbore markersets by wellbore ID',
234
235
236
237
             description=f"""Get all Wellbore Markersets objects using its relationship Wellbore ID.  <p>All Markers linked to this
            specific ID will be returned</p>
            <p>The Wellbore Markerset kind is {OSDU_WELLBOREMARKERSET_KIND} returns all records directly based on existing schemas</p>{REQUIRED_ROLES_READ}""",
             response_model=CursorQueryResponse)
fabian serin's avatar
fabian serin committed
238
async def query_markers_bywellbore(wellboreId: str, body: SearchQueryRequest = DEFAULT_QUERYREQUEST,
239
                                   ctx: Context = Depends(get_ctx)):
fabian serin's avatar
fabian serin committed
240
241
    body.query = added_relationships_query(wellboreId, WELLBORE_RELATIONSHIP, body.query)
    return await query_request(query_type, OSDU_WELLBOREMARKERSET_KIND, ctx, body)