tracktolib 0.35.0__tar.gz → 0.37.0__tar.gz

{tracktolib-0.35.0 → tracktolib-0.37.0}/PKG-INFO

@@ -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 → tracktolib-0.37.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tracktolib"
-version = "0.35.0"
+version = "0.37.0"
 description = "Utility library for python"
 authors = ["Julien Brayere <julien.brayere@tracktor.fr>"]
 license = "MIT"
@@ -70,7 +70,7 @@ pythonPlatform = "Linux"
 
 [tool.commitizen]
 name = "cz_conventional_commits"
-version = "0.35.0"
+version = "0.37.0"
 tag_format = "$version"
 version_files = [
     "pyproject.toml:version"

{tracktolib-0.35.0 → tracktolib-0.37.0}/setup.py

@@ -19,7 +19,7 @@ extras_require = \
 
 setup_kwargs = {
     'name': 'tracktolib',
-    'version': '0.35.0',
+    'version': '0.37.0',
     'description': 'Utility library for python',
     'long_description': "# Tracktolib\n\n[![Python versions](https://img.shields.io/pypi/pyversions/tracktolib)](https://pypi.python.org/pypi/tracktolib)\n[![Latest PyPI version](https://img.shields.io/pypi/v/tracktolib?logo=pypi)](https://pypi.python.org/pypi/tracktolib)\n[![CircleCI](https://circleci.com/gh/Tracktor/tracktolib/tree/master.svg?style=shield)](https://app.circleci.com/pipelines/github/Tracktor/tracktolib?branch=master)\n\nUtility library for python\n\n# Installation\n\nYou can choose to not install all the dependencies by specifying\nthe [extra](https://python-poetry.org/docs/cli/#options-4) parameter such as:\n\n```bash\npoetry add tracktolib@latest -E pg-sync -E tests --group dev \n```\n\nHere we only install the utilities using `psycopg` (pg-sync) and `deepdiff` (tests) for the dev environment.\n\n# Utilities\n\n- **log**\n\nUtility functions for logging.\n\n```python\nimport logging\nfrom tracktolib.logs import init_logging\n\nlogger = logging.getLogger()\nformatter, stream_handler = init_logging(logger, 'json', version='0.0.1')\n```\n\n- **pg**\n\nUtility functions for [asyncpg](https://github.com/MagicStack/asyncpg)\n\n- **pg-sync**\n\nUtility functions based on psycopg such as `fetch_one`, `insert_many`, `fetch_count` ...\n\nTo use the functions, create a `Connection` using psycopg: `conn = psycopg2.connect()`\n\n*fetch_one*\n\n```python\nfrom pg.pg_sync import (\n    insert_many, fetch_one, fetch_count, fetch_all\n)\n\ndata = [\n    {'foo': 'bar', 'value': 1},\n    {'foo': 'baz', 'value': 2}\n]\ninsert_many(conn, 'public.test', data)  # Will insert the 2 dict\nquery = 'SELECT foo from public.test order by value asc'\nvalue = fetch_one(conn, query, required=True)  # Will return {'foo': 'bar'}, raise an error is not found\nassert fetch_count(conn, 'public.test') == 2\nquery = 'SELECT * from public.test order by value asc'\nassert fetch_all(conn, query) == data\n\n```\n\n- **tests**\n\nUtility functions for testing\n\n- **s3-minio**\n\nUtility functions for [minio](https://min.io/docs/minio/linux/developers/python/API.html)\n\n- **s3**\n\nUtility functions for [aiobotocore](https://github.com/aio-libs/aiobotocore)\n\n- **logs**\n\nUtility functions to initialize the logging formatting and streams\n\n- **http**\n\nUtility functions using [httpx](https://www.python-httpx.org/)\n\n- **api**\n\nUtility functions using [fastapi](https://fastapi.tiangolo.com/)\n",
     'author': 'Julien Brayere',

{tracktolib-0.35.0 → tracktolib-0.37.0}/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-0.35.0 → tracktolib-0.37.0}/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-0.37.0/tracktolib/pg_utils.py

@@ -0,0 +1,58 @@
+from typing_extensions import LiteralString
+from typing import Iterable
+from typing import cast
+
+
+def get_tmp_table_query(schema: 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;
+    """
+
+    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/tracktolib/pg_utils.py

@@ -1,18 +0,0 @@
-from typing_extensions import LiteralString
-
-
-def get_tmp_table_query(schema: LiteralString,
-                        table: LiteralString):
-    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;
-    """
-    return tmp_table_name, create_tmp_table_query, insert_query