gh-91602: Add iterdump() support for filtering database objects.

felixxm · felixxm · commit 46302ad91345 · 2024-01-24T05:34:46.000+01:00
diff --git a/Lib/sqlite3/dump.py b/Lib/sqlite3/dump.py
@@ -15,7 +15,7 @@ def _quote_value(value):
     return "'{0}'".format(value.replace("'", "''"))
 
 
-def _iterdump(connection):
+def _iterdump(connection, *, filter=None):
     """
     Returns an iterator to the dump of the database in an SQL text format.
 
@@ -32,15 +32,23 @@ def _iterdump(connection):
         yield('PRAGMA foreign_keys=OFF;')
     yield('BEGIN TRANSACTION;')
 
+    if filter:
+        # Return database objects which match the filter pattern.
+        filter_name_clause = 'AND "name" LIKE %s'
+        params = [filter]
+    else:
+        filter_name_clause = ""
+        params = []
     # sqlite_master table contains the SQL CREATE statements for the database.
-    q = """
+    q = f"""
         SELECT "name", "type", "sql"
         FROM "sqlite_master"
             WHERE "sql" NOT NULL AND
             "type" == 'table'
+            {filter_name_clause}
             ORDER BY "name"
         """
-    schema_res = cu.execute(q)
+    schema_res = cu.execute(q, params)
     sqlite_sequence = []
     for table_name, type, sql in schema_res.fetchall():
         if table_name == 'sqlite_sequence':
@@ -82,11 +90,12 @@ def _iterdump(connection):
             yield("{0};".format(row[0]))
 
     # Now when the type is 'index', 'trigger', or 'view'
-    q = """
+    q = f"""
         SELECT "name", "type", "sql"
         FROM "sqlite_master"
             WHERE "sql" NOT NULL AND
             "type" IN ('index', 'trigger', 'view')
+            {filter_name_clause}
         """
     schema_res = cu.execute(q)
     for name, type, sql in schema_res.fetchall():
diff --git a/Lib/test/test_sqlite3/test_dump.py b/Lib/test/test_sqlite3/test_dump.py
@@ -54,6 +54,29 @@ def test_table_dump(self):
         [self.assertEqual(expected_sqls[i], actual_sqls[i])
             for i in range(len(expected_sqls))]
 
+    def test_table_dump_filter(self):
+        all_table_sqls = [
+            """CREATE TABLE "test_table_1" ("id_2" INTEGER);""",
+            """INSERT INTO "test_table_1" VALUES(1);""",
+            """INSERT INTO "test_table_1" VALUES(2);""",
+            """CREATE TABLE "some_table_2" ("id_1" INTEGER);""",
+            """INSERT INTO "some_table_2" VALUES(3);""",
+            """INSERT INTO "some_table_2" VALUES(4);""",
+        ]
+        all_views_sqls = [
+            """CREATE VIEW view_1 AS SELECT * FROM some_table_2""",
+            """CREATE VIEW view_2 AS SELECT * FROM test_table_1""",
+        ]
+        # Create database structure.
+        for sql in [*all_table_sqls, *all_views_sqls]:
+            self.cu.execute(sql)
+        # %_table_% matches all tables.
+        dump_sqls = list(self.cx.iterdump(filter="%_table_%"))
+        self.assertEqual(dump_sqls, all_table_sqls)
+        # view_% matches all views.
+        dump_sqls = list(self.cx.iterdump(filter="view_%"))
+        self.assertEqual(dump_sqls, all_views_sqls)
+
     def test_dump_autoincrement(self):
         expected = [
             'CREATE TABLE "t1" (id integer primary key autoincrement);',
diff --git a/Modules/_sqlite/clinic/connection.c.h b/Modules/_sqlite/clinic/connection.c.h
diff --git a/Modules/_sqlite/connection.c b/Modules/_sqlite/connection.c
@@ -1979,12 +1979,17 @@ pysqlite_connection_interrupt_impl(pysqlite_Connection *self)
 /*[clinic input]
 _sqlite3.Connection.iterdump as pysqlite_connection_iterdump
 
+    *
+    filter: object = None
+        An optional LIKE pattern for database objects to dump
+
 Returns iterator to the dump of the database in an SQL text format.
 [clinic start generated code]*/
 
 static PyObject *
-pysqlite_connection_iterdump_impl(pysqlite_Connection *self)
-/*[clinic end generated code: output=586997aaf9808768 input=1911ca756066da89]*/
+pysqlite_connection_iterdump_impl(pysqlite_Connection *self,
+                                  PyObject *filter)
+/*[clinic end generated code: output=fd81069c4bdeb6b0 input=4ae6d9a898f108df]*/
 {
     if (!pysqlite_check_connection(self)) {
         return NULL;
@@ -1998,8 +2003,7 @@ pysqlite_connection_iterdump_impl(pysqlite_Connection *self)
         }
         return NULL;
     }
-
-    PyObject *retval = PyObject_CallOneArg(iterdump, (PyObject *)self);
+    PyObject *retval = PyObject_Call(iterdump, (PyObject *)self, filter);
     Py_DECREF(iterdump);
     return retval;
 }
