search_v3.py 10.3 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)


fabian serin's avatar
fabian serin committed
129

fabian serin's avatar
fabian serin committed
130
131
def update_query_with_names_based_search(names: str = None, user_query: str = None, name_field = "data.FacilityName") -> str:
    generated_query = f"{name_field}:{names}"
fabian serin's avatar
fabian serin committed
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
    return added_query(generated_query, user_query)


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
    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


async def query_request_with_cursor(query_type: str, kind: str, ctx: Context, query: SimpleCursorQueryRequest = None):
    returned_fields = query_type_returned_fields(query_type)
    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)


async def query_request_with_offset(query_type: str, kind: str, ctx: Context, query: SimpleOffsetQueryRequest = None):
    returned_fields = query_type_returned_fields(query_type)

    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)


async def query_request(query_type: str, kind: str, ctx: Context, query: SearchQueryRequest = None):
    # 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)
        return await query_request_with_offset(query_type, kind, ctx, cursor_query)

    cursor_query = SimpleCursorQueryRequest(**query_as_dict)
    return await query_request_with_cursor(query_type, kind, ctx, cursor_query)


186
187
188
189
@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
190
191
async def query_wellbores(body: SearchQueryRequest = DEFAULT_QUERYREQUEST, ctx: Context = Depends(get_ctx)):
    return await query_request(query_type, OSDU_WELLBORE_KIND, ctx, body)
192

fabian serin's avatar
fabian serin committed
193

194
195
196
197
198
@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
199
async def query_welllogs_bywellbore(wellboreId: str, body: SearchQueryRequest = DEFAULT_QUERYREQUEST,
fabian serin's avatar
fabian serin committed
200
                                    ctx: Context = Depends(get_ctx)):
fabian serin's avatar
fabian serin committed
201
202
    body.query = added_relationships_query(wellboreId, WELLBORE_RELATIONSHIP, body.query)
    return await query_request(query_type, OSDU_WELLLOG_KIND, ctx, body)
203
204
205
206
207
208
209
210


@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)
211
async def query_welllogs_bywellboreattribute(wellboreAttribute: str, body: SearchQuery = SearchQuery(query=None),
fabian serin's avatar
fabian serin committed
212
213
214
                                             ctx: Context = Depends(get_ctx)):
    return await query_request_with_specific_attribute(query_type, wellboreAttribute, OSDU_WELLBORE_KIND,
                                                       OSDU_WELLLOG_KIND,
215
216
217
218
                                                       WELLBORE_RELATIONSHIP, ctx,
                                                       body.query)


fabian serin's avatar
fabian serin committed
219
220
@router.post('/query/wellbores/{wellboreId}/wellboremarkersets',
             summary='Query with cursor, search wellbore markersets by wellbore ID',
221
222
223
224
             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
225
async def query_markers_bywellbore(wellboreId: str, body: SearchQueryRequest = DEFAULT_QUERYREQUEST,
226
                                   ctx: Context = Depends(get_ctx)):
fabian serin's avatar
fabian serin committed
227
228
    body.query = added_relationships_query(wellboreId, WELLBORE_RELATIONSHIP, body.query)
    return await query_request(query_type, OSDU_WELLBOREMARKERSET_KIND, ctx, body)