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

felixxm · felixxm · commit a03f2981dc58 · 2024-01-23T21:51:02.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
@@ -1983,7 +1983,7 @@ 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)
+pysqlite_connection_iterdump_impl(pysqlite_Connection *self, const char *filter)
 /*[clinic end generated code: output=586997aaf9808768 input=1911ca756066da89]*/
 {
     if (!pysqlite_check_connection(self)) {
@@ -1998,8 +1998,9 @@ pysqlite_connection_iterdump_impl(pysqlite_Connection *self)
         }
         return NULL;
     }
-
-    PyObject *retval = PyObject_CallOneArg(iterdump, (PyObject *)self);
+    PyObject *py_filter = NULL;
+    py_filter = PyUnicode_FromString(filter);
+    PyObject *retval = PyObject_Call(iterdump, (PyObject *)self, py_filter);
     Py_DECREF(iterdump);
     return retval;
 }

Original file line number	Diff line number	Diff line change
`@@ -1983,7 +1983,7 @@ Returns iterator to the dump of the database in an SQL text format.`
`1983`	`1983`	`[clinic start generated code]*/`
`1984`	`1984`
`1985`	`1985`	`static PyObject *`
`1986`		`-pysqlite_connection_iterdump_impl(pysqlite_Connection *self)`
	`1986`	`+pysqlite_connection_iterdump_impl(pysqlite_Connection self, const char filter)`
`1987`	`1987`	`/[clinic end generated code: output=586997aaf9808768 input=1911ca756066da89]/`
`1988`	`1988`	`{`
`1989`	`1989`	`if (!pysqlite_check_connection(self)) {`
`@@ -1998,8 +1998,9 @@ pysqlite_connection_iterdump_impl(pysqlite_Connection *self)`
`1998`	`1998`	`}`
`1999`	`1999`	`return NULL;`
`2000`	`2000`	`}`
`2001`		`-`
`2002`		`- PyObject retval = PyObject_CallOneArg(iterdump, (PyObject )self);`
	`2001`	`+ PyObject *py_filter = NULL;`
	`2002`	`+ py_filter = PyUnicode_FromString(filter);`
	`2003`	`+ PyObject retval = PyObject_Call(iterdump, (PyObject )self, py_filter);`
`2003`	`2004`	`Py_DECREF(iterdump);`
`2004`	`2005`	`return retval;`
`2005`	`2006`	`}`