Eng 5335 sql cell describe to enumerate constraints (#68)

Author: James Robinson

noteable_magics/sql/meta_commands.py

@@ -455,9 +455,14 @@ def run(self, invoked_as: str, args: List[str]) -> Tuple[DataFrame, bool]:
 
                 display(HTML('\n'.join(html_buf)))
         else:
-            # Is a table. Let's go get indices, foreign keys. If meaningful dataframe returned, transform to
+            # Is a table. Let's go get indices, foreign keys, other table constraints.
+            # If meaningful dataframe returned for any of these, transform to
             # HTML for presentation (DEX only expects at most a single DF display()ed per cell) and display it.
-            for secondary_function in (index_dataframe, foreignkeys_dataframe):
+            for secondary_function in (
+                index_dataframe,
+                foreignkeys_dataframe,
+                constraints_dataframe,
+            ):
                 secondary_df = secondary_function(inspector, relation_name, schema)
                 if len(secondary_df):
                     display(secondary_dataframe_to_html(secondary_df))
@@ -478,6 +483,32 @@ def _split_schema_table(schema_table: str) -> Tuple[Optional[str], str]:
         return (schema, table)
 
 
+def constraints_dataframe(
+    inspector: SchemaStrippingInspector, table_name: str, schema: Optional[str]
+) -> DataFrame:
+    """Transform results from inspector.get_check_constraints() into a single dataframe for display() purposes"""
+
+    names: List[str] = []
+    definitions: List[str] = []
+
+    constraint_dicts: List[dict] = inspector.get_check_constraints(table_name, schema)
+
+    for constraint_dict in sorted(constraint_dicts, key=lambda d: d['name']):
+        names.append(constraint_dict['name'])
+        definitions.append(constraint_dict['sqltext'])
+
+    df = DataFrame(
+        {
+            'Constraint': names,
+            'Definition': definitions,
+        }
+    )
+
+    title = f'Table <code>{displayable_relation_name(schema, table_name)}</code> Check Constraints'
+
+    return set_dataframe_metadata(df, title=title)
+
+
 def foreignkeys_dataframe(
     inspector: SchemaStrippingInspector, table_name: str, schema: Optional[str]
 ) -> DataFrame:
@@ -736,6 +767,9 @@ def get_pk_constraint(self, table_name: str, schema: Optional[str] = None) -> di
     def get_foreign_keys(self, table_name: str, schema: Optional[str] = None) -> List[dict]:
         return self.underlying_inspector.get_foreign_keys(table_name, schema=schema)
 
+    def get_check_constraints(self, table_name: str, schema: Optional[str] = None) -> List[dict]:
+        return self.underlying_inspector.get_check_constraints(table_name, schema=schema)
+
     def get_indexes(self, table_name: str, schema: Optional[str] = None) -> List[dict]:
         return self.underlying_inspector.get_indexes(table_name, schema=schema)
 

tests/conftest.py

@@ -158,10 +158,17 @@ def populate_database(connection: Connection, include_comments=False):
         db.execute('insert into int_table (a, b, c) values (1, 2, 3), (4, 5, 6)')
 
         db.execute(
-            "create table str_table(str_id text not null default 'foonly', int_col int default 22)"
+            """create table str_table(
+                str_id text not null default 'f'
+                    constraint single_char_str_id check (length(str_id) = 1),
+                int_col int default 22
+                    constraint only_even_int_col_values check (int_col % 2 = 0),
+                constraint never_f_10 check (not (str_id = 'f' and int_col = 10))
+             )
+            """
         )
         db.execute(
-            "insert into str_table(str_id, int_col) values ('a', 1), ('b', 2), ('c', 3), ('d', null)"
+            "insert into str_table(str_id, int_col) values ('a', 2), ('b', 2), ('c', 4), ('d', null)"
         )
 
         db.execute(

tests/test_sql_magic.py

@@ -233,7 +233,7 @@ def test_against_string(self, sql_magic, ipython_shell):
         )
 
         # Scalar result.
-        assert results == 1
+        assert results == 2
 
     @pytest.mark.parametrize('ret_col,expected_value', [('a', 1), ('b', 2)])
     def test_sqlsafe(self, sql_magic, ipython_shell, ret_col, expected_value):

tests/test_sql_magic_meta_commands.py

@@ -490,15 +490,68 @@ def test_against_table_without_a_primary_key(self, sql_magic, ipython_namespace,
         # str_table on sqlite will not have any primary key or any indices at all
         # (all tables in cockroach have an implicit PK, so can't test with it)
 
-        # So the only output should be the dataframe. Not a subsequent HTML blob describing indices.
+        # The output should NOT include an HTML blob describing indices.
         sql_magic.execute(r'@sqlite \d str_table')
 
-        assert len(mock_display.call_args_list) == 1
+        assert len(mock_display.call_args_list) == 2  # main df, constraints df-as-html
 
         df = mock_display.call_args_list[0].args[0]
+        assert isinstance(df, pd.DataFrame)
         assert df.attrs['noteable']['decoration']['title'] == 'Table "str_table" Structure'
 
-        assert isinstance(df, pd.DataFrame)
+        # Test test_constraints() will exercise this further. Only mention it here
+        # because will be returned and should not be talking about primary key / indices.
+        constraint_html = mock_display.call_args_list[1].args[0].data
+        assert constraint_html.startswith(
+            '<br />\n<h2>Table <code>str_table</code> Check Constraints</h2>'
+        )
+
+    # All CRDB tables have a primary key, so conditionally expect it to be described.
+    @pytest.mark.parametrize(
+        'handle,expected_display_callcount', [('@cockroach', 3), ('@sqlite', 2)]
+    )
+    def test_constraints(
+        self, handle, expected_display_callcount, sql_magic, ipython_namespace, mock_display
+    ):
+
+        sql_magic.execute(rf'{handle} \d str_table')
+
+        assert (
+            len(mock_display.call_args_list) == expected_display_callcount
+        )  # main df, maybe index html, constraints df-as-html
+
+        # The constraints HTML blob will be the final one always.
+        constraint_html = mock_display.call_args_list[-1].args[0].data
+        assert constraint_html.startswith(
+            '<br />\n<h2>Table <code>str_table</code> Check Constraints</h2>'
+        )
+
+        # Convert back to dataframe
+        constraint_df = pd.read_html(constraint_html)[0]
+
+        assert len(constraint_df) == 3  # Three check constraints on this table
+
+        assert constraint_df.columns.tolist() == [
+            'Constraint',
+            'Definition',
+        ]
+
+        # Should be alpha sorted by constraint name.
+        assert constraint_df['Constraint'].tolist() == [
+            'never_f_10',
+            'only_even_int_col_values',
+            'single_char_str_id',
+        ]
+
+        # The SQL dialects convert the constraint expressions back to strings with slightly
+        # varying spellings (as expected), so can't simply blindly assert all of them.
+        constraint_definitions = constraint_df['Definition'].tolist()
+
+        # This one happens to be regurgitated consistently between sqlite and CRDB.
+        assert 'length(str_id) = 1' in constraint_definitions
+        # Little gentler substring matching for the other two.
+        assert any("str_id = 'f'" in cd for cd in constraint_definitions)
+        assert any("int_col % 2" in cd for cd in constraint_definitions)
 
     def test_no_args_gets_table_list(self, sql_magic, ipython_namespace):
         sql_magic.execute(r'@sqlite \d')