tracktolib 0.49.1__py3-none-any.whl → 0.51.0__py3-none-any.whl

tracktolib/pg/__init__.py

@@ -11,5 +11,6 @@ from .query import (
     update_returning,
     update_one,
     insert_pg,
+    OnConflict,
 )
 from .utils import iterate_pg, upsert_csv, safe_pg, safe_pg_context, PGError, PGException

tracktolib/pg/query.py

@@ -1,6 +1,7 @@
 import typing
 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:
@@ -32,6 +33,7 @@ def _get_on_conflict_query(
     constraint: K | None,
     on_conflict: K | None,
     where: K | None,
+    merge_columns: Iterable[K] | None,
 ) -> str:
     _on_conflict = get_conflict_query(
         columns=columns,
@@ -40,6 +42,7 @@ def _get_on_conflict_query(
         constraint=constraint,
         on_conflict=on_conflict,
         where=where,
+        merge_columns=merge_columns,
     )
     return f"{query} {_on_conflict}"
 
@@ -72,6 +75,8 @@ class PGConflictQuery(Generic[K]):
     query: str | None = None
     constraint: str | None = None
     where: str | None = None
+    """JSONB keys to merge (like jsonb1 || newjsonb2)"""
+    merge_keys: Iterable[K] | None = None
 
     def __post_init__(self):
         _has_keys = 1 if (self.keys or self.ignore_keys) else 0
@@ -172,6 +177,7 @@ class PGInsertQuery(PGQuery):
                 self.on_conflict.constraint,
                 self.on_conflict.query,
                 self.on_conflict.where,
+                self.on_conflict.merge_keys,
             )
 
         # Returning
@@ -193,17 +199,30 @@ def get_update_fields(
     start_from: int = 0,
     ignore_keys: list[str] | None = None,
     quote_columns: bool = False,
+    merge_keys: list[str] | None = None,
 ) -> tuple[str, list]:
     values, fields, where_values = [], [], []
     counter = 0
+    _merge_keys = set(merge_keys or [])
+    _ignore_keys = ignore_keys or []
+
+    _invalid_merge_keys = _merge_keys - set(keys)
+    if _invalid_merge_keys:
+        raise ValueError(f"Merge keys not in keys found: {_invalid_merge_keys}")
+
     for k in keys:
         v = item[k]
-        if ignore_keys and k in ignore_keys:
+        if k in _ignore_keys:
             where_values.append(v)
             continue
         values.append(v)
         _col = f'"{k}"' if quote_columns else k
-        fields.append(f"{_col} = ${counter + start_from + 1}")
+        _counter = counter + start_from + 1
+        fields.append(
+            f"{_col} = ${_counter}"
+            if k not in _merge_keys
+            else f"{_col} = COALESCE(t.{_col}, jsonb_build_object()) || " f"${_counter}"
+        )
         counter += 1
     return ",\n".join(fields), values + where_values
 
@@ -222,6 +241,8 @@ class PGUpdateQuery(PGQuery):
     returning: str | list[str] | None = None
     """If True, the query will return all the updated fields"""
     return_keys: bool = False
+    """Values to update using merge (like {}::jsonb || {}::jsonb)"""
+    merge_keys: list[str] | None = None
 
     _update_fields: str | None = field(init=False, default=None)
     _values: list | None = field(init=False, default=None)
@@ -238,6 +259,7 @@ class PGUpdateQuery(PGQuery):
             start_from=self.start_from or 0,
             ignore_keys=self.where_keys,
             quote_columns=self.quote_columns,
+            merge_keys=self.merge_keys,
         )
         if self.returning and self.return_keys:
             raise ValueError("Please choose either returning or return_keys")
@@ -263,7 +285,7 @@ class PGUpdateQuery(PGQuery):
             raise ValueError("No update fields found")
 
         query = f"""
-        UPDATE {self.table}
+        UPDATE {self.table} t
             SET {self._update_fields}
         {self._get_where_query()}
         """
@@ -370,8 +392,13 @@ async def update_one(
     keys: list[str] | None = None,
     start_from: int | None = None,
     where: str | None = None,
+    merge_keys: list[str] | None = None,
 ):
-    query = PGUpdateQuery(table=table, items=[item], start_from=start_from, where_keys=keys, where=where)
+    query = PGUpdateQuery(
+        table=table, items=[item], start_from=start_from, where_keys=keys, where=where, merge_keys=merge_keys
+    )
+    print(query.query)
+    print(query.values)
     await conn.execute(query.query, *args, *query.values)
 
 
@@ -386,6 +413,7 @@ async def update_returning(
     where: str | None = None,
     keys: list[str] | None = None,
     start_from: int | None = None,
+    merge_keys: list[str] | None = None,
 ) -> Any | None: ...
 
 
@@ -400,6 +428,7 @@ async def update_returning(
     where: str | None = None,
     keys: list[str] | None = None,
     start_from: int | None = None,
+    merge_keys: list[str] | None = None,
 ) -> asyncpg.Record | None: ...
 
 
@@ -414,6 +443,7 @@ async def update_returning(
     where: str | None = None,
     keys: list[str] | None = None,
     start_from: int | None = None,
+    merge_keys: list[str] | None = None,
 ) -> asyncpg.Record | None: ...
 
 
@@ -427,6 +457,7 @@ async def update_returning(
     where: str | None = None,
     keys: list[str] | None = None,
     start_from: int | None = None,
+    merge_keys: list[str] | None = None,
 ) -> Any | asyncpg.Record | None:
     if returning is not None:
         returning_values = [returning] if isinstance(returning, str) else returning
@@ -440,6 +471,7 @@ async def update_returning(
         where_keys=keys,
         return_keys=return_keys,
         returning=returning_values,
+        merge_keys=merge_keys,
     )
     fn = conn.fetchval if len(returning_values or []) == 1 else conn.fetchrow
     return await fn(query.query, *args, *query.values)

tracktolib/pg_utils.py

@@ -1,7 +1,8 @@
-from typing_extensions import LiteralString
 from typing import Iterable
 from typing import cast
 
+from typing_extensions import LiteralString
+
 
 def get_tmp_table_query(
     schema: LiteralString,
@@ -41,6 +42,7 @@ def get_conflict_query(
     constraint: str | None = None,
     on_conflict: str | None = None,
     where: str | None = None,
+    merge_columns: Iterable[str] | None = None,
 ) -> LiteralString:
     if on_conflict:
         return cast(LiteralString, on_conflict)
@@ -55,8 +57,20 @@ def get_conflict_query(
     else:
         raise NotImplementedError("update_keys or constraint must be set")
 
-    _ignore_columns = [*(update_columns or []), *(ignore_columns or [])]
+    _update_columns = update_columns or []
+    _ignore_columns = ignore_columns or []
+    _merge_columns = merge_columns or []
+
+    if set(_merge_columns) & set(_update_columns):
+        raise ValueError("Duplicate keys found between merge and update")
+    if set(_merge_columns) & set(_ignore_columns):
+        raise ValueError("Merge column cannot be ignored")
+
+    _ignore_columns = [*_update_columns, *_ignore_columns, *_merge_columns]
     fields = ", ".join(f"{x} = COALESCE(EXCLUDED.{x}, t.{x})" for x in columns if x not in _ignore_columns)
+    if merge_columns:
+        fields = fields + ", " if fields else fields
+        fields += ", ".join(f"{x} = COALESCE(t.{x}, jsonb_build_object()) || EXCLUDED.{x}" for x in merge_columns)
     if not fields:
         raise ValueError("No fields set")
 

{tracktolib-0.49.1.dist-info → tracktolib-0.51.0.dist-info}/METADATA

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

{tracktolib-0.49.1.dist-info → tracktolib-0.51.0.dist-info}/RECORD

@@ -3,17 +3,17 @@ tracktolib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tracktolib/api.py,sha256=xArVgRj_g7bw2tEDbjC9qs9f55b9X0-kJoc8-s1rsYo,9616
 tracktolib/http_utils.py,sha256=c10JGmHaBw3VSDMYhz2dvVw2lo4PUAq1xMub74I7xDc,2625
 tracktolib/logs.py,sha256=M5RZ8OYKgLEBJeC1AwUCSRbFAH09hwdTjjs-AQ9QGg8,2204
-tracktolib/pg/__init__.py,sha256=32kKBJzIDyppszZMAOSMgM3ggz7Ow0TaNKiMs740gyw,333
-tracktolib/pg/query.py,sha256=KESY4g0XePUrRIakygmzhPP2vIOfJGWc2PAZsIj_PfI,13580
+tracktolib/pg/__init__.py,sha256=j67e3B3gBbCHLD20QBybptmNdbbVMzNhZE6XjIPuKVo,349
+tracktolib/pg/query.py,sha256=XX5eD02TRWRIiBVj0ENooQ893aycan2fJ55_1J207Rw,14711
 tracktolib/pg/utils.py,sha256=cL24KEt4SWJQ7LJPzaO3c8Xg0ZLmjhn22DtTWg86nwc,6324
 tracktolib/pg_sync.py,sha256=z6EB0hELq6t129gnVnX0rNXv2McZXxI8khy-TEG65Ag,5165
-tracktolib/pg_utils.py,sha256=8DDlZVSgfNie9M-zL7wilFufXMMvSPw5d7aIDs_pQCM,1976
+tracktolib/pg_utils.py,sha256=VXPpy1jGq6aCgTlfFJDIrq6JDujR83JN5ZRiCi8Lx4E,2582
 tracktolib/s3/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tracktolib/s3/minio.py,sha256=wMEjkSes9Fp39fD17IctALpD6zB2xwDRQEmO7Vzan3g,1387
 tracktolib/s3/s3.py,sha256=d0Q63Zb62ef4jAt05zQwpgYvAoxHC9kSQuHzzK90VvE,4825
 tracktolib/tests.py,sha256=Pbc4yGQrIWFLBEgX-kDBxhMCqH-bfOJtkH2agvaM6ZQ,382
 tracktolib/utils.py,sha256=jwLww8bqDu8zEip9uN4yW0lE5_YMWfrAHYtagr8sYOA,5295
-tracktolib-0.49.1.dist-info/LICENSE,sha256=uUanH0X7SeZEPdsRTHegMSMTiIHMurt9H0jSwEwKE1Y,1081
-tracktolib-0.49.1.dist-info/METADATA,sha256=_HBMlEoVz25S6viwoubRgSxNhZjG51Madd4d3A0F-XA,3641
-tracktolib-0.49.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-tracktolib-0.49.1.dist-info/RECORD,,
+tracktolib-0.51.0.dist-info/LICENSE,sha256=uUanH0X7SeZEPdsRTHegMSMTiIHMurt9H0jSwEwKE1Y,1081
+tracktolib-0.51.0.dist-info/METADATA,sha256=v2MsoN7EIuclo58jhPp40YJespn9Oy4rYKM8kg7VRsQ,3641
+tracktolib-0.51.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+tracktolib-0.51.0.dist-info/RECORD,,