fix(db-manager): broken select/fetch logic (#92)

* fix(db-manager): broken select/fetch logic

* fix tests

* fix tests

* fix tests

* fix tests

* fix tests

* fix tests

* fix tests

* fix tests

* fix tests

* fix tests

* fix tests

Author: azarz

changelog.md

@@ -1,5 +1,9 @@
 # CHANGELOG
 
+## 3.1.6
+FIXED :
+- wrong database manager logic for select + batch fetch
+
 ## 3.1.5
 FIXED :
 - database manager does not raie error when retries have failed
@@ -18,7 +22,7 @@ CHANGED:
 
 ## 3.1.0
 ADDED:
-- fields vla_par_defaut, cout_penalites, vehicule_leger_interdit, cout_vehicule_prioritaire 
+- fields vla_par_defaut, cout_penalites, vehicule_leger_interdit, cout_vehicule_prioritaire
 
 ## 3.0.0
 REMOVED:

r2gg/__about__.py

@@ -34,7 +34,7 @@
 __uri_tracker__ = f"{__uri_repository__}issues/"
 __uri__ = __uri_repository__
 
-__version__ = "3.1.5"
+__version__ = "3.1.6"
 __version_info__ = tuple(
     [
         int(num) if num.isdigit() else num

r2gg/_database.py

@@ -108,56 +108,70 @@ def ensure_connection(self):
             time.sleep(DELAY)
             self._connection = self.connect_working_db()
 
-    def execute_select_query(self, cursor, query, show_duration):
-        if TIMEOUT:
-            cursor.execute("SET statement_timeout = %s", (1000 * TIMEOUT,))  # timeout in milliseconds
-
-        if show_duration:
-            self.logger.info("SQL: {}".format(query))
-            st_execute = time.time()
-            cursor.execute(query)
-            et_execute = time.time()
-            self.logger.info("Execution ended. Elapsed time : %s seconds." % (et_execute - st_execute))
-        else:
-            cursor.execute(query)
-
-    @database_retry_decorator
+    # IMPORTANT:
+    # Streaming SELECTs must NOT use retry logic.
+    # If the connection drops, the cursor state is unrecoverable.
     def execute_select_fetch_multiple(self, query, batchsize=1, show_duration=False):
-        with self._connection.cursor(cursor_factory=DictCursor) as cursor:
-            self.execute_select_query(cursor, query, show_duration)
-            rows = cursor.fetchmany(batchsize)
+        """
+        Streaming SELECT using a named server-side cursor.
+        No retry. No reconnect. No commit.
+        Fail fast if the connection drops (old behavior).
+        """
+        self.ensure_connection()
+        cursor_name = f"cursor_{int(time.time() * 1000)}"
+        with self._connection.cursor(cursor_factory=DictCursor, name=cursor_name) as cursor:
+            if TIMEOUT:
+                cursor.execute("SET statement_timeout = %s", (1000 * TIMEOUT,))
+            if show_duration:
+                self.logger.info(f"SQL: {query}")
+                st = time.time()
+                cursor.execute(query)
+                self.logger.info(
+                    "Execution ended. Elapsed time : %s seconds.",
+                    time.time() - st
+                )
+            else:
+                cursor.execute(query)
+
             count = cursor.rowcount
-            while rows:
+
+            while True:
+                rows = cursor.fetchmany(batchsize)
+                if not rows:
+                    break
                 if batchsize == 1:
                     rows = rows.pop()
                 yield rows, count
-                rows = cursor.fetchmany(batchsize)
-            self._connection.commit()
-            return
 
     # the method below should be used as a generator function otherwise use execute_update
     @database_retry_decorator
     def execute_update_query(self, query, params=None, isolation_level=None, show_duration=False):
+        self.ensure_connection()
         if show_duration :
             self.logger.info("SQL: {}".format(query))
         st_execute = time.time()
         with self._connection.cursor(cursor_factory=DictCursor) as cursor:
             old_isolation_level = self._connection.isolation_level
-            if isolation_level is not None:
-                self._connection.set_isolation_level(isolation_level)
-            cursor.execute(query, params)
-            self._connection.commit()
+            try:
+                if isolation_level is not None:
+                    self._connection.set_isolation_level(isolation_level)
+                cursor.execute(query, params)
+                self._connection.commit()
+            finally:
+                self._connection.set_isolation_level(old_isolation_level)
             if show_duration:
                 et_execute = time.time()
                 self.logger.info("Execution ended. Elapsed time : %s seconds." % (et_execute - st_execute))
-            self._connection.set_isolation_level(old_isolation_level)
         yield  # the decorator database_retry_decorator only supports generators
         return
 
     def execute_update(self, query, params=None, isolation_level=None):
         next(self.execute_update_query(query, params=params, isolation_level=isolation_level), None)
 
     def execute_select_fetch_one(self, query, show_duration=False):
-        gen = self.execute_select_fetch_multiple(query, 1, show_duration)
-        row, count = next(gen, (None, None))
-        return row, count
+        try:
+            gen = self.execute_select_fetch_multiple(query, 1, show_duration)
+            row, count = next(gen, (None, None))
+            return row, count
+        finally:
+            gen.close()  # Ensure the generator is closed to free resources

r2gg/_pivot_to_osm.py

@@ -81,10 +81,13 @@ def pivot_to_osm(config, source, db_configs, database: DatabaseManager, logger,
                     sql_query_nodes += " LIMIT {} OFFSET {}".format(batchsize, offset)
                     offset += batchsize
                     logger.info("Writing nodes")
-                    for row, count in database.execute_select_fetch_multiple(sql_query_nodes, show_duration=True):
-                        nodeEl = writeNode(row, extraction_date)
-                        xf.write(nodeEl, pretty_print=True)
-
+                    gen = database.execute_select_fetch_multiple(sql_query_nodes, show_duration=True)
+                    try:
+                        for row, count in gen:
+                            nodeEl = writeNode(row, extraction_date)
+                            xf.write(nodeEl, pretty_print=True)
+                    finally:
+                        gen.close()
                     logger.info("%s / %s nodes ajoutés" % (offset, nodesize))
                 et_nodes = time.time()
                 logger.info("Writing nodes ended. Elapsed time : %s seconds." % (et_nodes - st_nodes))
@@ -124,16 +127,19 @@ def pivot_to_osm(config, source, db_configs, database: DatabaseManager, logger,
                 logger.info("Writing restrictions")
                 st_execute = time.time()
                 i = 1
-                for row, count in database.execute_select_fetch_multiple(sql_query_non_comm, show_duration=True):
-                    if row['common_vertex_id'] == -1:
+                gen = database.execute_select_fetch_multiple(sql_query_non_comm, show_duration=True)
+                try:
+                    for row, count in gen:
+                        if row['common_vertex_id'] == -1:
+                            i += 1
+                            continue
+                        ResEl = writeRes(row, i, extraction_date)
+                        xf.write(ResEl, pretty_print=True)
+                        if count > 0 and (i % ceil(count / 10) == 0):
+                            logger.info("%s / %s restrictions ajoutés" % (i, count))
                         i += 1
-                        continue
-                    ResEl = writeRes(row, i, extraction_date)
-                    xf.write(ResEl, pretty_print=True)
-                    if (i % ceil(count / 10) == 0):
-                        logger.info("%s / %s restrictions ajoutés" % (i, count))
-                    i += 1
-
+                finally:
+                    gen.close()
                 et_execute = time.time()
                 logger.info("Writing restrictions ended. Elapsed time : %s seconds." % (et_execute - st_execute))
 

r2gg/_pivot_to_pgr.py

@@ -115,38 +115,41 @@ def pivot_to_pgr(source, cost_calculation_file_path, database_work: DatabaseMana
     index = 0
     batchsize = 10000
     generator = database_work.execute_select_fetch_multiple(tr_query, show_duration=True, batchsize=batchsize)
-    rows, count = next(generator,(None, None))
-    # Insertion petit à petit -> plus performant
-
-    logger.info("SQL: Inserting or updating {} values in out db".format(count))
-
-    st_execute = time.time()
-
-    while rows:
-        values_str = ""
-        # Tuple des valuers à insérer
-        values_tuple = ()
-        for row in rows:
-            values_str += "(%s, %s, %s),"
-            values_tuple += (index, row['id_from'], row['id_to'])
-            index += 1
-        values_str = values_str[:-1]
-
-        set_on_conflict = (
-            "id_from = excluded.id_from,id_to = excluded.id_to"
-        )
+    try:
+        rows, count = next(generator,(None, None))
+        # Insertion petit à petit -> plus performant
+
+        logger.info("SQL: Inserting or updating {} values in out db".format(count))
+
+        st_execute = time.time()
+
+        while rows:
+            values_str = ""
+            # Tuple des valuers à insérer
+            values_tuple = ()
+            for row in rows:
+                values_str += "(%s, %s, %s),"
+                values_tuple += (index, row['id_from'], row['id_to'])
+                index += 1
+            values_str = values_str[:-1]
+
+            set_on_conflict = (
+                "id_from = excluded.id_from,id_to = excluded.id_to"
+            )
 
-        sql_insert = """
-            INSERT INTO {}.turn_restrictions (id, id_from, id_to)
-            VALUES {}
-            ON CONFLICT (id) DO UPDATE
-              SET {};
-        """.format(schema, values_str, set_on_conflict)
-        database_out.execute_update(sql_insert, values_tuple)
+            sql_insert = """
+                INSERT INTO {}.turn_restrictions (id, id_from, id_to)
+                VALUES {}
+                ON CONFLICT (id) DO UPDATE
+                SET {};
+            """.format(schema, values_str, set_on_conflict)
+            database_out.execute_update(sql_insert, values_tuple)
 
-        rows, _ = next(generator,(None, None))
+            rows, _ = next(generator,(None, None))
 
-    et_execute = time.time()
+        et_execute = time.time()
+    finally:
+        generator.close()  # Ensure the generator is closed to free resources
     logger.info("Writing turn restrinctions Done. Elapsed time : %s seconds." %(et_execute - st_execute))
 
     # Noeuds ---------------------------------------------------------------------------------------
@@ -171,30 +174,32 @@ def pivot_to_pgr(source, cost_calculation_file_path, database_work: DatabaseMana
     index = 0
     batchsize = 10000
     generator = database_work.execute_select_fetch_multiple(nd_query, show_duration=True, batchsize=batchsize)
-    rows, count = next(generator, (None, None))
-    while rows:
-        values_str = ""
-        # Tuple des valeurs à insérer
-        values_tuple = ()
-        for row in rows:
-            values_str += "(%s, %s),"
-            values_tuple += (row['id'], row['geom'])
-            index += 1
-        values_str = values_str[:-1]
-
-        set_on_conflict = (
-            "the_geom = excluded.the_geom"
-        )
-
-        sql_insert = """
-            INSERT INTO {}_vertices_pgr (id, the_geom)
-            VALUES {}
-            ON CONFLICT (id) DO UPDATE
-              SET {};
-        """.format(ways_table_name, values_str, set_on_conflict)
-        database_out.execute_update(sql_insert, values_tuple)
-        rows, _ = next(generator,(None, None))
+    try:
+        rows, count = next(generator, (None, None))
+        while rows:
+            values_str = ""
+            # Tuple des valeurs à insérer
+            values_tuple = ()
+            for row in rows:
+                values_str += "(%s, %s),"
+                values_tuple += (row['id'], row['geom'])
+                index += 1
+            values_str = values_str[:-1]
+
+            set_on_conflict = (
+                "the_geom = excluded.the_geom"
+            )
 
+            sql_insert = """
+                INSERT INTO {}_vertices_pgr (id, the_geom)
+                VALUES {}
+                ON CONFLICT (id) DO UPDATE
+                SET {};
+            """.format(ways_table_name, values_str, set_on_conflict)
+            database_out.execute_update(sql_insert, values_tuple)
+            rows, _ = next(generator,(None, None))
+    finally:
+        generator.close()  # Ensure the generator is closed to free resources
 
     et_execute = time.time()
     logger.info("Writing vertices Done. Elapsed time : %s seconds." %(et_execute - st_execute))
@@ -266,45 +271,48 @@ def pivot_to_pgr(source, cost_calculation_file_path, database_work: DatabaseMana
     # logger.info("SQL: Inserting or updating {} values in out db".format(cursor_in.rowcount))
     st_execute = time.time()
     percent = 0
-    rows, count = next(generator, (None, None))
-    while rows:
-        percent += 1000000 / count
-        # Chaîne permettant l'insertion de valeurs via psycopg
-        values_str = ""
-        # Tuple des valuers à insérer
-        values_tuple = ()
-        for row in rows:
-            values_str += "(" + single_value_str + "),"
-            output_costs = output_costs_from_costs_config(costs, row)
-            values_tuple += tuple(
-                row[ output_columns_name ] for output_columns_name in output_columns_names
-            ) + output_costs
-        values_str = values_str[:-1]
-
-        output_columns = "("
-        for output_columns_name in output_columns_names:
-            output_columns += output_columns_name + ','
-        output_columns = output_columns[:-1]
-
-        set_on_conflict = ''
-        for output_columns_name in output_columns_names:
-            set_on_conflict += "{0} = excluded.{0},".format(output_columns_name)
-        set_on_conflict = set_on_conflict[:-1]
-
-        for output in costs["outputs"]:
-            output_columns += "," + output["name"] + ",reverse_" + output["name"]
-            set_on_conflict += ",{0} = excluded.{0}".format(output["name"])
-            set_on_conflict += ",{0} = excluded.{0}".format("reverse_" + output["name"])
-
-        output_columns += ")"
-        sql_insert = """
-            INSERT INTO {} {}
-            VALUES {}
-            ON CONFLICT (id) DO UPDATE
-              SET {};
-            """.format(ways_table_name, output_columns, values_str, set_on_conflict)
-        database_out.execute_update(sql_insert, values_tuple)
-        rows, _ = next(generator,(None, None))
+    try:
+        rows, count = next(generator, (None, None))
+        while rows:
+            percent += 1000000 / count
+            # Chaîne permettant l'insertion de valeurs via psycopg
+            values_str = ""
+            # Tuple des valuers à insérer
+            values_tuple = ()
+            for row in rows:
+                values_str += "(" + single_value_str + "),"
+                output_costs = output_costs_from_costs_config(costs, row)
+                values_tuple += tuple(
+                    row[ output_columns_name ] for output_columns_name in output_columns_names
+                ) + output_costs
+            values_str = values_str[:-1]
+
+            output_columns = "("
+            for output_columns_name in output_columns_names:
+                output_columns += output_columns_name + ','
+            output_columns = output_columns[:-1]
+
+            set_on_conflict = ''
+            for output_columns_name in output_columns_names:
+                set_on_conflict += "{0} = excluded.{0},".format(output_columns_name)
+            set_on_conflict = set_on_conflict[:-1]
+
+            for output in costs["outputs"]:
+                output_columns += "," + output["name"] + ",reverse_" + output["name"]
+                set_on_conflict += ",{0} = excluded.{0}".format(output["name"])
+                set_on_conflict += ",{0} = excluded.{0}".format("reverse_" + output["name"])
+
+            output_columns += ")"
+            sql_insert = """
+                INSERT INTO {} {}
+                VALUES {}
+                ON CONFLICT (id) DO UPDATE
+                SET {};
+                """.format(ways_table_name, output_columns, values_str, set_on_conflict)
+            database_out.execute_update(sql_insert, values_tuple)
+            rows, _ = next(generator,(None, None))
+    finally:
+        generator.close()  # Ensure the generator is closed to free resources
 
     et_execute = time.time()
     logger.info("Writing ways ended. Elapsed time : %s seconds." %(et_execute - st_execute))