SQL: Limit tables in dropdown

VesnaT · VesnaT · commit b1d65e09cc27 · 2025-12-19T10:33:06.000+01:00
diff --git a/Orange/data/sql/backend/base.py b/Orange/data/sql/backend/base.py
@@ -68,6 +68,37 @@ def list_tables(self, schema=None):
                 tables.append(TableDesc(name, schema, sql))
             return tables
 
+    def n_tables_query(self, schema=None) -> str:
+        """Return a query to count tables in database.
+
+        Parameters
+        ----------
+        schema : Optional[str]
+            If set, only tables from schema should be listed
+
+        Returns
+        -------
+        Query string.
+        """
+        raise NotImplementedError
+
+    def n_tables(self, schema=None) -> int:
+        """Return number of tables in database.
+
+        Parameters
+        ----------
+        schema : Optional[str]
+            If set, only tables from given schema will be listed.
+
+        Returns
+        -------
+        Number of tables in the database.
+        """
+        query = self.n_tables_query(schema)
+        with self.execute_sql_query(query) as cur:
+            res = cur.fetchone()
+        return res[0]
+
     def get_fields(self, table_name):
         """Return a list of field names and metadata in the given table
 
diff --git a/Orange/data/sql/backend/mssql.py b/Orange/data/sql/backend/mssql.py
@@ -43,6 +43,9 @@ def list_tables_query(self, schema=None):
       ORDER BY [TABLE_NAME]
         """
 
+    def n_tables_query(self, _=None) -> str:
+        return "SELECT COUNT(*) FROM information_schema.tables"
+
     def quote_identifier(self, name):
         return "[{}]".format(name)
 
diff --git a/Orange/data/sql/backend/postgres.py b/Orange/data/sql/backend/postgres.py
@@ -113,6 +113,12 @@ def list_tables_query(self, schema=None):
                         AND NOT c.relname LIKE '\\_\\_%'
                    ORDER BY 1,2;""".format(schema_clause)
 
+    def n_tables_query(self, schema=None) -> str:
+        query = "SELECT COUNT(*) FROM information_schema.tables"
+        if schema:
+            query += f" WHERE table_schema = '{schema}'"
+        return query
+
     def create_variable(self, field_name, field_metadata,
                         type_hints, inspect_table=None):
         if field_name in type_hints:
diff --git a/Orange/widgets/data/owsql.py b/Orange/widgets/data/owsql.py
@@ -1,5 +1,5 @@
 from AnyQt.QtWidgets import QComboBox, QTextEdit, QMessageBox, QApplication, \
-    QGridLayout
+    QGridLayout, QLineEdit
 from AnyQt.QtGui import QCursor
 from AnyQt.QtCore import Qt
 
@@ -17,6 +17,7 @@
 from Orange.widgets.widget import Output, Msg
 
 MAX_DL_LIMIT = 1000000
+MAX_TABLES = 1000
 
 
 def is_postgres(backend):
@@ -77,6 +78,7 @@ def __init__(self):
         self.backendcombo = None
         self.tables = None
         self.tablecombo = None
+        self.tabletext = None
         self.sqltext = None
         self.custom_sql = None
         super().__init__()
@@ -126,6 +128,11 @@ def _add_tables_controls(self):
         self.tablecombo.setToolTip('table')
         self.tablecombo.activated[int].connect(self.select_table)
 
+        self.tabletext = QLineEdit(placeholderText='TABLE_NAME')
+        self.tabletext.setToolTip('table')
+        self.tabletext.editingFinished.connect(self.select_table)
+        self.tabletext.setVisible(False)
+
         self.custom_sql = gui.vBox(box)
         self.custom_sql.setVisible(self.data_source == self.CUSTOM_SQL)
         self.sqltext = QTextEdit(self.custom_sql)
@@ -142,6 +149,7 @@ def _add_tables_controls(self):
 
         form.addWidget(radio_table, 1, 0, Qt.AlignLeft)
         form.addWidget(self.tablecombo, 1, 1)
+        form.addWidget(self.tabletext, 1, 1)
         form.addWidget(radio_custom_sql, 2, 0, Qt.AlignLeft)
 
         gui.checkBox(box, self, "guess_values",
@@ -185,9 +193,15 @@ def refresh_tables(self):
             return
 
         self.tables.append("Select a table")
-        self.tables.extend(self.backend.list_tables(self.schema))
-        index = self.tablecombo.findText(str(self.table))
-        self.tablecombo.setCurrentIndex(index if index != -1 else 0)
+        if self.backend.n_tables(self.schema) <= MAX_TABLES:
+            self.tables.extend(self.backend.list_tables(self.schema))
+            index = self.tablecombo.findText(str(self.table))
+            self.tablecombo.setCurrentIndex(index if index != -1 else 0)
+            self.tablecombo.setVisible(True)
+            self.tabletext.setVisible(False)
+        else:
+            self.tablecombo.setVisible(False)
+            self.tabletext.setVisible(True)
         self.tablecombo.repaint()
 
     # Called on tablecombo selection change:
@@ -204,14 +218,16 @@ def select_table(self):
 
     def get_table(self):
         curIdx = self.tablecombo.currentIndex()
-        if self.data_source == self.TABLE and curIdx <= 0:
+        if self.data_source == self.TABLE and curIdx <= 0 and \
+                self.tabletext.text() == "":
             if self.database_desc:
                 self.database_desc["Table"] = "(None)"
             self.data_desc_table = None
             return None
 
         if self.data_source == self.TABLE:
-            self.table = self.tables[curIdx]
+            self.table = self.tables[curIdx] if curIdx > 0 else \
+                self.tabletext.text()
             self.database_desc["Table"] = self.table
             if "Query" in self.database_desc:
                 del self.database_desc["Query"]
@@ -291,7 +307,8 @@ def get_table(self):
                 confirm.setIcon(QMessageBox.Warning)
                 confirm.setText("Data appears to be big. Do you really "
                                 "want to download it to local memory?\n"
-                                "Table length: {:,}. Limit {:,}".format(table.approx_len(), MAX_DL_LIMIT))
+                                "Table length: {:,}. Limit {:,}".format(
+                    table.approx_len(), MAX_DL_LIMIT))
 
                 if table.approx_len() <= MAX_DL_LIMIT:
                     confirm.addButton("Yes", QMessageBox.YesRole)
diff --git a/Orange/widgets/data/tests/test_owsql.py b/Orange/widgets/data/tests/test_owsql.py
@@ -10,6 +10,13 @@
 from Orange.tests.sql.base import DataBaseTest as dbt
 
 
+def mock_sqltable(*args, **_):
+    table = Table(args[1])
+    table.get_domain = lambda **_: table.domain
+    table.download_data = lambda *_: 1
+    return table
+
+
 class TestOWSqlConnected(WidgetTest, dbt):
     def setUpDB(self):
         # pylint: disable=attribute-defined-outside-init
@@ -72,6 +79,7 @@ def test_restore_table(self, mock_backends, mock_sqltable):
         backend().display_name = "database"
         del backend().missing_extension
         backend().list_tables.return_value = ["a", "b", "c"]
+        backend().n_tables.return_value = 3
         mock_backends.available_backends.return_value = [backend]
         mock_sqltable().approx_len.return_value = 100
 
@@ -113,27 +121,55 @@ def test_data_source(self, mocked_backends: mock.Mock):
         backend = mock.Mock()
         backend().display_name = "Dummy Backend"
         backend().list_tables.return_value = ["a", "b", "c"]
+        backend().n_tables.return_value = 3
         mocked_backends.available_backends.return_value = [backend]
 
         settings = {"selected_backend": "Dummy Backend",
                     "host": "host", "port": "port", "database": "DB",
                     "schema": "", "username": "username",
-                    "password": "password", "data_source": OWSql.TABLE}
+                    "password": "password"}
         widget: OWSql = self.create_widget(OWSql, stored_settings=settings)
         self.assertEqual(widget.tablecombo.currentText(), "Select a table")
         self.assertFalse(widget.tablecombo.isHidden())
+        self.assertTrue(widget.tabletext.isHidden())
         self.assertTrue(widget.custom_sql.isHidden())
 
         widget.controls.data_source.buttons[OWSql.CUSTOM_SQL].click()
         self.assertEqual(widget.tablecombo.currentText(), "Select a table")
         self.assertFalse(widget.tablecombo.isHidden())
+        self.assertTrue(widget.tabletext.isHidden())
         self.assertFalse(widget.custom_sql.isHidden())
 
         widget.controls.data_source.buttons[OWSql.TABLE].click()
         self.assertEqual(widget.tablecombo.currentText(), "Select a table")
         self.assertFalse(widget.tablecombo.isHidden())
+        self.assertTrue(widget.tabletext.isHidden())
         self.assertTrue(widget.custom_sql.isHidden())
 
+    @mock.patch('Orange.widgets.data.owsql.MAX_TABLES', 2)
+    @mock.patch('Orange.widgets.data.owsql.SqlTable',
+                mock.Mock(side_effect=mock_sqltable))
+    @mock.patch('Orange.widgets.data.owsql.Backend')
+    def test_table_text(self, mocked_backends: mock.Mock):
+        backend = mock.Mock()
+        backend().display_name = "Dummy Backend"
+        backend().list_tables.return_value = ["iris", "zoo", "titanic"]
+        backend().n_tables.return_value = 3
+        mocked_backends.available_backends.return_value = [backend]
+
+        settings = {"selected_backend": "Dummy Backend",
+                    "host": "host", "port": "port", "database": "DB",
+                    "schema": "", "username": "username",
+                    "password": "password"}
+        widget: OWSql = self.create_widget(OWSql, stored_settings=settings)
+        self.assertTrue(widget.tablecombo.isHidden())
+        self.assertFalse(widget.tabletext.isHidden())
+        widget.tabletext.setText("zoo")
+        widget.select_table()
+        output = self.get_output(widget.Outputs.data, widget=widget)
+        self.assertIsInstance(output, Table)
+        self.assertEqual(len(output), 101)
+
 
 if __name__ == "__main__":
     unittest.main()
