Add Columns endpoint

db/columns.py

@@ -1,7 +1,17 @@
-from sqlalchemy import Column, Integer, ForeignKey, Table, DDL, MetaData
-from db import constants
+import logging
+import warnings
+from alembic.migration import MigrationContext
+from alembic.operations import Operations
+from sqlalchemy import (
+    Column, Integer, ForeignKey, Table, MetaData, and_, select
+)
+from db import constants, tables
+from db.types import alteration
 
+logger = logging.getLogger(__name__)
 
+
+NAME = "name"
 NULLABLE = "nullable"
 PRIMARY_KEY = "primary_key"
 TYPE = "type"
@@ -77,8 +87,7 @@ def get_default_mathesar_column_list():
     return [
         MathesarColumn(
             c,
-            DEFAULT_COLUMNS[c][TYPE],
-            primary_key=DEFAULT_COLUMNS[c][PRIMARY_KEY]
+            DEFAULT_COLUMNS[c][TYPE], primary_key=DEFAULT_COLUMNS[c][PRIMARY_KEY]
         )
         for c in DEFAULT_COLUMNS
     ]
@@ -90,17 +99,110 @@ def init_mathesar_table_column_list_with_defaults(column_list):
     return default_columns + given_columns
 
 
-def rename_column(schema, table_name, column_name, new_column_name, engine):
-    _preparer = engine.dialect.identifier_preparer
+def get_column_index_from_name(table_oid, column_name, engine):
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", message="Did not recognize type")
+        pg_attribute = Table("pg_attribute", MetaData(), autoload_with=engine)
+    sel = select(pg_attribute.c.attnum).where(
+        and_(
+            pg_attribute.c.attrelid == table_oid,
+            pg_attribute.c.attname == column_name
+        )
+    )
     with engine.begin() as conn:
-        metadata = MetaData(bind=engine, schema=schema)
-        table = Table(table_name, metadata, schema=schema, autoload_with=engine)
-        column = table.columns[column_name]
-        prepared_table_name = _preparer.format_table(table)
-        prepared_column_name = _preparer.format_column(column)
-        prepared_new_column_name = _preparer.quote(new_column_name)
-        alter_stmt = f"""
-        ALTER TABLE {prepared_table_name}
-        RENAME {prepared_column_name} TO {prepared_new_column_name}
-        """
-        conn.execute(DDL(alter_stmt))
+        result = conn.execute(sel).fetchone()[0]
+    return result - 1
+
+
+def create_column(engine, table_oid, column_data):
+    column_type = column_data[TYPE]
+    column_nullable = column_data.get(NULLABLE, True)
+    supported_types = alteration.get_supported_alter_column_types(engine)
+    sa_type = supported_types.get(column_type.lower())
+    if sa_type is None:
+        logger.warning("Requested type not supported. falling back to String")
+        sa_type = supported_types[alteration.STRING]
+    table = tables.reflect_table_from_oid(table_oid, engine)
+    column = MathesarColumn(
+        column_data[NAME], sa_type, nullable=column_nullable,
+    )
+    with engine.begin() as conn:
+        ctx = MigrationContext.configure(conn)
+        op = Operations(ctx)
+        op.add_column(table.name, column, schema=table.schema)
+    return tables.reflect_table_from_oid(table_oid, engine).columns[column_data[NAME]]
+
+
+def alter_column(
+        engine,
+        table_oid,
+        column_index,
+        column_definition_dict,
+):
+    assert len(column_definition_dict) == 1
+    column_def_key = list(column_definition_dict.keys())[0]
+    column_index = int(column_index)
+    attribute_alter_map = {
+        NAME: rename_column,
+        TYPE: retype_column,
+        NULLABLE: change_column_nullable,
+    }
+    return attribute_alter_map[column_def_key](
+        table_oid, column_index, column_definition_dict[column_def_key], engine,
+    )
+
+
+def rename_column(table_oid, column_index, new_column_name, engine):
+    table = tables.reflect_table_from_oid(table_oid, engine)
+    column = table.columns[column_index]
+    with engine.begin() as conn:
+        ctx = MigrationContext.configure(conn)
+        op = Operations(ctx)
+        op.alter_column(
+            table.name,
+            column.name,
+            new_column_name=new_column_name,
+            schema=table.schema
+        )
+    return tables.reflect_table_from_oid(table_oid, engine).columns[column_index]
+
+
+def retype_column(table_oid, column_index, new_type, engine):
+    table = tables.reflect_table_from_oid(table_oid, engine)
+    alteration.alter_column_type(
+        table.schema,
+        table.name,
+        table.columns[column_index].name,
+        new_type,
+        engine,
+    )
+    return tables.reflect_table_from_oid(table_oid, engine).columns[column_index]
+
+
+def change_column_nullable(table_oid, column_index, nullable, engine):
+    table = tables.reflect_table_from_oid(table_oid, engine)
+    column = table.columns[column_index]
+    with engine.begin() as conn:
+        ctx = MigrationContext.configure(conn)
+        op = Operations(ctx)
+        op.alter_column(
+            table.name,
+            column.name,
+            nullable=nullable,
+            schema=table.schema
+        )
+    return tables.reflect_table_from_oid(table_oid, engine).columns[column_index]
+
+
+def drop_column(
+        engine,
+        table_oid,
+        column_index,
+):
+    column_index = int(column_index)
+    table = tables.reflect_table_from_oid(table_oid, engine)
+    column = table.columns[column_index]
+    with engine.begin() as conn:
+        ctx = MigrationContext.configure(conn)
+        op = Operations(ctx)
+        op.drop_column(table.name, column.name, schema=table.schema)

db/tests/test_columns.py

@@ -1,6 +1,11 @@
 import re
+from unittest.mock import patch
+from psycopg2.errors import NotNullViolation
 import pytest
-from sqlalchemy import String, Integer, ForeignKey, Column, select, Table, MetaData
+from sqlalchemy import (
+    String, Integer, ForeignKey, Column, select, Table, MetaData, create_engine
+)
+from sqlalchemy.exc import IntegrityError
 from db import columns, tables, constants
 
 
@@ -22,7 +27,9 @@ def _rename_column(schema, table_name, old_col_name, new_col_name, engine):
     """
     Renames the colum of a table and assert the change went through
     """
-    columns.rename_column(schema, table_name, old_col_name, new_col_name, engine)
+    table_oid = tables.get_oid_from_table(table_name, schema, engine)
+    column_index = columns.get_column_index_from_name(table_oid, old_col_name, engine)
+    columns.rename_column(table_oid, column_index, new_col_name, engine)
     table = tables.reflect_table(table_name, schema, engine)
     assert new_col_name in table.columns
     assert old_col_name not in table.columns
@@ -137,6 +144,31 @@ def test_MC_is_default_when_false_for_pk():
         assert not col.is_default
 
 
+@pytest.mark.parametrize(
+    "column_dict,func_name",
+    [
+        ({"name": "blah"}, "rename_column"),
+        ({"type": "blah"}, "retype_column"),
+        ({"nullable": True}, "change_column_nullable"),
+    ]
+)
+def test_alter_column_chooses_wisely(column_dict, func_name):
+    engine = create_engine("postgresql://")
+    with patch.object(columns, func_name) as mock_alterer:
+        columns.alter_column(
+            engine,
+            1234,
+            5678,
+            column_dict
+        )
+    mock_alterer.assert_called_with(
+        1234,
+        5678,
+        list(column_dict.values())[0],
+        engine,
+    )
+
+
 def test_rename_column(engine_with_schema):
     old_col_name = "col1"
     new_col_name = "col2"
@@ -202,6 +234,175 @@ def test_rename_column_index(engine_with_schema):
     assert new_col_name in index_columns
 
 
+def test_get_column_index_from_name(engine_with_schema):
+    engine, schema = engine_with_schema
+    table_name = "table_with_columns"
+    zero_name = "colzero"
+    one_name = "colone"
+    table = Table(
+        table_name,
+        MetaData(bind=engine, schema=schema),
+        Column(zero_name, Integer),
+        Column(one_name, String),
+    )
+    table.create()
+    table_oid = tables.get_oid_from_table(table_name, schema, engine)
+    assert columns.get_column_index_from_name(table_oid, zero_name, engine) == 0
+    assert columns.get_column_index_from_name(table_oid, one_name, engine) == 1
+
+
+def test_retype_column_correct_column(engine_with_schema):
+    engine, schema = engine_with_schema
+    table_name = "atableone"
+    target_type = "boolean"
+    target_column_name = "thecolumntochange"
+    nontarget_column_name = "notthecolumntochange"
+    table = Table(
+        table_name,
+        MetaData(bind=engine, schema=schema),
+        Column(target_column_name, Integer),
+        Column(nontarget_column_name, String),
+    )
+    table.create()
+    table_oid = tables.get_oid_from_table(table_name, schema, engine)
+    with patch.object(columns.alteration, "alter_column_type") as mock_retyper:
+        columns.retype_column(table_oid, 0, target_type, engine)
+    mock_retyper.assert_called_with(
+        schema,
+        table_name,
+        target_column_name,
+        "boolean",
+        engine
+    )
+
+
+def test_create_column(engine_with_schema):
+    engine, schema = engine_with_schema
+    table_name = "atableone"
+    target_type = "boolean"
+    initial_column_name = "original_column"
+    new_column_name = "added_column"
+    table = Table(
+        table_name,
+        MetaData(bind=engine, schema=schema),
+        Column(initial_column_name, Integer),
+    )
+    table.create()
+    table_oid = tables.get_oid_from_table(table_name, schema, engine)
+    column_data = {"name": new_column_name, "type": target_type}
+    created_col = columns.create_column(engine, table_oid, column_data)
+    altered_table = tables.reflect_table_from_oid(table_oid, engine)
+    assert len(altered_table.columns) == 2
+    assert created_col.name == new_column_name
+    assert created_col.type.compile(engine.dialect) == "BOOLEAN"
+
+
+nullable_changes = [(True, True), (False, False), (True, False), (False, True)]
+
+
+@pytest.mark.parametrize("nullable_tup", nullable_changes)
+def test_change_column_nullable_changes(engine_with_schema, nullable_tup):
+    engine, schema = engine_with_schema
+    table_name = "atablefornulling"
+    target_column_name = "thecolumntochange"
+    nontarget_column_name = "notthecolumntochange"
+    table = Table(
+        table_name,
+        MetaData(bind=engine, schema=schema),
+        Column(target_column_name, Integer, nullable=nullable_tup[0]),
+        Column(nontarget_column_name, String),
+    )
+    table.create()
+    table_oid = tables.get_oid_from_table(table_name, schema, engine)
+    changed_column = columns.change_column_nullable(
+        table_oid,
+        0,
+        nullable_tup[1],
+        engine
+    )
+    assert changed_column.nullable is nullable_tup[1]
+
+
+@pytest.mark.parametrize("nullable_tup", nullable_changes)
+def test_change_column_nullable_with_data(engine_with_schema, nullable_tup):
+    engine, schema = engine_with_schema
+    table_name = "atablefornulling"
+    target_column_name = "thecolumntochange"
+    table = Table(
+        table_name,
+        MetaData(bind=engine, schema=schema),
+        Column(target_column_name, Integer, nullable=nullable_tup[0]),
+    )
+    table.create()
+    ins = table.insert().values(
+        [
+            {target_column_name: 1},
+            {target_column_name: 2},
+            {target_column_name: 3},
+        ]
+    )
+    with engine.begin() as conn:
+        conn.execute(ins)
+    table_oid = tables.get_oid_from_table(table_name, schema, engine)
+    changed_column = columns.change_column_nullable(
+        table_oid,
+        0,
+        nullable_tup[1],
+        engine
+    )
+    assert changed_column.nullable is nullable_tup[1]
+
+
+def test_change_column_nullable_changes_raises_with_null_data(engine_with_schema):
+    engine, schema = engine_with_schema
+    table_name = "atablefornulling"
+    target_column_name = "thecolumntochange"
+    table = Table(
+        table_name,
+        MetaData(bind=engine, schema=schema),
+        Column(target_column_name, Integer, nullable=True),
+    )
+    table.create()
+    ins = table.insert().values(
+        [
+            {target_column_name: 1},
+            {target_column_name: 2},
+            {target_column_name: None},
+        ]
+    )
+    with engine.begin() as conn:
+        conn.execute(ins)
+    table_oid = tables.get_oid_from_table(table_name, schema, engine)
+    with pytest.raises(IntegrityError) as e:
+        columns.change_column_nullable(
+            table_oid,
+            0,
+            False,
+            engine
+        )
+        assert type(e.orig) == NotNullViolation
+
+
+def test_drop_column_correct_column(engine_with_schema):
+    engine, schema = engine_with_schema
+    table_name = "atable"
+    target_column_name = "thecolumntodrop"
+    nontarget_column_name = "notthecolumntodrop"
+    table = Table(
+        table_name,
+        MetaData(bind=engine, schema=schema),
+        Column(target_column_name, Integer),
+        Column(nontarget_column_name, String),
+    )
+    table.create()
+    table_oid = tables.get_oid_from_table(table_name, schema, engine)
+    columns.drop_column(engine, table_oid, 0)
+    altered_table = tables.reflect_table_from_oid(table_oid, engine)
+    assert len(altered_table.columns) == 1
+    assert nontarget_column_name in altered_table.columns
+    assert target_column_name not in altered_table.columns
+
+
 def get_mathesar_column_init_args():
     init_code = columns.MathesarColumn.__init__.__code__
     return init_code.co_varnames[1:init_code.co_argcount]