Skip to content

Commit

Permalink
Merge branch 'dev'
Browse files Browse the repository at this point in the history
  • Loading branch information
boonhapus committed Nov 22, 2024
2 parents 16a0fab + 011c470 commit 0573705
Show file tree
Hide file tree
Showing 5 changed files with 31 additions and 7 deletions.
2 changes: 1 addition & 1 deletion cs_tools/__project__.py
Original file line number Diff line number Diff line change
@@ -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/"
Expand Down
24 changes: 23 additions & 1 deletion cs_tools/sync/sqlite/syncer.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import logging
import pathlib

from sqlalchemy.dialects.sqlite import insert
import pydantic
import sqlalchemy as sa

Expand Down Expand Up @@ -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):
# """ """
Expand Down Expand Up @@ -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()
6 changes: 3 additions & 3 deletions cs_tools/sync/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []
Expand All @@ -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()

Expand Down
2 changes: 1 addition & 1 deletion cs_tools/updater/_updater.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
4 changes: 3 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,6 @@ dependencies = [
"thoughtspot_tml",
"awesomeversion",
"httpx >= 0.27.0",
"pydantic >= 2.6.4",
"pydantic-settings",
"email-validator",
"pendulum >= 3.0.0",
Expand All @@ -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",

Expand Down

0 comments on commit 0573705

Please sign in to comment.