fix PG_LOGICAL_SLOT_PEEK_CHANGES query crashe #546

examples/book/schema.json

@@ -55,8 +55,7 @@
             "Hero",
             "Villain",
             "Geometry",
-            "Infinity",
-            "TextEmbedding3Small"
+            "Infinity"
         ],
         "nodes": {
             "table": "book",

pgsync/base.py

@@ -529,6 +529,42 @@ def _logical_slot_changes(
             statement = statement.offset(offset)
         return statement
 
+    def max_lsn(
+        self,
+        slot_name: str,
+        txmin: t.Optional[int] = None,
+        txmax: t.Optional[int] = None,
+    ):
+        filters: list = []
+        statement: sa.sql.Select = sa.select(
+            sa.func.MAX(sa.text("lsn")),
+        ).select_from(
+            sa.func.PG_LOGICAL_SLOT_PEEK_CHANGES(
+                slot_name,
+                None,
+                None,
+            )
+        )
+        if txmin is not None:
+            filters.append(
+                sa.cast(
+                    sa.cast(sa.column("xid"), sa.Text),
+                    sa.BigInteger,
+                )
+                >= txmin
+            )
+        if txmax is not None:
+            filters.append(
+                sa.cast(
+                    sa.cast(sa.column("xid"), sa.Text),
+                    sa.BigInteger,
+                )
+                < txmax
+            )
+        if filters:
+            statement = statement.where(sa.and_(*filters))
+        return self.fetchone(statement)[0]
+
     def logical_slot_get_changes(
         self,
         slot_name: str,

pgsync/querybuilder.py

@@ -34,7 +34,7 @@ def _eval_expression(
             ):
                 # handle UUID typed expressions:
                 # psycopg2.errors.UndefinedFunction: operator does not exist: uuid = integer
-                return expression.left == None
+                return expression.left is None
 
         return expression
 

pgsync/sync.py

@@ -377,25 +377,21 @@ def logical_slot_changes(
         # minimize the tmp file disk usage when calling
         # PG_LOGICAL_SLOT_PEEK_CHANGES and PG_LOGICAL_SLOT_GET_CHANGES
         # by limiting to a smaller batch size.
-        offset: int = 0
-        total: int = 0
-        limit: int = settings.LOGICAL_SLOT_CHUNK_SIZE
-        count: int = self.logical_slot_count_changes(
-            self.__name,
-            txmin=txmin,
-            txmax=txmax,
-            upto_nchanges=upto_nchanges,
-        )
+
+        upto_nchanges: int = upto_nchanges or settings.LOGICAL_SLOT_CHUNK_SIZE
+
+        # this is the max lsn we can go upto
+        max_lsn: int = self.max_lsn(self.__name, txmin=txmin, txmax=txmax)
+
         while True:
             changes: int = self.logical_slot_peek_changes(
                 self.__name,
                 txmin=txmin,
                 txmax=txmax,
                 upto_nchanges=upto_nchanges,
-                limit=limit,
-                offset=offset,
+                upto_lsn=max_lsn,
             )
-            if not changes or total > count:
+            if not changes:
                 break
 
             rows: list = []
@@ -450,11 +446,8 @@ def logical_slot_changes(
                 txmin=txmin,
                 txmax=txmax,
                 upto_nchanges=upto_nchanges,
-                limit=limit,
-                offset=offset,
+                upto_lsn=max_lsn,
             )
-            offset += limit
-            total += len(changes)
             self.count["xlog"] += len(rows)
 
     def _root_primary_key_resolver(

requirements/base.txt

@@ -6,9 +6,9 @@
 #
 async-timeout==4.0.3
     # via redis
-boto3==1.34.124
+boto3==1.34.138
     # via -r requirements/base.in
-botocore==1.34.124
+botocore==1.34.138
     # via
     #   boto3
     #   s3transfer
@@ -61,31 +61,31 @@ python-dotenv==1.0.1
     # via
     #   -r requirements/base.in
     #   environs
-redis==5.0.5
+redis==5.0.7
     # via -r requirements/base.in
 requests==2.32.3
     # via
     #   opensearch-py
     #   requests-aws4auth
 requests-aws4auth==1.2.3
     # via -r requirements/base.in
-s3transfer==0.10.1
+s3transfer==0.10.2
     # via boto3
 six==1.16.0
     # via
     #   opensearch-dsl
     #   opensearch-py
     #   python-dateutil
     #   requests-aws4auth
-sqlalchemy==2.0.30
+sqlalchemy==2.0.31
     # via -r requirements/base.in
 sqlparse==0.5.0
     # via -r requirements/base.in
 typing-extensions==4.12.2
     # via
     #   elasticsearch-dsl
     #   sqlalchemy
-urllib3==1.26.18
+urllib3==1.26.19
     # via
     #   botocore
     #   elastic-transport

requirements/dev.txt

@@ -8,9 +8,9 @@ async-timeout==4.0.3
     # via redis
 black==24.4.2
     # via -r requirements/dev.in
-boto3==1.34.124
+boto3==1.34.138
     # via -r requirements/base.in
-botocore==1.34.124
+botocore==1.34.138
     # via
     #   boto3
     #   s3transfer
@@ -27,7 +27,7 @@ click==8.1.7
     # via
     #   -r requirements/base.in
     #   black
-coverage[toml]==7.5.3
+coverage[toml]==7.5.4
     # via
     #   -r requirements/dev.in
     #   pytest-cov
@@ -47,11 +47,11 @@ events==0.5
     # via opensearch-py
 exceptiongroup==1.2.1
     # via pytest
-faker==25.8.0
+faker==26.0.0
     # via -r requirements/dev.in
-filelock==3.14.0
+filelock==3.15.4
     # via virtualenv
-flake8==7.0.0
+flake8==7.1.0
     # via -r requirements/dev.in
 freezegun==1.5.1
     # via -r requirements/dev.in
@@ -100,7 +100,7 @@ pre-commit==3.5.0
     # via -r requirements/dev.in
 psycopg2-binary==2.9.9
     # via -r requirements/base.in
-pycodestyle==2.11.1
+pycodestyle==2.12.0
     # via flake8
 pyflakes==3.2.0
     # via flake8
@@ -127,23 +127,23 @@ python-dotenv==1.0.1
     #   environs
 pyyaml==6.0.1
     # via pre-commit
-redis==5.0.5
+redis==5.0.7
     # via -r requirements/base.in
 requests==2.32.3
     # via
     #   opensearch-py
     #   requests-aws4auth
 requests-aws4auth==1.2.3
     # via -r requirements/base.in
-s3transfer==0.10.1
+s3transfer==0.10.2
     # via boto3
 six==1.16.0
     # via
     #   opensearch-dsl
     #   opensearch-py
     #   python-dateutil
     #   requests-aws4auth
-sqlalchemy==2.0.30
+sqlalchemy==2.0.31
     # via -r requirements/base.in
 sqlparse==0.5.0
     # via -r requirements/base.in
@@ -157,11 +157,11 @@ typing-extensions==4.12.2
     #   black
     #   elasticsearch-dsl
     #   sqlalchemy
-urllib3==1.26.18
+urllib3==1.26.19
     # via
     #   botocore
     #   elastic-transport
     #   opensearch-py
     #   requests
-virtualenv==20.26.2
+virtualenv==20.26.3
     # via pre-commit

tests/test_sync.py

@@ -81,9 +81,8 @@ def test_logical_slot_changes(self, mock_logger, sync):
                     "testdb_testdb",
                     txmin=None,
                     txmax=None,
-                    upto_nchanges=None,
-                    limit=settings.LOGICAL_SLOT_CHUNK_SIZE,
-                    offset=0,
+                    upto_nchanges=settings.LOGICAL_SLOT_CHUNK_SIZE,
+                    upto_lsn=None,
                 )
                 mock_sync.assert_not_called()
 
@@ -98,9 +97,8 @@ def test_logical_slot_changes(self, mock_logger, sync):
                     "testdb_testdb",
                     txmin=None,
                     txmax=None,
-                    upto_nchanges=None,
-                    limit=settings.LOGICAL_SLOT_CHUNK_SIZE,
-                    offset=0,
+                    upto_nchanges=settings.LOGICAL_SLOT_CHUNK_SIZE,
+                    upto_lsn=None,
                 )
                 mock_sync.assert_not_called()
 
@@ -127,9 +125,8 @@ def test_logical_slot_changes(self, mock_logger, sync):
                         "testdb_testdb",
                         txmin=None,
                         txmax=None,
-                        upto_nchanges=None,
-                        limit=settings.LOGICAL_SLOT_CHUNK_SIZE,
-                        offset=0,
+                        upto_nchanges=settings.LOGICAL_SLOT_CHUNK_SIZE,
+                        upto_lsn=None,
                     )
                     mock_get.assert_called_once()
                     mock_sync.assert_called_once()