pg-batch 1.2.0__py3-none-any.whl

pg_batch/__main__.py

@@ -0,0 +1,4 @@
+from . import pg_batch
+
+if __name__ == '__main__':
+    pg_batch.main()

pg_batch/pg_batch.py

@@ -0,0 +1,283 @@
+#!/usr/bin/env python3
+
+# Author: Gabriel Bordeaux (gabfl)
+# Github: https://github.com/gabfl/pg-batch
+
+import sys
+import time
+import psycopg2
+import psycopg2.extras
+import argparse
+
+
+def update_batch(ids, table, set_, sleep=0, primary_key='id'):
+    """
+        Update a batch of rows
+    """
+
+    global confirmed_write
+
+    # Leave if ids is empty
+    if not ids or len(ids) == 0:
+        return None
+
+    # Prepare update
+    print('* Updating %i rows...' % len(ids))
+    sql = "UPDATE " + table + " SET " + set_ + \
+        " WHERE {0} IN (".format(primary_key) + \
+        ', '.join([str(x) for x in ids]) + ")"
+    print("   query: " + sql)
+
+    if confirmed_write or query_yes_no("* Start updating?"):
+        # Switch confirmed_write skip the question for the next update
+        confirmed_write = True
+
+        # Execute query
+        run_query(sql, sleep)
+    else:  # answered "no"
+        print("Error: Update declined.")
+        sys.exit()
+
+    return True
+
+
+def delete_batch(ids, table, sleep=0, primary_key='id'):
+    """
+        Delete a batch of rows
+    """
+
+    global confirmed_write
+
+    # Leave if ids is empty
+    if not ids or len(ids) == 0:
+        return None
+
+    # Prepare delete
+    print('* Deleting %i rows...' % len(ids))
+    sql = "DELETE FROM " + table + \
+        " WHERE {0} IN (".format(primary_key) + \
+        ', '.join([str(x) for x in ids]) + ")"
+    print("   query: " + sql)
+
+    if confirmed_write or query_yes_no("* Start deleting?"):
+        # Switch confirmed_write skip the question for the next delete
+        confirmed_write = True
+
+        # Execute query
+        run_query(sql, sleep)
+    else:  # answered "no"
+        print("Error: Delete declined.")
+        sys.exit()
+
+    return True
+
+
+def run_query(sql, sleep=0):
+    """Execute a write query"""
+
+    # Execute query
+    with connection.cursor() as cursorUpd:
+        cursorUpd.execute(sql)
+        connection.commit()
+
+    # Optional Sleep
+    if sleep > 0:
+        time.sleep(sleep)
+
+    return True
+
+
+def get_input():
+    """
+        Get user input
+    """
+
+    return input().lower()
+
+
+def query_yes_no(question, default="yes"):
+    """Ask a yes/no question via input() and return their answer.
+
+    "question" is a string that is presented to the user.
+    "default" is the presumed answer if the user just hits <Enter>.
+        It must be "yes" (the default), "no" or None (meaning
+        an answer is required of the user).
+
+    The "answer" return value is True for "yes" or False for "no".
+
+    (thanks https://code.activestate.com/recipes/577058/)
+    """
+    valid = {"yes": True, "y": True, "ye": True,
+             "no": False, "n": False}
+    if default is None:
+        prompt = " [y/n] "
+    elif default == "yes":
+        prompt = " [Y/n] "
+    elif default == "no":
+        prompt = " [y/N] "
+    else:
+        raise ValueError("invalid default answer: '%s'" % default)
+
+    while True:
+        sys.stdout.write(question + prompt)
+        choice = get_input()
+
+        if default is not None and choice == '':
+            return valid[default]
+        elif choice in valid:
+            return valid[choice]
+        else:
+            sys.stdout.write("Please respond with 'yes' or 'no' "
+                             "(or 'y' or 'n').\n")
+
+
+def connect(host, user, port, password, database):
+    """
+        Connect to a PostgreSQL database
+    """
+
+    # Connect to the database
+    try:
+        return psycopg2.connect(host=host,
+                                user=user,
+                                port=port,
+                                password=password,
+                                dbname=database)
+    except Exception:
+        raise RuntimeError('Error: PostgreSQL connection failed.')
+
+
+def execute(host, user, port, password, database, action, table, where, set_=None, no_confirm=False, primary_key='id', read_batch_size=10000, write_batch_size=50, sleep=0):
+    """
+        Execute batch update or delete
+    """
+
+    global confirmed_write, connection
+
+    # Make sure we have a SET clause for updates
+    if action == 'update' and set_ is None:
+        raise RuntimeError('Error: argument -s/--set is required for updates.')
+
+    # Connect to the database
+    connection = connect(host, user, port, password, database)
+
+    try:
+        # confirmed_write default value
+        confirmed_write = False
+        if no_confirm:
+            confirmed_write = True
+
+        with connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
+            # Default vars
+            min_id = 0
+
+            while 1:  # Infinite loop, will be broken by sys.exit()
+                # Get rows to modify
+                print("* Selecting data...")
+                sql = "SELECT {0} as id FROM ".format(primary_key) + table + " WHERE " + where + \
+                    " AND {0} > %s ORDER BY {1} LIMIT %s".format(
+                        primary_key, primary_key)
+                print("   query: " + sql % (min_id, read_batch_size))
+                cursor.execute(sql, (min_id, read_batch_size))
+
+                # Row count
+                count = cursor.rowcount
+
+                # No more rows
+                if count == 0:
+                    print("* No more rows to modify!")
+                    sys.exit()
+
+                # Loop thru rows
+                print("* Preparing to modify %s rows..." % count)
+                ids = []
+                for result in cursor:
+                    # Append ID to batch
+                    ids.append(result.get('id'))
+                    # print(result)
+
+                    # Minimum ID for future select
+                    min_id = result.get('id')
+
+                    # Process write when batch size if reached
+                    if len(ids) >= write_batch_size:
+                        if action == 'delete':
+                            # Process delete
+                            delete_batch(ids, table, sleep, primary_key)
+                        else:
+                            # Process update
+                            update_batch(ids, table, set_, sleep, primary_key)
+
+                        # Reset ids
+                        ids = []
+
+                # Process final batch
+                if ids:
+                    if action == 'delete':
+                        # Process delete
+                        delete_batch(ids, table, sleep, primary_key)
+                    else:
+                        # Process update
+                        update_batch(ids, table, set_, sleep, primary_key)
+    except SystemExit:
+        print("* Program exited")
+    # except:
+    #    print("Unexpected error:", sys.exc_info()[0])
+    finally:
+        connection.close()
+
+    return True
+
+
+def main():
+    # Parse arguments
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-H", "--host", default="127.0.0.1",
+                        help="PostgreSQL server host")
+    parser.add_argument("-P", "--port", type=int, default=5432,
+                        help="PostgreSQL server port")
+    parser.add_argument("-U", "--user", required=True,
+                        help="PostgreSQL user")
+    parser.add_argument("-p", "--password", default='',
+                        help="PostgreSQL password")
+    parser.add_argument("-d", "--database", required=True,
+                        help="PostgreSQL database name")
+    parser.add_argument("-t", "--table", required=True,
+                        help="PostgreSQL table")
+    parser.add_argument("-id", "--primary_key", default='id',
+                        help="Name of the primary key column")
+    parser.add_argument("-w", "--where", required=True,
+                        help="Select WHERE clause")
+    parser.add_argument("-s", "--set",
+                        help="Update SET clause")
+    parser.add_argument("-rbz", "--read_batch_size", type=int, default=10000,
+                        help="Select batch size")
+    parser.add_argument("-wbz", "--write_batch_size", type=int, default=50,
+                        help="Update/delete batch size")
+    parser.add_argument("-S", "--sleep", type=float, default=0.00,
+                        help="Sleep after each batch")
+    parser.add_argument("-a", "--action", default='update', choices=['update', 'delete'],
+                        help="Action ('update' or 'delete')")
+    parser.add_argument("-n", "--no_confirm", action='store_true',
+                        help="Don't ask for confirmation before to run the write queries")
+    args = parser.parse_args()
+
+    execute(host=args.host,
+            user=args.user,
+            port=args.port,
+            password=args.password,
+            database=args.database,
+            action=args.action,
+            table=args.table,
+            where=args.where,
+            set_=args.set,
+            no_confirm=args.no_confirm,
+            primary_key=args.primary_key,
+            read_batch_size=args.read_batch_size,
+            write_batch_size=args.write_batch_size,
+            sleep=args.sleep
+            )
+
+
+if __name__ == '__main__':
+    main()

pg_batch-1.2.0.dist-info/METADATA

@@ -0,0 +1,173 @@
+Metadata-Version: 2.4
+Name: pg_batch
+Version: 1.2.0
+Summary: Run large PostgreSQL UPDATE and DELETE queries with small batches to prevent locks
+Author-email: Gabriel Bordeaux <pypi@gab.lc>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/gabfl/pg-batch
+Project-URL: Issues, https://github.com/gabfl/pg-batch/issues
+Classifier: Topic :: Software Development
+Classifier: Topic :: Database
+Classifier: Topic :: Database :: Database Engines/Servers
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Natural Language :: English
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Development Status :: 5 - Production/Stable
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: psycopg2-binary<3,>=2.9
+Dynamic: license-file
+
+# pg-batch
+
+[![Pypi](https://img.shields.io/pypi/v/pg-batch.svg)](https://pypi.org/project/pg-batch)
+[![Build Status](https://github.com/gabfl/pg-batch/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/gabfl/pg-batch/actions)
+[![codecov](https://codecov.io/gh/gabfl/pg-batch/branch/main/graph/badge.svg)](https://codecov.io/gh/gabfl/pg-batch)
+[![MIT licensed](https://img.shields.io/badge/license-MIT-green.svg)](https://raw.githubusercontent.com/gabfl/pg-batch/main/LICENSE)
+
+PostgreSQL equivalent of https://github.com/gabfl/mysql-batch.
+
+Updating or deleting a large amount of rows in PostgreSQL will create locks that will paralyze other queries running in parallel.
+
+This tool will run UPDATE and DELETE queries in small batches to limit locking. If a large number of rows has to be updated or deleted, it is also possible to limit the number of rows selected at once.
+
+## Installation
+
+```
+pip3 install pg_batch
+```
+
+## UPDATE example
+
+You can run this example with the schema available in [sample_table/schema.sql](sample_table/schema.sql)
+
+The following example will be identical to the following update:
+
+```sql
+UPDATE batch_test SET date = NOW() WHERE number > 30 AND date is NULL;
+```
+
+This is the equivalent to process this update with batches of 20 rows:
+
+```bash
+pg_batch --host localhost \
+         --user postgres \
+         --password secret_password \
+         --database "test" \
+         --table "batch_test" \
+         --write_batch_size 20 \
+         --where "number > 30 AND date IS NULL" \
+         --set "date = NOW()"
+```
+
+Output sample:
+
+```bash
+* Selecting data...
+   query: SELECT id as id FROM batch_test WHERE number > 30 AND date IS NULL AND id > 0 ORDER BY id LIMIT 10000
+* Preparing to modify 70 rows...
+* Updating 20 rows...
+   query: UPDATE batch_test SET date = NOW() WHERE id IN (31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50)
+* Start updating? [Y/n]
+* Updating 20 rows...
+   query: UPDATE batch_test SET date = NOW() WHERE id IN (51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70)
+* Updating 20 rows...
+   query: UPDATE batch_test SET date = NOW() WHERE id IN (71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90)
+* Updating 10 rows...
+   query: UPDATE batch_test SET date = NOW() WHERE id IN (91, 92, 93, 94, 95, 96, 97, 98, 99, 100)
+* Selecting data...
+   query: SELECT id as id FROM batch_test WHERE number > 30 AND date IS NULL AND id > 100 ORDER BY id LIMIT 10000
+* No more rows to modify!
+* Program exited
+```
+
+## DELETE example
+
+The following example will be identical to the following delete:
+
+```sql
+DELETE FROM batch_test WHERE number > 30 AND date is NULL;
+```
+
+This is the equivalent to process this delete with batches of 20 rows:
+
+```bash
+pg_batch --host localhost \
+         --user postgres \
+         --password secret_password \
+         --database "test" \
+         --table "batch_test" \
+         --write_batch_size 20 \
+         --where "number > 30 AND date IS NULL" \
+         --action "delete"
+```
+
+Output sample:
+
+```bash
+* Selecting data...
+   query: SELECT id as id FROM batch_test WHERE number > 30 AND date IS NULL AND id > 0 ORDER BY id LIMIT 10000
+* Preparing to modify 70 rows...
+* Deleting 20 rows...
+   query: DELETE FROM batch_test WHERE id IN (31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50)
+* Start deleting? [Y/n]
+* Deleting 20 rows...
+   query: DELETE FROM batch_test WHERE id IN (51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70)
+* Deleting 20 rows...
+   query: DELETE FROM batch_test WHERE id IN (71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90)
+* Deleting 10 rows...
+   query: DELETE FROM batch_test WHERE id IN (91, 92, 93, 94, 95, 96, 97, 98, 99, 100)
+* Selecting data...
+   query: SELECT id as id FROM batch_test WHERE number > 30 AND date IS NULL AND id > 100 ORDER BY id LIMIT 10000
+* No more rows to modify!
+* Program exited
+```
+
+## Usage
+
+```bash
+usage: pg_batch [-h] [-H HOST] [-P PORT] -U USER [-p PASSWORD] -d DATABASE -t
+                TABLE [-id PRIMARY_KEY] -w WHERE [-s SET]
+                [-rbz READ_BATCH_SIZE] [-wbz WRITE_BATCH_SIZE] [-S SLEEP]
+                [-a {update,delete}] [-n]
+
+options:
+  -h, --help            show this help message and exit
+  -H HOST, --host HOST  PostgreSQL server host
+  -P PORT, --port PORT  PostgreSQL server port
+  -U USER, --user USER  PostgreSQL user
+  -p PASSWORD, --password PASSWORD
+                        PostgreSQL password
+  -d DATABASE, --database DATABASE
+                        PostgreSQL database name
+  -t TABLE, --table TABLE
+                        PostgreSQL table
+  -id PRIMARY_KEY, --primary_key PRIMARY_KEY
+                        Name of the primary key column
+  -w WHERE, --where WHERE
+                        Select WHERE clause
+  -s SET, --set SET     Update SET clause
+  -rbz READ_BATCH_SIZE, --read_batch_size READ_BATCH_SIZE
+                        Select batch size
+  -wbz WRITE_BATCH_SIZE, --write_batch_size WRITE_BATCH_SIZE
+                        Update/delete batch size
+  -S SLEEP, --sleep SLEEP
+                        Sleep after each batch
+  -a {update,delete}, --action {update,delete}
+                        Action ('update' or 'delete')
+  -n, --no_confirm      Don't ask for confirmation before to run the write
+                        queries
+```
+
+## License
+
+This program is under MIT license ([view license](LICENSE)).

pg_batch-1.2.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+pg_batch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pg_batch/__main__.py,sha256=0XMLusKiaJR-vIpHD3JChZGFhyJL7KNyRO1FTB623XY,71
+pg_batch/pg_batch.py,sha256=MQU2Nz5qvhJUXyhpwbFro5-rUkOnTS7mRWmZhwqIf-c,9034
+pg_batch-1.2.0.dist-info/licenses/LICENSE,sha256=WqS7GgVJk_T1GouVCLKKsR1q83EOah9Hpb72mq0pQ64,1073
+pg_batch-1.2.0.dist-info/METADATA,sha256=R6j7tFu4olPfvEXnDfwPh5KWonY5977Og4k81nMh1Aw,6690
+pg_batch-1.2.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+pg_batch-1.2.0.dist-info/entry_points.txt,sha256=cV5U089FFWzVfM7SgVPHqwWjA8gMp7au1SCq_BtNf4M,52
+pg_batch-1.2.0.dist-info/top_level.txt,sha256=MezhyRgDIbdatUfJA5a2GbYU0avk33zBuGqAa2ZKM-w,9
+pg_batch-1.2.0.dist-info/RECORD,,

pg_batch-1.2.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

pg_batch-1.2.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+pg_batch = pg_batch.pg_batch:main

pg_batch-1.2.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2017 Gabriel Bordeaux
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pg_batch-1.2.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+pg_batch