async def get_collection(self, id: str, **kwargs) -> ORJSONResponse: """Get collection by id. Called with `GET /collections/{collectionId}`. Args: id: Id of the collection. Returns: Collection. """ request = kwargs["request"] pool = kwargs["request"].app.state.readpool async with pool.acquire() as conn: q, p = render( """ SELECT * FROM get_collection(:id::text); """, id=id, ) collection = await conn.fetchval(q, *p) if collection is None: raise NotFoundError links = await CollectionLinks(collection_id=id, request=request).get_links() collection["links"] = links return ORJSONResponse( Collection.construct(**collection).dict(exclude_none=True))
async def get_collection(self, collection_id: str, **kwargs) -> Collection: """Get collection by id. Called with `GET /collections/{collection_id}`. Args: id: Id of the collection. Returns: Collection. """ collection: Optional[Dict[str, Any]] request: Request = kwargs["request"] pool = request.app.state.readpool async with pool.acquire() as conn: q, p = render( """ SELECT * FROM get_collection(:id::text); """, id=collection_id, ) collection = await conn.fetchval(q, *p) if collection is None: raise NotFoundError(f"Collection {id} does not exist.") collection["links"] = await CollectionLinks( collection_id=collection_id, request=request).get_links(extra_links=collection.get("links")) return Collection(**collection)
async def fetch(self, query, kwargs): pool = await self.pool() start = time.time() logger.debug("Start time: %s Query: %s Args:%s", start, query, kwargs) rquery, args = render(query, **kwargs) async with pool.acquire() as con: try: r = await con.fetch(rquery, *args) except asyncpg.exceptions.UndefinedColumnError as e: raise ValueError(f"{e}") except asyncpg.exceptions.DataError as e: raise ValueError(f"{e}") except asyncpg.exceptions.CharacterNotInRepertoireError as e: raise ValueError(f"{e}") except Exception as e: logger.debug(f"Database Error: {e}") if str(e).startswith("ST_TileEnvelope"): raise HTTPException(status_code=422, detail=f"{e}") raise HTTPException(status_code=500, detail=f"{e}") logger.debug( "query took: %s results_firstrow: %s", time.time() - start, str(r and r[0])[0:500], ) return r
def test_recursion(): # without JoinComponent this would cause a recursion error, see #35 values = [f'foo_{i}' for i in range(5000)] query, query_args = render(':t', t=funcs.comma_sep(*values)) assert query.startswith('$1, $2') assert query.endswith('$4998, $4999, $5000') assert len(query_args) == 5000
async def test_manual_logic(conn): query, params = render( 'SELECT :select FROM users WHERE :where ORDER BY :order_by', select=select_fields('first_name'), where=V('created') > datetime(2021, 1, 1), order_by=V('last_name'), ) v = await conn.fetch(query, *params) assert ['Fred', 'Joe'] == [r[0] for r in v]
async def _search_base(self, search_request: PgstacSearch, **kwargs) -> Dict[str, Any]: """Cross catalog search (POST). Called with `POST /search`. Args: search_request: search request parameters. Returns: ItemCollection containing items which match the search criteria. """ request = kwargs["request"] pool = request.app.state.readpool # pool = kwargs["request"].app.state.readpool req = search_request.json(exclude_none=True) async with pool.acquire() as conn: q, p = render( """ SELECT * FROM search(:req::text::jsonb); """, req=req, ) items = await conn.fetchval(q, *p) next = items.pop("next", None) prev = items.pop("prev", None) collection = ItemCollection.construct(**items) cleaned_features = [] if collection.features is None or len(collection.features) == 0: raise NotFoundError("No features found") for feature in collection.features: feature = Item.construct(**feature) if "links" not in search_request.fields.exclude: links = await ItemLinks( collection_id=feature.collection, item_id=feature.id, request=request, ).get_links() feature.links = links exclude = search_request.fields.exclude if len(exclude) == 0: exclude = None include = search_request.fields.include if len(include) == 0: include = None feature = feature.dict(exclude_none=True, ) cleaned_features.append(feature) collection.features = cleaned_features collection.links = await PagingLinks( request=request, next=next, prev=prev, ).get_links() return collection
async def dbfunc(pool: pool, func: str, arg: Union[str, Dict]): """Wrap PLPGSQL Functions. Keyword arguments: pool -- the asyncpg pool to use to connect to the database func -- the name of the PostgreSQL function to call arg -- the argument to the PostgreSQL function as either a string or a dict that will be converted into jsonb """ try: if isinstance(arg, str): async with pool.acquire() as conn: q, p = render( f""" SELECT * FROM {func}(:item::text); """, item=arg, ) return await conn.fetchval(q, *p) else: async with pool.acquire() as conn: q, p = render( f""" SELECT * FROM {func}(:item::text::jsonb); """, item=arg.json(exclude_unset=True), ) return await conn.fetchval(q, *p) except exceptions.UniqueViolationError as e: raise ConflictError from e except exceptions.NoDataFoundError as e: raise NotFoundError from e except exceptions.NotNullViolationError as e: raise DatabaseError from e except exceptions.ForeignKeyViolationError as e: raise ForeignKeyError from e
async def test_manual_select(conn): query, params = render('SELECT :v FROM users ORDER BY first_name', v=select_fields('first_name', 'last_name')) v = await conn.fetch(query, *params) assert [ { 'first_name': 'Franks', 'last_name': 'spencer' }, { 'first_name': 'Fred', 'last_name': 'blogs' }, { 'first_name': 'Joe', 'last_name': None }, ] == [dict(r) for r in v]
async def get_anomaly(table_name: str, record_id: int): pool = app.state.pool table = Table(table_name) async with pool.acquire() as conn: record = await conn.fetchrow( str((Query.from_(table).select( table.val, table.dat).where(table.id == record_id)))) target_date = record[1] matrix = record[0] matrix = matrix[1:-1] current_day = list(matrix.split(', ')) query, args = buildpg.render(""" SELECT avg(transponed_arrays.element :: numeric) FROM :table_name, LATERAL ( SELECT val ->> length_series.idx element, length_series.idx idx FROM ( SELECT generate_series(0, jsonb_array_length(val) - 1) ) length_series(idx) ) transponed_arrays WHERE (to_char(dat, 'DD') = to_char(:target::date, 'DD') and to_char(dat, 'MM') = to_char(:target::date, 'MM') ) GROUP BY transponed_arrays.idx ORDER BY transponed_arrays.idx; """, table_name=buildpg.V(table_name), target=target_date) async with pool.acquire() as conn: calculated_values = await conn.fetch(query, *args) avg_values = [rec[0] for rec in calculated_values] anomaly = [int(a) - int(b) for a, b in zip(current_day, avg_values)] bigdict = await get_bigdict_from_matrix(anomaly) return JSONResponse(bigdict)
async def get_average_for_values( *, table_name: str, start_date: date = Form(...), end_date: date = Form(...), ) -> JSONResponse: conn: Connection = app.state.connection query, args = buildpg.render( """ SELECT avg(transponed_arrays.element :: numeric) FROM :table_name, LATERAL ( SELECT val ->> length_series.idx element, length_series.idx idx FROM ( SELECT generate_series(0, jsonb_array_length(val) - 1) ) length_series(idx) ) transponed_arrays WHERE :table_name.dat BETWEEN :start_date AND :end_date GROUP BY transponed_arrays.idx ORDER BY transponed_arrays.idx; """, table_name=buildpg.V(table_name), start_date=start_date, end_date=end_date, ) calculated_values = await conn.fetch(query, *args) values = [rec[0] for rec in calculated_values] bigdict = await get_bigdict_from_matrix(values) return JSONResponse(bigdict)
def test_args(): query, params = render('WHERE :a', a=V('a') == True) # noqa: E712 assert query == 'WHERE a = $1' assert params == [True] assert params[0] is True
def test_falsey_values(value, expected): query, params = render('WHERE :a', a=V('a') == value) assert query == 'WHERE a = $1' assert params == [expected] assert type(params[0]) == type(expected)
def test_where(): query, params = render(':v', v=clauses.Where((V('x') == 4) & (V('y').like('xxx')))) assert 'WHERE x = $1 AND y LIKE $2' == query assert [4, 'xxx'] == params
def test_simple_blocks(block, expected_query, expected_params): query, params = render(':v', v=block()) assert expected_query == query assert expected_params == params
def test_render(template, ctx, expected_query, expected_params): ctx = ctx() query, params = render(template, **ctx) assert expected_query == query assert expected_params == params
def test_errors(query, ctx, msg): with pytest.raises(BuildError) as exc_info: render(query, **ctx) assert msg in str(exc_info.value)
def test_simple_blocks(block, expected_query): query, _ = render(':v', v=block()) assert expected_query == query
async def _search_base(self, search_request: PgstacSearch, **kwargs: Any) -> ItemCollection: """Cross catalog search (POST). Called with `POST /search`. Args: search_request: search request parameters. Returns: ItemCollection containing items which match the search criteria. """ items: Dict[str, Any] request: Request = kwargs["request"] pool = request.app.state.readpool # pool = kwargs["request"].app.state.readpool req = search_request.json(exclude_none=True) try: async with pool.acquire() as conn: q, p = render( """ SELECT * FROM search(:req::text::jsonb); """, req=req, ) items = await conn.fetchval(q, *p) except InvalidDatetimeFormatError: raise InvalidQueryParameter( f"Datetime parameter {search_request.datetime} is invalid.") next: Optional[str] = items.pop("next", None) prev: Optional[str] = items.pop("prev", None) collection = ItemCollection(**items) cleaned_features: List[Item] = [] for feature in collection.get("features") or []: feature = Item(**feature) if (search_request.fields.exclude is None or "links" not in search_request.fields.exclude): # TODO: feature.collection is not always included # This code fails if it's left outside of the fields expression # I've fields extension updated test cases to always include feature.collection feature["links"] = await ItemLinks( collection_id=feature["collection"], item_id=feature["id"], request=request, ).get_links(extra_links=feature.get("links")) exclude = search_request.fields.exclude if exclude and len(exclude) == 0: exclude = None include = search_request.fields.include if include and len(include) == 0: include = None cleaned_features.append(feature) collection["features"] = cleaned_features collection["links"] = await PagingLinks( request=request, next=next, prev=prev, ).get_links() return collection
def test_render(template, var, expected_query, expected_params): query, params = render(template, var=var()) assert expected_query == query assert expected_params == params