Describe outbound foreign keys in SQL Cell \describe output. (#67)

Author: James Robinson

noteable_magics/sql/meta_commands.py

@@ -455,11 +455,12 @@ 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 and transform to a dataframe for
-            # presentation.
-            index_df = index_dataframe(inspector, relation_name, schema)
-            if len(index_df):
-                display(secondary_dataframe_to_html(index_df))
+            # Is a table. Let's go get indices, foreign keys. If meaningful dataframe returned, 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):
+                secondary_df = secondary_function(inspector, relation_name, schema)
+                if len(secondary_df):
+                    display(secondary_dataframe_to_html(secondary_df))
 
         return main_relation_df, False
 
@@ -477,6 +478,44 @@ def _split_schema_table(schema_table: str) -> Tuple[Optional[str], str]:
         return (schema, table)
 
 
+def foreignkeys_dataframe(
+    inspector: SchemaStrippingInspector, table_name: str, schema: Optional[str]
+) -> DataFrame:
+    """Transform results from inspector.get_indexes() into a single dataframe for display() purposes"""
+
+    names: List[str] = []  # Will be '(unnamed)' if the constraint was not named
+    constrained_columns: List[str] = []  # Will be comma separated list for compound FKs
+    referenced_qualified_tables: List[str] = []
+    referenced_columns: List[str] = []  # Will be comma separated list for compound FKs
+
+    fkey_dicts = inspector.get_foreign_keys(table_name, schema)
+
+    for fk_dict in fkey_dicts:
+        if fk_dict['referred_schema']:
+            # Schema qualify the table.
+            referred_table = f"{fk_dict['referred_schema']}.{fk_dict['referred_table']}"
+        else:
+            referred_table = fk_dict['referred_table']
+
+        referenced_qualified_tables.append(referred_table)
+        names.append(fk_dict.get('name', '(unnamed)'))
+        constrained_columns.append(', '.join(fk_dict.get('constrained_columns')))
+        referenced_columns.append(', '.join(fk_dict.get('referred_columns')))
+
+    df = DataFrame(
+        {
+            'Foreign Key': names,
+            'Columns': constrained_columns,
+            'Referenced Table': referenced_qualified_tables,
+            'Referenced Columns': referenced_columns,
+        }
+    )
+
+    title = f'Table <code>{displayable_relation_name(schema, table_name)}</code> Foreign Keys'
+
+    return set_dataframe_metadata(df, title=title)
+
+
 def index_dataframe(
     inspector: SchemaStrippingInspector, table_name: str, schema: Optional[str]
 ) -> DataFrame:
@@ -694,6 +733,9 @@ def get_view_definition(self, view_name: str, schema: Optional[str] = None) -> s
     def get_pk_constraint(self, table_name: str, schema: Optional[str] = None) -> dict:
         return self.underlying_inspector.get_pk_constraint(table_name, schema=schema)
 
+    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_indexes(self, table_name: str, schema: Optional[str] = None) -> List[dict]:
         return self.underlying_inspector.get_indexes(table_name, schema=schema)
 

tests/conftest.py

@@ -200,10 +200,14 @@ def cleanup_any_extra_tables(connection: Connection):
 
     unexpected_relations = relations - KNOWN_TABLES
 
+    # CRDB needs 'cascade' to be able to drop tables referenced with FKs. SQLite does
+    # not recognize, however.
+    maybe_cascade = 'cascade' if 'cockroach' in str(connection._engine) else ''
+
     for unexpected_relation in unexpected_relations:
         try:
             with Session(connection._engine) as db:
-                db.execute(f'drop table {unexpected_relation}')
+                db.execute(f'drop table {unexpected_relation} {maybe_cascade}')
                 db.commit()
         except Exception:
             # Maybe it was a view?

tests/test_sql_magic_meta_commands.py

@@ -424,6 +424,68 @@ def test_against_schema_qualified_view(self, sql_magic, ipython_namespace, mock_
             '<br />\n<h2>View <code>public.str_int_view</code> Definition</h2>'
         ), html_contents
 
+    @pytest.mark.parametrize(
+        'handle,schema', [('@cockroach', ''), ('@cockroach', 'public'), ('@sqlite', '')]
+    )
+    def test_foreign_keys(self, sql_magic, ipython_namespace, mock_display, handle, schema):
+        """Describing table `references_int_table` should talk about a foreign key over to int_table"""
+
+        # If was asked with schema qualification, then various outputs will also be schema qualified.
+        qualified_references_int_table = (
+            f'{schema}.references_int_table' if schema else 'references_int_table'
+        )
+        qualified_int_table = f'{schema}.int_table' if schema else 'int_table'
+
+        sql_magic.execute(rf'{handle} \d {qualified_references_int_table}')
+
+        assert (
+            len(mock_display.call_args_list) == 3
+        )  # main structure DF, index DF-as-html, foreign key DF-as-html.
+        fk_html = mock_display.call_args_list[2].args[0]
+        assert isinstance(fk_html, HTML)
+        html_contents: str = fk_html.data
+
+        assert html_contents.startswith(
+            f'<br />\n<h2>Table <code>{qualified_references_int_table}</code> Foreign Keys</h2>'
+        ), html_contents
+
+        # Convert the HTML table back to dataframe to complete test.
+        fk_df = pd.read_html(html_contents)[0]
+
+        assert fk_df.columns.tolist() == [
+            'Foreign Key',
+            'Columns',
+            'Referenced Table',
+            'Referenced Columns',
+        ]
+        assert fk_df['Columns'].tolist() == ['a_id']
+        assert fk_df['Referenced Table'].tolist() == [qualified_int_table]
+        assert fk_df['Referenced Columns'].tolist() == ['a']
+
+        # Also test against a table with a compound foreign key. Must create table pair ad hoc. Will be cleaned
+        # up upon test cleanup.
+
+        sql_magic.execute(
+            f'{handle}\ncreate table {qualified_int_table}_2 (a int, b int, primary key(a, b))'
+        )
+
+        sql_magic.execute(
+            f'{handle}\ncreate table {qualified_references_int_table}_2 (a_ref int primary key, b_ref int, constraint a_b_fk foreign key (a_ref, b_ref) references {qualified_int_table}_2(a, b))'
+        )
+
+        mock_display.reset_mock()
+
+        sql_magic.execute(fr'{handle} \describe {qualified_references_int_table}_2')
+
+        assert (
+            len(mock_display.call_args_list) == 3
+        )  # main structure DF, index DF-as-html, foreign key DF-as-html.
+        fk2_html = mock_display.call_args_list[2].args[0]
+
+        fk2_df = pd.read_html(fk2_html.data)[0]
+        assert fk2_df['Columns'][0] == 'a_ref, b_ref'
+        assert fk2_df['Referenced Columns'][0] == 'a, b'
+
     def test_against_table_without_a_primary_key(self, sql_magic, ipython_namespace, mock_display):
         # 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)