tracktolib 0.35.0__py3-none-any.whl → 0.37.0__py3-none-any.whl

tracktolib/pg/query.py

@@ -3,6 +3,7 @@ from dataclasses import dataclass, field
 from typing import (
     TypeVar, Iterable, Callable, Generic, Iterator, TypeAlias,
     overload, Any, Literal)
+from ..pg_utils import get_conflict_query
 
 try:
     import asyncpg
@@ -29,27 +30,11 @@ def _get_on_conflict_query(query: str,
                            keys: Iterable[K],
                            update_keys: Iterable[K] | None,
                            ignore_keys: Iterable[K] | None,
-                           constraint: str | None,
-                           on_conflict: str | None) -> str:
-    if on_conflict:
-        return f'{query} {on_conflict}'
-
-    if constraint:
-        query = f'{query} ON CONFLICT ON CONSTRAINT {constraint}'
-    elif update_keys:
-        update_keys_str = ', '.join(sorted(update_keys))
-        query = f'{query} ON CONFLICT ({update_keys_str})'
-    else:
-        raise NotImplementedError('update_keys or constraint must be set')
-
-    _ignore_keys = [*(update_keys or []), *(ignore_keys or [])]
-    fields = ', '.join(f'{x} = COALESCE(EXCLUDED.{x}, t.{x})'
-                       for x in keys
-                       if x not in _ignore_keys)
-    if not fields:
-        raise ValueError('No fields set')
-
-    return f'{query} DO UPDATE SET {fields}'
+                           constraint: K | None,
+                           on_conflict: K | None) -> str:
+    _on_conflict = get_conflict_query(keys=keys, update_keys=update_keys, ignore_keys=ignore_keys,
+                                      constraint=constraint, on_conflict=on_conflict)
+    return f'{query} {_on_conflict}'
 
 
 ReturningFn = Callable[[Iterable[K] | None, K | None], None]

tracktolib/pg/utils.py

@@ -3,8 +3,9 @@ import datetime as dt
 import functools
 import logging
 from pathlib import Path
-from typing import AsyncIterator
+from typing import AsyncIterator, Iterable, Sequence
 from typing_extensions import LiteralString
+from ..pg_utils import get_conflict_query
 
 try:
     import asyncpg
@@ -89,14 +90,27 @@ async def upsert_csv(conn: asyncpg.Connection,
                      *,
                      chunk_size: int = 5_000,
                      show_progress: bool = False,
-                     nb_lines: int | None = None):
+                     nb_lines: int | None = None,
+                     on_conflict_keys: Iterable[LiteralString] | None = None,
+                     delimiter: str = ',',
+                     col_names: Sequence[str] | None = None,
+                     skip_header: bool = False):
     infos = await get_table_infos(conn, schema, table)
 
+    on_conflict_str = 'ON CONFLICT DO NOTHING'
+    if on_conflict_keys is not None:
+        on_conflict_str = get_conflict_query(keys=infos.keys(),
+                                             update_keys=on_conflict_keys)
+
     with csv_path.open('r') as f:
-        reader = csv.DictReader(f)
-        _columns = [x.lower() for x in (reader.fieldnames or [])]
+        reader = csv.DictReader(f, delimiter=delimiter, fieldnames=col_names)
+        if skip_header:
+            next(reader)
+        _columns = col_names if col_names else [x.lower() for x in (reader.fieldnames or [])]
         async with conn.transaction():
-            _tmp_table, _tmp_query, _insert_query = get_tmp_table_query(schema, table)
+            _tmp_table, _tmp_query, _insert_query = get_tmp_table_query(schema, table,
+                                                                        columns=infos.keys(),
+                                                                        on_conflict=on_conflict_str)
             logger.info(f'Creating tmp table: {_tmp_table!r}')
             await conn.execute(_tmp_query)
             logger.info(f'Inserting data from {csv_path!r} to {_tmp_table!r}')

tracktolib/pg_utils.py

@@ -1,18 +1,58 @@
 from typing_extensions import LiteralString
+from typing import Iterable
+from typing import cast
 
 
 def get_tmp_table_query(schema: LiteralString,
-                        table: LiteralString):
+                        table: LiteralString,
+                        columns: Iterable[LiteralString] | None = None,
+                        on_conflict: LiteralString = 'ON CONFLICT DO NOTHING'):
     tmp_table_name = f'{schema}_{table}_tmp'
     create_tmp_table_query = f"""
     CREATE TEMP TABLE {tmp_table_name}
     (LIKE {schema}.{table} INCLUDING DEFAULTS)
     ON COMMIT DROP;
     """
-    insert_query = f"""
-        INSERT INTO {schema}.{table}
-        SELECT *
-        FROM {tmp_table_name}
-        ON CONFLICT DO NOTHING;
-    """
+
+    if columns:
+        _columns = ','.join(columns)
+        insert_query = f"""
+            INSERT INTO {schema}.{table} as t({_columns})
+            SELECT *
+            FROM {tmp_table_name}
+            {on_conflict};
+        """
+    else:
+        insert_query = f"""
+            INSERT INTO {schema}.{table}
+            SELECT *
+            FROM {tmp_table_name}
+            {on_conflict};
+        """
     return tmp_table_name, create_tmp_table_query, insert_query
+
+
+def get_conflict_query(keys: Iterable[str],
+                       update_keys: Iterable[str] | None = None,
+                       ignore_keys: Iterable[str] | None = None,
+                       constraint: str | None = None,
+                       on_conflict: str | None = None) -> LiteralString:
+    if on_conflict:
+        return cast(LiteralString, on_conflict)
+
+    if constraint:
+        query = f'ON CONFLICT ON CONSTRAINT {constraint}'
+    elif update_keys:
+        update_keys_str = ', '.join(sorted(update_keys))
+        query = f'ON CONFLICT ({update_keys_str})'
+    else:
+        raise NotImplementedError('update_keys or constraint must be set')
+
+    _ignore_keys = [*(update_keys or []), *(ignore_keys or [])]
+    fields = ', '.join(f'{x} = COALESCE(EXCLUDED.{x}, t.{x})'
+                       for x in keys
+                       if x not in _ignore_keys)
+    if not fields:
+        raise ValueError('No fields set')
+
+    return cast(LiteralString, f'{query} DO UPDATE SET {fields}')

{tracktolib-0.35.0.dist-info → tracktolib-0.37.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tracktolib
-Version: 0.35.0
+Version: 0.37.0
 Summary: Utility library for python
 Home-page: https://github.com/tracktor/tracktolib
 License: MIT

{tracktolib-0.35.0.dist-info → tracktolib-0.37.0.dist-info}/RECORD

@@ -4,16 +4,16 @@ tracktolib/api.py,sha256=2hAtq7Cjx7GqiqO2ATJfiCK3rM66kDfLKMBVZ_I8ttU,6743
 tracktolib/http.py,sha256=GHWvD9rHTF1oAnG6Ptna_oeMDcwCxYYpPxLOME-oZ3Y,1574
 tracktolib/logs.py,sha256=nbI9LO2NaXDAv9aoUkad54V_x8B8zga-MwqUfYb-pkk,2523
 tracktolib/pg/__init__.py,sha256=mK_EAUN9wTLQC38n3iaBFTP2YzaEo7S17AyFiKHgu6k,241
-tracktolib/pg/query.py,sha256=6IfRqID_RIqnhW4cyZKsZNMA5t9QtXWay3Hllu3xolw,15398
-tracktolib/pg/utils.py,sha256=A7YVFKRwolxTP7S4T6XtR2vyj4oq3jwOITBJZtdDEys,4228
+tracktolib/pg/query.py,sha256=E42F2pq_kv465XO2HQo52BhKhORsZ-mdvsGsqKdnbVg,14969
+tracktolib/pg/utils.py,sha256=Yfd-6G5KfHbwIReYrymzdozFs8flQaP-R9K5CP42t_4,5059
 tracktolib/pg_sync.py,sha256=M_9DMVXWozhxD8KGVWKxLblOG4K-npbK5SKoO-Rz0yM,5203
-tracktolib/pg_utils.py,sha256=rWYCkuOk_Qq2ShNvbudEvcgFNKTu2vJd6PsUmElVu9Q,545
+tracktolib/pg_utils.py,sha256=e5QdJ5clFtap949U57sSjplaWi6LNv-Rs2cuz4NRIxk,2012
 tracktolib/s3/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tracktolib/s3/minio.py,sha256=gkpI2bvVwfXaAFY7M9e0Mlv82x0PjuriS8LdRmEJtMo,1258
 tracktolib/s3/s3.py,sha256=Kfmz0uLl5yGQxH4XvAfIe1yhOieTXS9Ck85nhj_aMnQ,4011
 tracktolib/tests.py,sha256=GycCVRUOdzazxqBXKBWWzhrlkV8yyD7NaActz_A2WXU,590
 tracktolib/utils.py,sha256=kU35LCJaFDNg0IR9ekAi4uEJ3unZCC3VPzqgkctt06Q,5437
-tracktolib-0.35.0.dist-info/LICENSE,sha256=uUanH0X7SeZEPdsRTHegMSMTiIHMurt9H0jSwEwKE1Y,1081
-tracktolib-0.35.0.dist-info/WHEEL,sha256=vxFmldFsRN_Hx10GDvsdv1wroKq8r5Lzvjp6GZ4OO8c,88
-tracktolib-0.35.0.dist-info/METADATA,sha256=6nszCklyCJfTvWH4_HYIVOkIqQL4nVwiSsQPtVgbGZk,3630
-tracktolib-0.35.0.dist-info/RECORD,,
+tracktolib-0.37.0.dist-info/LICENSE,sha256=uUanH0X7SeZEPdsRTHegMSMTiIHMurt9H0jSwEwKE1Y,1081
+tracktolib-0.37.0.dist-info/WHEEL,sha256=vxFmldFsRN_Hx10GDvsdv1wroKq8r5Lzvjp6GZ4OO8c,88
+tracktolib-0.37.0.dist-info/METADATA,sha256=Fu_O-M1gdC3VCihdzJN3PVTRx9VjkO5TacWTLfNQC4c,3630
+tracktolib-0.37.0.dist-info/RECORD,,