Merge branch 'dev'

Author: boonhapus

cs_tools/__project__.py

@@ -1,4 +1,4 @@
-__version__ = "1.5.11"
+__version__ = "1.5.12"
 __docs__ = "https://thoughtspot.github.io/cs_tools/"
 __repo__ = "https://github.com/thoughtspot/cs_tools"
 __help__ = f"{__repo__}/discussions/"

cs_tools/sync/sqlite/syncer.py

@@ -5,6 +5,7 @@
 import logging
 import pathlib
 
+from sqlalchemy.dialects.sqlite import insert
 import pydantic
 import sqlalchemy as sa
 
@@ -40,6 +41,21 @@ def __init__(self, **kwargs):
     def __repr__(self):
         return f"<SQLiteSyncer conn_string='{self.engine.url}'>"
 
+    def insert_on_conflict(self, data: TableRows, *, table: sa.Table) -> Union[sa.Insert, sa.Update]:
+        """UPSERT."""
+        stmt = insert(table).values(data)
+
+        if table.columns == table.primary_key:
+            set_ = {c.key: getattr(stmt.excluded, c.key) for c in table.columns}
+        else:
+            set_ = {c.key: getattr(stmt.excluded, c.key) for c in table.columns if c.key not in table.primary_key}
+
+        stmt = stmt.on_conflict_do_update(
+            index_elements=table.primary_key,
+            set_=set_,
+        )
+        return stmt
+
     # @contextlib.contextmanager
     # def pragma_speedy_insert(self):
     #     """ """
@@ -91,6 +107,12 @@ def dump(self, tablename: str, *, data: TableRows) -> None:
             )
 
         if self.load_strategy == "UPSERT":
-            sync_utils.generic_upsert(table, session=self.session, data=data)
+            sync_utils.batched(
+                self.insert_on_conflict,
+                session=self.session,
+                data=data,
+                max_parameters=const.SQLITE_MAX_VARIABLES,
+                table=table,
+            )
 
         self.session.commit()

cs_tools/sync/utils.py

@@ -50,7 +50,7 @@ def format_datetime_values(row: dict[str, Any], *, dt_format: str = DATETIME_FOR
     return out
 
 
-def batched(prepared_statement, *, session: sa.orm.Session, data: TableRows, max_parameters: int = 999) -> None:
+def batched(prepared_statement, *, session: sa.orm.Session, data: TableRows, max_parameters: int = 999, **kw) -> None:
     """Split data across multiple transactions."""
     batchsize = min(5000, max_parameters // len(data[0]))
     rows = []
@@ -60,14 +60,14 @@ def batched(prepared_statement, *, session: sa.orm.Session, data: TableRows, max
 
         # Commit every so often.
         if row_number % batchsize == 0:
-            stmt = prepared_statement(rows)
+            stmt = prepared_statement(rows, **kw)
             session.execute(stmt)
             session.commit()
             rows = []
 
     # Final commit, grab the rest of the data rows.
     if rows:
-        stmt = prepared_statement(rows)
+        stmt = prepared_statement(rows, **kw)
         session.execute(stmt)
         session.commit()
 

cs_tools/updater/_updater.py

@@ -195,7 +195,7 @@ def pip(self, command: str, *args, with_system_python: bool = False, **kwargs) -
             # don't ping pypi for new versions of pip -- it doesn't matter and is noisy
             "--disable-pip-version-check",
             # trust installs from the official python package index and the thoughtspot github repos
-            "--trusted-host", "files.pythonhost.org",
+            "--trusted-host", "files.pythonhosted.org",
             "--trusted-host", "pypi.org",
             "--trusted-host", "pypi.python.org",
             "--trusted-host", "github.com",

pyproject.toml

@@ -39,7 +39,6 @@ dependencies = [
     "thoughtspot_tml",
     "awesomeversion",
     "httpx >= 0.27.0",
-    "pydantic >= 2.6.4",
     "pydantic-settings",
     "email-validator",
     "pendulum >= 3.0.0",
@@ -48,6 +47,9 @@ dependencies = [
     "toml",
     "packaging",
 
+    # TODO: https://github.com/pydantic/pydantic/issues/10910
+    "pydantic == 2.9.2",
+
     # TODO: https://github.com/thoughtspot/thoughtspot_tml/issues/24
     "betterproto[compiler] == 2.0.0b6",