tablemaster 2.1.9__tar.gz → 2.1.10__tar.gz

{tablemaster-2.1.9 → tablemaster-2.1.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tablemaster
-Version: 2.1.9
+Version: 2.1.10
 Summary: tablemaster is a Python toolkit for moving and managing tabular data across databases, Feishu/Lark, Google Sheets, and local files with one consistent API.
 Author-email: Livid <livid.su@gmail.com>
 Project-URL: Homepage, https://github.com/ilivid/tablemaster

{tablemaster-2.1.9 → tablemaster-2.1.10}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tablemaster"
-version = "2.1.9"
+version = "2.1.10"
 description = "tablemaster is a Python toolkit for moving and managing tabular data across databases, Feishu/Lark, Google Sheets, and local files with one consistent API."
 readme = "README.md"
 requires-python = ">=3.9"

{tablemaster-2.1.9 → tablemaster-2.1.10}/tablemaster/database.py

@@ -130,6 +130,19 @@ def _safe_identifier(identifier: str) -> str:
     return identifier
 
 
+def _safe_postgresql_table(table: str) -> str:
+    return '.'.join(_safe_identifier(part.strip()) for part in table.split('.'))
+
+
+def _execute_postgresql_values(connection: Any, sql: str, rows: List[Tuple[Any, ...]], page_size: int) -> None:
+    from psycopg2.extras import execute_values
+
+    proxied_connection = connection.connection
+    dbapi_connection = getattr(proxied_connection, 'driver_connection', proxied_connection)
+    with dbapi_connection.cursor() as cursor:
+        execute_values(cursor, sql, rows, page_size=page_size)
+
+
 def _safe_mysql_type(data_type: str) -> str:
     """
     Ensure a MySQL data type expression is safe from SQL injection.
@@ -383,6 +396,7 @@ class ManageTable:
                                     
                                 safe_keys = [_safe_identifier(k) for k in keys]
                                 safe_columns = [_safe_identifier(col) for col in columns]
+                                safe_table = _safe_postgresql_table(self.table)
                                 quoted_columns = ', '.join([f'"{col}"' for col in safe_columns])
                                 update_columns = ', '.join(
                                     [f'"{col}"=EXCLUDED."{col}"' for col in safe_columns if col not in safe_keys]
@@ -391,23 +405,45 @@ class ManageTable:
                                 
                                 if update_columns:
                                     insert_sql = f"""
-                                    INSERT INTO {self.table} ({quoted_columns})
-                                    VALUES ({value_placeholders})
+                                    INSERT INTO {safe_table} ({quoted_columns})
+                                    VALUES %s
                                     ON CONFLICT ({conflict_keys_str}) DO UPDATE SET {update_columns}
                                     """
                                 else:
                                     insert_sql = f"""
-                                    INSERT INTO {self.table} ({quoted_columns})
-                                    VALUES ({value_placeholders})
+                                    INSERT INTO {safe_table} ({quoted_columns})
+                                    VALUES %s
                                     ON CONFLICT ({conflict_keys_str}) DO NOTHING
                                     """
+                                data_frame = chunk.astype(object).where(pd.notna(chunk), None)
+                                data = [tuple(row) for row in data_frame.itertuples(index=False, name=None)]
+                                _execute_postgresql_values(connection, insert_sql, data, page_size=chunk_size)
+                                pbar.update(1)
+                                continue
                             else:
                                 raise ValueError(f'Unsupported db_type for upsert: {db_type}')
                         else:
-                            insert_sql = f"""
-                            INSERT IGNORE INTO {self.table} ({', '.join([f'`{col}`' for col in columns])})
-                            VALUES ({value_placeholders})
-                            """
+                            if db_type in ('mysql', 'tidb'):
+                                insert_sql = f"""
+                                INSERT IGNORE INTO {self.table} ({', '.join([f'`{col}`' for col in columns])})
+                                VALUES ({value_placeholders})
+                                """
+                            elif db_type == 'postgresql':
+                                safe_columns = [_safe_identifier(col) for col in columns]
+                                safe_table = _safe_postgresql_table(self.table)
+                                quoted_columns = ', '.join([f'"{col}"' for col in safe_columns])
+                                insert_sql = f"""
+                                INSERT INTO {safe_table} ({quoted_columns})
+                                VALUES %s
+                                ON CONFLICT DO NOTHING
+                                """
+                                data_frame = chunk.astype(object).where(pd.notna(chunk), None)
+                                data = [tuple(row) for row in data_frame.itertuples(index=False, name=None)]
+                                _execute_postgresql_values(connection, insert_sql, data, page_size=chunk_size)
+                                pbar.update(1)
+                                continue
+                            else:
+                                raise ValueError(f'Unsupported db_type for upsert: {db_type}')
 
                         data = chunk.astype(object).where(pd.notna(chunk), None).to_dict(orient='records')
                         connection.execute(text(insert_sql), data)

{tablemaster-2.1.9 → tablemaster-2.1.10}/tablemaster.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tablemaster
-Version: 2.1.9
+Version: 2.1.10
 Summary: tablemaster is a Python toolkit for moving and managing tabular data across databases, Feishu/Lark, Google Sheets, and local files with one consistent API.
 Author-email: Livid <livid.su@gmail.com>
 Project-URL: Homepage, https://github.com/ilivid/tablemaster

{tablemaster-2.1.9 → tablemaster-2.1.10}/tablemaster.egg-info/SOURCES.txt

@@ -32,4 +32,5 @@ tablemaster/schema/dialects/mysql.py
 tablemaster/schema/dialects/postgresql.py
 tablemaster/schema/dialects/tidb.py
 tests/test_error_visibility.py
+tests/test_postgresql_upsert.py
 tests/test_schema_core.py

tablemaster-2.1.10/tests/test_postgresql_upsert.py

@@ -0,0 +1,63 @@
+from types import SimpleNamespace
+from unittest import TestCase
+from unittest.mock import patch
+
+import pandas as pd
+
+from tablemaster.database import ManageTable
+
+
+class _BeginContext:
+    def __enter__(self):
+        return SimpleNamespace()
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+
+class _Engine:
+    def begin(self):
+        return _BeginContext()
+
+
+def _postgresql_cfg():
+    return SimpleNamespace(
+        name='pg',
+        user='u',
+        password='p',
+        host='127.0.0.1',
+        database='d',
+        db_type='postgresql',
+    )
+
+
+class PostgreSQLUpsertTests(TestCase):
+    def test_postgresql_upsert_uses_execute_values(self):
+        table = ManageTable('orders', _postgresql_cfg())
+        df = pd.DataFrame({'id': [1, 2], 'name': ['alpha', None]})
+
+        with patch('tablemaster.database._resolve_engine', return_value=_Engine()):
+            with patch('tablemaster.database._execute_postgresql_values') as execute_values:
+                table.upsert_data(df, key='id')
+
+        execute_values.assert_called_once()
+        _, sql, rows = execute_values.call_args.args
+        self.assertIn('INSERT INTO orders ("id", "name")', sql)
+        self.assertIn('VALUES %s', sql)
+        self.assertIn('ON CONFLICT ("id") DO UPDATE SET "name"=EXCLUDED."name"', sql)
+        self.assertEqual([(1, 'alpha'), (2, None)], rows)
+        self.assertEqual(10000, execute_values.call_args.kwargs['page_size'])
+
+    def test_postgresql_ignore_uses_on_conflict_do_nothing(self):
+        table = ManageTable('public.orders', _postgresql_cfg())
+        df = pd.DataFrame({'id': [1], 'name': ['alpha']})
+
+        with patch('tablemaster.database._resolve_engine', return_value=_Engine()):
+            with patch('tablemaster.database._execute_postgresql_values') as execute_values:
+                table.upsert_data(df, ignore=True)
+
+        execute_values.assert_called_once()
+        _, sql, rows = execute_values.call_args.args
+        self.assertIn('INSERT INTO public.orders ("id", "name")', sql)
+        self.assertIn('ON CONFLICT DO NOTHING', sql)
+        self.assertEqual([(1, 'alpha')], rows)