philorg/gisaf-backend
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages 6 Wiki Activity Actions

Fix/update baskets:

Browse source 
- get_file
- import
- Basket importers must return BasketImportResult
Add API points
Fix utils delete_df and upsert_df, also making them async friendly
Auth: add helper functions to UserRead
This commit is contained in:
phil 2024-04-09 16:16:04 +05:30
parent 52e1d2135b
commit d2ae5e4d7b
9 changed files with 323 additions and 182 deletions
Download patch file Download diff file Expand all files Collapse all files

41
pdm.lock generated
View file

@ -5,7 +5,20 @@
groups = ["default", "dev", "mqtt"] groups = ["default", "dev", "mqtt"]
strategy = ["cross_platform"] strategy = ["cross_platform"]
lock_version = "4.4.1" lock_version = "4.4.1"
content_hash = "sha256:0da68c7fed8db7a12e36002b8d6194c1651f9653fc7fbf7797c553774c9dbf32" content_hash = "sha256:581ae31055bb26abe5b7bf7cab172337258913d8960892dbd206e00421b309b7"
[[package]]
name = "aiofile"
version = "3.8.8"
requires_python = ">=3.7, <4"
summary = "Asynchronous file operations."
dependencies = [
"caio~=0.9.0",
]
files = [
{file = "aiofile-3.8.8-py3-none-any.whl", hash = "sha256:41e8845cce055779cd77713d949a339deb012eab605b857765e8f8e52a5ed811"},
{file = "aiofile-3.8.8.tar.gz", hash = "sha256:41f3dc40bd730459d58610476e82e5efb2f84ae6e9fa088a9545385d838b8a43"},
]
[[package]] [[package]]
name = "aiomqtt" name = "aiomqtt"
@ -20,6 +33,20 @@ files = [
{file = "aiomqtt-2.0.1.tar.gz", hash = "sha256:60f6451c8ab7235cfb392b1b0cab398e9bc6040f4b140628c0615371abcde15f"}, {file = "aiomqtt-2.0.1.tar.gz", hash = "sha256:60f6451c8ab7235cfb392b1b0cab398e9bc6040f4b140628c0615371abcde15f"},
] ]
[[package]]
name = "aiopath"
version = "0.6.11"
requires_python = ">=3.10"
summary = "📁 Async pathlib for Python"
dependencies = [
"aiofile<4,>=3.5.0",
"anyio<4,>=3.2.0",
]
files = [
{file = "aiopath-0.6.11-py2.py3-none-any.whl", hash = "sha256:7b1f1aa3acb422050908ac3c4755b5e43f625111be003f1bfc7dc2193027c45d"},
{file = "aiopath-0.6.11.tar.gz", hash = "sha256:2f0d4d9195281612c6508cbfa12ac3184c31540d13b9e6215a325897da59decd"},
]
[[package]] [[package]]
name = "aiosqlite" name = "aiosqlite"
version = "0.20.0" version = "0.20.0"
@ -175,6 +202,18 @@ files = [
{file = "bcrypt-4.0.1.tar.gz", hash = "sha256:27d375903ac8261cfe4047f6709d16f7d18d39b1ec92aaf72af989552a650ebd"}, {file = "bcrypt-4.0.1.tar.gz", hash = "sha256:27d375903ac8261cfe4047f6709d16f7d18d39b1ec92aaf72af989552a650ebd"},
] ]
[[package]]
name = "caio"
version = "0.9.13"
requires_python = ">=3.7, <4"
summary = "Asynchronous file IO for Linux MacOS or Windows."
files = [
{file = "caio-0.9.13-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:789f8b55f4a2b46be14361df3ac8d14b6c8f0a3730badd70cb1b7778fcdc7039"},
{file = "caio-0.9.13-cp311-cp311-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a914684bad2a757cf013ae88d785d81659a3add1885bad60cd20bfbd3068bd5a"},
{file = "caio-0.9.13-py3-none-any.whl", hash = "sha256:582cbfc6e203d1dedf662ba972a94db6e744fe0b6bb9e02922b0f86803006fc9"},
{file = "caio-0.9.13.tar.gz", hash = "sha256:26f1e08a442bef4526a66142ea4e325e22dca8f040800aecb3caf8fae0589e98"},
]
[[package]] [[package]]
name = "certifi" name = "certifi"
version = "2023.7.22" version = "2023.7.22"

3
pyproject.toml
View file

@ -18,7 +18,7 @@ dependencies = [
"pydantic-settings>=2.0.3", "pydantic-settings>=2.0.3",
"pyshp>=2.3.1", "pyshp>=2.3.1",
"python-jose[cryptography]>=3.3.0", "python-jose[cryptography]>=3.3.0",
"python-multipart>=0.0.6", "python-multipart>=0.0.9",
"pyyaml>=6.0.1", "pyyaml>=6.0.1",
"redis>=5.0.1", "redis>=5.0.1",
"sqlalchemy[asyncio]>=2.0.23", "sqlalchemy[asyncio]>=2.0.23",
@ -29,6 +29,7 @@ dependencies = [
"psycopg>=3.1.18", "psycopg>=3.1.18",
"plotly>=5.20.0", "plotly>=5.20.0",
"matplotlib>=3.8.3", "matplotlib>=3.8.3",
"aiopath>=0.6.11",
] ]
requires-python = ">=3.11,<4" requires-python = ">=3.11,<4"
readme = "README.md" readme = "README.md"

2
src/gisaf/_version.py
View file

@ -1 +1 @@
__version__: str = '2023.4.dev62+g08c53cf.d20240405' __version__: str = '2023.4.dev63+g52e1d21.d20240408'

60
src/gisaf/api/admin.py
View file

@ -1,9 +1,11 @@
import logging import logging
from fastapi import Depends, APIRouter, HTTPException, status, responses from fastapi import (Depends, APIRouter, HTTPException, status, responses,
UploadFile)
from fastapi.responses import FileResponse
from gisaf.models.admin import AdminBasket, BasketNameOnly from gisaf.models.admin import AdminBasket, BasketImportResult, BasketNameOnly
from gisaf.models.authentication import User from gisaf.models.authentication import User, UserRead
from gisaf.security import get_current_active_user from gisaf.security import get_current_active_user
from gisaf.admin import manager from gisaf.admin import manager
@ -40,3 +42,55 @@ async def get_basket(
## TODO: Fix projects ## TODO: Fix projects
# projects=getattr(basket, 'projects', None) # projects=getattr(basket, 'projects', None)
) )
@api.post('/basket/upload/{name}')
async def upload_basket_file(
name: str,
file: UploadFile,
user: UserRead = Depends(get_current_active_user),
):
try:
basket = manager.baskets[name]
except KeyError:
raise HTTPException(status.HTTP_404_NOT_FOUND, f'No basket named {name}')
fileItem = await basket.add_files(file, user)
return fileItem
@api.get('/basket/download/{name}/{file_id}/{file_name}')
async def download_basket_file(
name: str,
file_id: int,
file_name: str,
user: User = Depends(get_current_active_user),
) -> FileResponse:
try:
basket = manager.baskets[name]
except KeyError:
raise HTTPException(status.HTTP_404_NOT_FOUND, f'No basket named {name}')
if basket.role:
if not user.has_role(basket.role):
raise HTTPException(status.HTTP_401_UNAUTHORIZED)
file_record = await basket.get_file(file_id)
if file_record is None:
raise HTTPException(status.HTTP_404_NOT_FOUND, f'No import file id {file_id}')
return FileResponse(file_record.path)
@api.get('/basket/import/{basket}/{file_id}')
async def import_basket_file(
basket: str,
file_id: int,
dryRun: bool = False,
user: User = Depends(get_current_active_user),
) -> BasketImportResult:
if not (user and user.has_role('reviewer')):
raise HTTPException(status.HTTP_401_UNAUTHORIZED)
basket_ = manager.baskets[basket]
file_import = await basket_.get_file(file_id)
if file_import is None:
raise HTTPException(status.HTTP_404_NOT_FOUND, f'No import file id {file_id}')
try:
result = await basket_.import_file(file_import, dryRun)
## FIXME: shouldn't it be AdminImportError?
except ImportError as err:
raise HTTPException(status.HTTP_500_INTERNAL_SERVER_ERROR, err)
return result

158
src/gisaf/baskets.py
View file

@ -1,4 +1,5 @@
from pathlib import Path from pathlib import Path
from aiopath import AsyncPath
from collections import defaultdict from collections import defaultdict
from json import loads from json import loads
from datetime import datetime from datetime import datetime
@ -10,13 +11,15 @@ from typing import ClassVar
# from aiohttp.web import HTTPUnauthorized, HTTPForbidden # from aiohttp.web import HTTPUnauthorized, HTTPForbidden
from sqlmodel import select from sqlmodel import select
from sqlalchemy.orm import joinedload, QueryableAttribute
from sqlalchemy.exc import NoResultFound
from fastapi import UploadFile
from gisaf.config import conf from gisaf.config import conf
from gisaf.utils import ToMigrate
from gisaf.database import db_session from gisaf.database import db_session
from gisaf.importers import RawSurveyImporter, GeoDataImporter, LineWorkImporter, ImportError from gisaf.importers import RawSurveyImporter, GeoDataImporter, LineWorkImporter, ImportError
from gisaf.models.admin import FileImport, AdminBasketFile, BasketImportResult from gisaf.models.admin import FileImport, AdminBasketFile, BasketImportResult
from gisaf.models.authentication import User from gisaf.models.authentication import UserRead
from gisaf.models.survey import Surveyor, Equipment from gisaf.models.survey import Surveyor, Equipment
from gisaf.models.project import Project from gisaf.models.project import Project
@ -48,7 +51,7 @@ class Basket:
self.importer = self.importer_class() self.importer = self.importer_class()
self.importer.basket = self self.importer.basket = self
async def allowed_for(self, user: User): async def allowed_for(self, user: UserRead):
""" """
Return False if the basket is protected by a role Return False if the basket is protected by a role
Request: aiohttp.Request instance Request: aiohttp.Request instance
@ -101,23 +104,37 @@ class Basket:
# else: # else:
# return df # return df
async def get_file(self, id): async def get_file(self, file_id: int) -> FileImport | None:
df = await FileImport.get_df( async with db_session() as session:
where=FileImport.id==id, query = select(FileImport).where(FileImport.id==file_id).options(
with_related=True, joinedload(FileImport.project),
) joinedload(FileImport.surveyor),
df.rename(columns={ joinedload(FileImport.equipment),
'gisaf_admin_project_name': 'project', )
'gisaf_survey_surveyor_name': 'surveyor', res = await session.exec(query)
'gisaf_survey_equipment_name': 'equipment', try:
}, inplace=True) file = res.one()
df['dir'] = df.dir.fillna('.') except NoResultFound:
## Replace path with Path from pathlib return None
df['path'] = df.apply(lambda fi: self.base_dir/fi['dir']/fi['name'], axis=1) else:
file = df.iloc[0] return file
## Override file.name, which otherwise is the index of the item (hack)
file.name = file['name'] # df = await FileImport.get_df(
return file # where=FileImport.id==id,
# with_related=True,
# )
# df.rename(columns={
# 'gisaf_admin_project_name': 'project',
# 'gisaf_survey_surveyor_name': 'surveyor',
# 'gisaf_survey_equipment_name': 'equipment',
# }, inplace=True)
# df['dir'] = df.dir.fillna('.')
# ## Replace path with Path from pathlib
# df['path'] = df.apply(lambda fi: self.base_dir/fi['dir']/fi['name'], axis=1)
# file = df.iloc[0]
# ## Override file.name, which otherwise is the index of the item (hack)
# file.name = file['name']
# return file
async def delete_file(self, id): async def delete_file(self, id):
file = await FileImport.get(id) file = await FileImport.get(id)
@ -129,7 +146,7 @@ class Basket:
path.unlink() path.unlink()
await file.delete() await file.delete()
async def import_file(self, file_import, dry_run=True, return_data_info=False, **kwargs): async def import_file(self, file_import: FileImport, dry_run=True, **kwargs) -> BasketImportResult:
""" """
Import the file by calling the basket's importer's do_import. Import the file by calling the basket's importer's do_import.
Time stamp the FileImport. Time stamp the FileImport.
@ -139,44 +156,42 @@ class Basket:
return BasketImportResult( return BasketImportResult(
message=f'No import defined/required for {self.name} basket' message=f'No import defined/required for {self.name} basket'
) )
result: BasketImportResult
try: try:
import_result = await self.importer.do_import(file_import, dry_run=dry_run, **kwargs) result = await self.importer.do_import(file_import, dry_run=dry_run, **kwargs)
except ImportError as err: except ImportError as err:
raise raise err
except Exception as err: except Exception as err:
logger.exception(err) logger.exception(err)
raise ImportError(f'Unexpected import error: {err}') raise ImportError(f'Unexpected import error (details in the Gisaf logs): {err}')
if isinstance(import_result, BasketImportResult): if not isinstance(result, BasketImportResult):
result = import_result raise ImportError('Import error: the importer did not return a BasketImportResult')
else: # if import_result:
if import_result: # if isinstance(import_result, (tuple, list)):
if isinstance(import_result, (tuple, list)): # assert len(import_result) >= 2, \
assert len(import_result) >= 2, \ # 'do_import should return message or (message, details)'
'do_import should return message or (message, details)' # result = BasketImportResult(
result = BasketImportResult( # message=import_result[0],
message=import_result[0], # details=import_result[1],
details=import_result[1], # )
) # if len(import_result) > 2:
if len(import_result) > 2: # data = import_result[2]
data = import_result[2] # else:
else: # result = BasketImportResult(message=import_result)
result = BasketImportResult(message=import_result) # else:
else: # result = BasketImportResult(message='Import successful.')
result = BasketImportResult(message='Import successful.')
if dry_run: if dry_run:
result.time = file_import.time result.time = file_import.time
else: if not dry_run:
if not result.time:
result.time = datetime.now()
## Save time stamp ## Save time stamp
await (await FileImport.get(file_import.id)).update(time=result.time).apply() async with db_session() as session:
if return_data_info: file_import.time = result.time
return result, data session.add(file_import)
else: await session.commit()
return result return result
async def add_files(self, reader, request): async def add_files(self, file: UploadFile, user: UserRead):
""" """
File upload to basket. File upload to basket.
Typically called through an http POST view handler. Typically called through an http POST view handler.
@ -184,37 +199,26 @@ class Basket:
Note that the return dict has eventually numpy types and needs NumpyEncoder Note that the return dict has eventually numpy types and needs NumpyEncoder
to be json.dump'ed. to be json.dump'ed.
""" """
raise ToMigrate("basket add_files reader was aiohttp's MultipartReader")
## TODO: multiple items ## TODO: multiple items
## TODO: check if file already exists ## TODO: check if file already exists
## First part is the file # assert part.name == 'file'
part = await reader.next()
assert part.name == 'file'
file_name = part.filename
## Save file on filesystem ## Save file on filesystem
size = 0 path = AsyncPath(self.base_dir) / file.filename
path = Path(self.base_dir) / file_name
## Eventually create the directory ## Eventually create the directory
path.parent.mkdir(parents=True, exist_ok=True) await path.parent.mkdir(parents=True, exist_ok=True)
with path.open('wb') as f: async with path.open('wb') as f:
while True: ## No way to use async to stream the file content to write it?
chunk = await part.read_chunk() # 8192 bytes by default. await f.write(await file.read())
if not chunk:
break
size += len(chunk)
f.write(chunk)
## Read other parts ## Read other parts
parts = defaultdict(None) # parts = defaultdict(None)
while True: # while True:
part = await reader.next() # part = await reader.next()
if not part: # if not part:
break # break
value = (await part.read()).decode() # value = (await part.read()).decode()
if value != 'null': # if value != 'null':
parts[part.name] = value # parts[part.name] = value
## Find ids of project, surveyor, equipment ## Find ids of project, surveyor, equipment
if 'project' in parts: if 'project' in parts:
@ -238,7 +242,7 @@ class Basket:
else: else:
store_type_name = None store_type_name = None
fileImportRecord = await FileImport( fileImportRecord = await FileImport(
name=file_name, name=file.file_name,
dir=parts.get('dir', '.'), dir=parts.get('dir', '.'),
basket=self.name, basket=self.name,
project_id=project_id, project_id=project_id,
@ -251,7 +255,7 @@ class Basket:
admin_basket_file = AdminBasketFile( admin_basket_file = AdminBasketFile(
id=fileImportRecord.id, id=fileImportRecord.id,
name=file_name, name=file.file_name,
status=parts.get('status'), status=parts.get('status'),
store=store_type_name, store=store_type_name,
project=parts.get('project'), project=parts.get('project'),

5
src/gisaf/importers.py
View file

@ -15,8 +15,7 @@ from sqlalchemy.sql.schema import Column
from gisaf.config import conf from gisaf.config import conf
#from .models.admin import FileImport #from .models.admin import FileImport
from gisaf.models.admin import FeatureImportData from gisaf.models.admin import FeatureImportData, BasketImportResult
# from gisaf.models.graphql import BasketImportResult
from gisaf.models.raw_survey import RawSurveyModel from gisaf.models.raw_survey import RawSurveyModel
from gisaf.models.survey import Surveyor, Accuracy, Equipment, AccuracyEquimentSurveyorMapping from gisaf.models.survey import Surveyor, Accuracy, Equipment, AccuracyEquimentSurveyorMapping
from gisaf.models.tags import Tags from gisaf.models.tags import Tags
@ -42,7 +41,7 @@ class Importer:
""" """
basket = None basket = None
async def do_import(self, file_record, dry_run=False, **kwargs): async def do_import(self, file_record, dry_run=False, **kwargs) -> BasketImportResult:
""" """
Return: a BasketImportResult instance, or a message string, Return: a BasketImportResult instance, or a message string,
or a tuple or a list like (message, details for user feedback). or a tuple or a list like (message, details for user feedback).

4
src/gisaf/models/admin.py
View file

@ -157,9 +157,9 @@ class Basket(BasketNameOnly):
class BasketImportResult(BaseModel): class BasketImportResult(BaseModel):
time: datetime time: datetime = Field(default_factory=datetime.now)
message: str message: str
details: str details: dict[str, str | int | float | bool] | None = None
class AdminBasketFile(BaseModel): class AdminBasketFile(BaseModel):
id: int id: int

10
src/gisaf/models/authentication.py
View file

@ -86,6 +86,16 @@ class UserRead(UserBase):
email: str | None # type: ignore email: str | None # type: ignore
roles: list[RoleReadNoUsers] = [] roles: list[RoleReadNoUsers] = []
def can_view(self, model) -> bool:
role = getattr(model, 'viewable_role', None)
if role:
return self.has_role(role)
else:
return True
def has_role(self, role: str) -> bool:
return role in (role.name for role in self.roles)
# class ACL(BaseModel): # class ACL(BaseModel):
# user_id: int # user_id: int

222
src/gisaf/utils.py
View file

@ -11,11 +11,10 @@ from numpy import ndarray
import pandas as pd import pandas as pd
from sqlalchemy.dialects.postgresql import insert from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.sql.expression import delete from sqlmodel import SQLModel, delete
# from graphene import ObjectType
from gisaf.config import conf from gisaf.config import conf
from gisaf.database import db_session
class ToMigrate(Exception): class ToMigrate(Exception):
pass pass
@ -213,107 +212,142 @@ def atimeit(func):
return helper return helper
async def delete_df(df, model): async def delete_df(df: pd.DataFrame, model: SQLModel):
""" """
Delete all data in the model's table in the database Delete all data in the model's table in the database
that matches data in the pandas dataframe. that matches data in the pandas dataframe.
""" """
table = model.__table__ if len(df) == 0:
return
ids = df.reset_index()['id'].values ids = df.reset_index()['id'].values
delete_stmt = delete(table).where(model.id.in_(ids)) statement = delete(model).where(model.id.in_(ids))
async with db.bind.raw_pool.acquire() as conn: async with db_session() as session:
async with conn.transaction(): await session.exec(statement)
await conn.execute(str(delete_stmt), *ids) await session.commit()
async def upsert_df(df, model): # async def upsert_df(df, model):
""" # """
Insert or update all data in the model's table in the database # Insert or update all data in the model's table in the database
that's present in the pandas dataframe. # that's present in the pandas dataframe.
Use postgres insert ... on conflict update... # Use postgres insert ... on conflict update...
with a series of inserts with with one row at a time. # with a series of inserts with with one row at a time.
For GeoDataFrame: the "geometry" column (df._geometry_column_name) is not honnored # For GeoDataFrame: the "geometry" column (df._geometry_column_name) is not honnored
(yet). It's the caller's responsibility to have a proper column name # (yet). It's the caller's responsibility to have a proper column name
(typically "geom" in Gisaf models) with a EWKT or EWKB representation of the geometry. # (typically "geom" in Gisaf models) with a EWKT or EWKB representation of the geometry.
""" # """
## See: https://stackoverflow.com/questions/33307250/postgresql-on-conflict-in-sqlalchemy # ## See: https://stackoverflow.com/questions/33307250/postgresql-on-conflict-in-sqlalchemy
# if len(df) == 0:
# return df
# table = model.__table__
# ## Generate the 'upsert' statement, using fake values but defining columns
# columns = {c.name for c in table.columns}
# values = {col: None for col in df.columns if col in columns}
# insrt_stmnt = insert(table).inline().values(values).returning(table.primary_key.columns)
# df_columns = set(df.columns)
# do_update_stmt = insrt_stmnt.on_conflict_do_update(
# constraint=table.primary_key,
# set_={
# k.name: getattr(insrt_stmnt.excluded, k.name)
# for k in insrt_stmnt.excluded
# if k.name in df_columns and
# k.name not in [c.name for c in table.primary_key.columns]
# }
# )
# ## Filter and reorder the df columns
# ## in order to match the order of columns in the insert statement
# df = df[[col for col in do_update_stmt.compile().positiontup
# if col in df_columns]].copy()
# def convert_to_object(value):
# """
# Quick (but slow) and dirty: clean up values (nan, nat) for inserting to postgres via asyncpg
# """
# if isinstance(value, float) and isnan(value):
# return None
# elif pd.isna(value):
# return None
# else:
# return value
# # def encode_geometry(geometry):
# # if not hasattr(geometry, '__geo_interface__'):
# # raise TypeError('{g} does not conform to '
# # 'the geo interface'.format(g=geometry))
# # shape = shapely.geometry.asShape(geometry)
# # return shapely.wkb.dumps(shape)
# # def decode_geometry(wkb):
# # return shapely.wkb.loads(wkb)
# ## pks: list of dicts of primary keys
# pks = {pk.name: [] for pk in table.primary_key.columns}
# async with db.bind.raw_pool.acquire() as conn:
# ## Set standard encoder for HSTORE, geometry
# await conn.set_builtin_type_codec('hstore', codec_name='pg_contrib.hstore')
# #await conn.set_type_codec(
# # 'geometry', # also works for 'geography'
# # encoder=encode_geometry,
# # decoder=decode_geometry,
# # format='binary',
# #)
# #await conn.set_type_codec(
# # 'json',
# # encoder=json.dumps,
# # decoder=json.loads,
# # schema='pg_catalog'
# #)
# ## For a sequence of inserts:
# insrt_stmnt_single = await conn.prepare(str(do_update_stmt))
# async with conn.transaction():
# for row in df.itertuples(index=False):
# converted_row = [convert_to_object(v) for v in row]
# returned = await insrt_stmnt_single.fetch(*converted_row)
# for returned_single in returned:
# for pk, value in returned_single.items():
# pks[pk].append(value)
# ## Return a copy of the original df, with actual DB columns, data and the primary keys
# for pk, values in pks.items():
# df[pk] = values
# return df
def postgres_upsert(table, conn, keys, data_iter):
# See https://stackoverflow.com/questions/61366664/how-to-upsert-pandas-dataframe-to-postgresql-table
# Comment by @HopefullyThisHelps
data = [dict(zip(keys, row)) for row in data_iter]
insert_statement = insert(table.table).values(data)
upsert_statement = insert_statement.on_conflict_do_update(
constraint=f"{table.table.name}_pkey",
set_={c.key: c for c in insert_statement.excluded},
)
conn.execute(upsert_statement)
async def upsert_df(df: pd.DataFrame, model: SQLModel, chunksize: int = 1000):
if len(df) == 0: if len(df) == 0:
return df return df
from functools import partial
table = model.__table__ import concurrent.futures
import asyncio
## Generate the 'upsert' statement, using fake values but defining columns loop = asyncio.get_running_loop()
columns = {c.name for c in table.columns} with concurrent.futures.ProcessPoolExecutor() as pool:
values = {col: None for col in df.columns if col in columns} await loop.run_in_executor(
insrt_stmnt = insert(table, inline=True, values=values, returning=table.primary_key.columns) pool,
df_columns = set(df.columns) partial(
do_update_stmt = insrt_stmnt.on_conflict_do_update( df.to_sql,
constraint=table.primary_key, model.__tablename__,
set_={ conf.db.get_pg_url(), # Cannot use sync_engine in run_in_executor
k.name: getattr(insrt_stmnt.excluded, k.name) # because it's not pickable
for k in insrt_stmnt.excluded schema=model.__table__.schema, # type: ignore
if k.name in df_columns and if_exists="append",
k.name not in [c.name for c in table.primary_key.columns] index=False,
} method=postgres_upsert,
) chunksize=chunksize,
## Filter and reorder the df columns ),
## in order to match the order of columns in the insert statement )
df = df[[col for col in do_update_stmt.compile().positiontup
if col in df_columns]].copy()
def convert_to_object(value):
"""
Quick (but slow) and dirty: clean up values (nan, nat) for inserting to postgres via asyncpg
"""
if isinstance(value, float) and isnan(value):
return None
elif pd.isna(value):
return None
else:
return value
# def encode_geometry(geometry):
# if not hasattr(geometry, '__geo_interface__'):
# raise TypeError('{g} does not conform to '
# 'the geo interface'.format(g=geometry))
# shape = shapely.geometry.asShape(geometry)
# return shapely.wkb.dumps(shape)
# def decode_geometry(wkb):
# return shapely.wkb.loads(wkb)
## pks: list of dicts of primary keys
pks = {pk.name: [] for pk in table.primary_key.columns}
async with db.bind.raw_pool.acquire() as conn:
## Set standard encoder for HSTORE, geometry
await conn.set_builtin_type_codec('hstore', codec_name='pg_contrib.hstore')
#await conn.set_type_codec(
# 'geometry', # also works for 'geography'
# encoder=encode_geometry,
# decoder=decode_geometry,
# format='binary',
#)
#await conn.set_type_codec(
# 'json',
# encoder=json.dumps,
# decoder=json.loads,
# schema='pg_catalog'
#)
## For a sequence of inserts:
insrt_stmnt_single = await conn.prepare(str(do_update_stmt))
async with conn.transaction():
for row in df.itertuples(index=False):
converted_row = [convert_to_object(v) for v in row]
returned = await insrt_stmnt_single.fetch(*converted_row)
for returned_single in returned:
for pk, value in returned_single.items():
pks[pk].append(value)
## Return a copy of the original df, with actual DB columns, data and the primary keys
for pk, values in pks.items():
df[pk] = values
return df
#async def upsert_df(df, model): #async def upsert_df(df, model):