On SQL explorer page add a button to copy the create tables database schema

features/sql_explorer.py

@@ -3,7 +3,7 @@ class SQLExplorer:
         self.execute = db_connection_method
 
     def get_schema_info(self, schema='public'):
-        # Get all table names in the specified schema
+        # Get tables
         tables_result = self.execute("""
             SELECT table_name
             FROM information_schema.tables
@@ -18,10 +18,26 @@ class SQLExplorer:
             columns_result = self.execute("""
                 SELECT column_name, data_type
                 FROM information_schema.columns
-                WHERE table_schema = %s AND table_name = %s;
+                WHERE table_schema = %s AND table_name = %s
+                ORDER BY ordinal_position;
             """, [schema, table])
             columns = [(row['column_name'], row['data_type']) for row in columns_result]
 
+            # Get primary keys
+            # The constraint_type = 'PRIMARY KEY' check ensures we only get PK constraints
+            # This returns all columns that are part of the PK for this table.
+            primary_keys_result = self.execute("""
+                SELECT kcu.column_name
+                FROM information_schema.table_constraints tc
+                JOIN information_schema.key_column_usage kcu
+                    ON tc.constraint_name = kcu.constraint_name
+                AND tc.table_schema = kcu.table_schema
+                WHERE tc.constraint_type = 'PRIMARY KEY'
+                AND tc.table_schema = %s
+                AND tc.table_name = %s;
+            """, [schema, table])
+            primary_keys = [row['column_name'] for row in primary_keys_result]
+
             # Get foreign keys
             foreign_keys_result = self.execute("""
                 SELECT
@@ -31,15 +47,15 @@ class SQLExplorer:
                 FROM
                     information_schema.table_constraints AS tc
                     JOIN information_schema.key_column_usage AS kcu
-                    ON tc.constraint_name = kcu.constraint_name
-                    AND tc.table_schema = kcu.table_schema
+                        ON tc.constraint_name = kcu.constraint_name
+                        AND tc.table_schema = kcu.table_schema
                     JOIN information_schema.constraint_column_usage AS ccu
-                    ON ccu.constraint_name = tc.constraint_name
-                    AND ccu.table_schema = tc.table_schema
+                        ON ccu.constraint_name = tc.constraint_name
+                        AND ccu.table_schema = tc.table_schema
                 WHERE
-                    tc.constraint_type = 'FOREIGN KEY' AND
-                    tc.table_schema = %s AND
-                    tc.table_name = %s;
+                    tc.constraint_type = 'FOREIGN KEY'
+                    AND tc.table_schema = %s
+                    AND tc.table_name = %s;
             """, [schema, table])
             foreign_keys = [
                 (row['fk_column'], row['referenced_table'], row['referenced_column'])
@@ -48,10 +64,25 @@ class SQLExplorer:
 
             schema_info[table] = {
                 'columns': columns,
+                'primary_keys': primary_keys,
                 'foreign_keys': foreign_keys
             }
 
         return schema_info
+
+    def map_data_type_for_sql(self, postgres_type):
+        # This is naive. For real usage, you may handle numeric precision, etc.
+        # Or simply return the raw type since your DB is PostgreSQL anyway.
+        return {
+            'character varying': 'VARCHAR',
+            'varchar': 'VARCHAR',
+            'text': 'TEXT',
+            'integer': 'INTEGER',
+            'bigint': 'BIGINT',
+            'boolean': 'BOOLEAN',
+            'timestamp without time zone': 'TIMESTAMP',
+            'timestamp with time zone': 'TIMESTAMPTZ',
+        }.get(postgres_type, postgres_type.upper())
     
     def map_data_type(self, postgres_type):
         type_mapping = {
@@ -92,6 +123,42 @@ class SQLExplorer:
         
         return "\n".join(mermaid_lines)
     
+    def generate_create_script(self, schema_info):
+        lines = []
+
+        for table, info in schema_info.items():
+            columns = info['columns']
+            pks = info.get('primary_keys', [])
+            fks = info['foreign_keys']
+
+            column_defs = []
+            for column_name, data_type in columns:
+                sql_type = self.map_data_type_for_sql(data_type)
+                column_defs.append(f'  "{column_name}" {sql_type}')
+
+            if pks:
+                pk_columns = ", ".join(f'"{pk}"' for pk in pks)
+                column_defs.append(f'  PRIMARY KEY ({pk_columns})')
+
+            create_stmt = 'CREATE TABLE "{}" (\n'.format(table)
+            create_stmt += ",\n".join(column_defs)
+            create_stmt += '\n);'
+            lines.append(create_stmt)
+
+            # Foreign keys
+            for fk_column, ref_table, ref_col in fks:
+                alter_stmt = (
+                    f'ALTER TABLE "{table}" '
+                    f'ADD CONSTRAINT "fk_{table}_{fk_column}" '
+                    f'FOREIGN KEY ("{fk_column}") '
+                    f'REFERENCES "{ref_table}" ("{ref_col}");'
+                )
+                lines.append(alter_stmt)
+
+            lines.append("")  # separate blocks
+
+        return "\n".join(lines)
+    
     def execute_sql(self, query):
         results = None
         columns = []