Skip to content

Commit ae37007

Browse files
cbornetJesse Whitehouse
cbornet
and
Jesse Whitehouse
authored
Add support for column comments (#306)
--------- Signed-off-by: Christophe Bornet <cbornet@hotmail.com> Signed-off-by: Jesse Whitehouse <jesse.whitehouse@databricks.com> Co-authored-by: Jesse Whitehouse <jesse.whitehouse@databricks.com>
1 parent 918752f commit ae37007

4 files changed

Lines changed: 112 additions & 5 deletions

File tree

‎src/databricks/sqlalchemy/_ddl.py‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import re
2-
from sqlalchemy.sql import compiler
2+
from sqlalchemy.sql import compiler, sqltypes
33
import logging
44

55
logger = logging.getLogger(__name__)
@@ -39,17 +39,40 @@ def visit_identity_column(self, identity, **kw):
3939
)
4040
return text
4141

42+
def visit_set_column_comment(self, create, **kw):
43+
return "ALTER TABLE %s ALTER COLUMN %s COMMENT %s" % (
44+
self.preparer.format_table(create.element.table),
45+
self.preparer.format_column(create.element),
46+
self.sql_compiler.render_literal_value(
47+
create.element.comment, sqltypes.String()
48+
),
49+
)
50+
51+
def visit_drop_column_comment(self, create, **kw):
52+
return "ALTER TABLE %s ALTER COLUMN %s COMMENT ''" % (
53+
self.preparer.format_table(create.element.table),
54+
self.preparer.format_column(create.element),
55+
)
56+
4257
def get_column_specification(self, column, **kwargs):
43-
"""Currently we override this method only to emit a log message if a user attempts to set
44-
autoincrement=True on a column. See comments in test_suite.py. We may implement implicit
45-
IDENTITY using this feature in the future, similar to the Microsoft SQL Server dialect.
58+
"""
59+
Emit a log message if a user attempts to set autoincrement=True on a column.
60+
See comments in test_suite.py. We may implement implicit IDENTITY using this
61+
feature in the future, similar to the Microsoft SQL Server dialect.
4662
"""
4763
if column is column.table._autoincrement_column or column.autoincrement is True:
4864
logger.warn(
4965
"Databricks dialect ignores SQLAlchemy's autoincrement semantics. Use explicit Identity() instead."
5066
)
5167

52-
return super().get_column_specification(column, **kwargs)
68+
colspec = super().get_column_specification(column, **kwargs)
69+
if column.comment is not None:
70+
literal = self.sql_compiler.render_literal_value(
71+
column.comment, sqltypes.STRINGTYPE
72+
)
73+
colspec += " COMMENT " + literal
74+
75+
return colspec
5376

5477

5578
class DatabricksStatementCompiler(compiler.SQLCompiler):

‎src/databricks/sqlalchemy/_parse.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,7 @@ def parse_column_info_from_tgetcolumnsresponse(thrift_resp_row) -> ReflectedColu
354354
"type": final_col_type,
355355
"nullable": bool(thrift_resp_row.NULLABLE),
356356
"default": thrift_resp_row.COLUMN_DEF,
357+
"comment": thrift_resp_row.REMARKS,
357358
}
358359

359360
# TODO: figure out how to return sqlalchemy.interfaces in a way that mypy respects

‎src/databricks/sqlalchemy/test_local/e2e/test_basic.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from sqlalchemy.engine import Engine
1919
from sqlalchemy.engine.reflection import Inspector
2020
from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column
21+
from sqlalchemy.schema import DropColumnComment, SetColumnComment
2122
from sqlalchemy.types import BOOLEAN, DECIMAL, Date, DateTime, Integer, String
2223

2324
try:
@@ -188,6 +189,41 @@ def test_create_table_not_null(db_engine, metadata_obj: MetaData):
188189
metadata_obj.drop_all(db_engine)
189190

190191

192+
def test_column_comment(db_engine, metadata_obj: MetaData):
193+
table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
194+
195+
column = Column("name", String(255), comment="some comment")
196+
SampleTable = Table(table_name, metadata_obj, column)
197+
198+
metadata_obj.create_all(db_engine)
199+
connection = db_engine.connect()
200+
201+
columns = db_engine.dialect.get_columns(
202+
connection=connection, table_name=table_name
203+
)
204+
205+
assert columns[0].get("comment") == "some comment"
206+
207+
column.comment = "other comment"
208+
connection.execute(SetColumnComment(column))
209+
210+
columns = db_engine.dialect.get_columns(
211+
connection=connection, table_name=table_name
212+
)
213+
214+
assert columns[0].get("comment") == "other comment"
215+
216+
connection.execute(DropColumnComment(column))
217+
218+
columns = db_engine.dialect.get_columns(
219+
connection=connection, table_name=table_name
220+
)
221+
222+
assert columns[0].get("comment") == ""
223+
224+
metadata_obj.drop_all(db_engine)
225+
226+
191227
def test_bulk_insert_with_core(db_engine, metadata_obj, session):
192228
import random
193229

‎src/databricks/sqlalchemy/test_local/test_ddl.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
import pytest
2+
from sqlalchemy import Column, MetaData, String, Table, create_engine
3+
from sqlalchemy.schema import CreateTable, DropColumnComment, SetColumnComment
4+
5+
6+
class TestTableCommentDDL:
7+
engine = create_engine(
8+
"databricks://token:****@****?http_path=****&catalog=****&schema=****"
9+
)
10+
11+
def compile(self, stmt):
12+
return str(stmt.compile(bind=self.engine))
13+
14+
@pytest.fixture
15+
def metadata(self) -> MetaData:
16+
"""Assemble a metadata object with one table containing one column."""
17+
metadata = MetaData()
18+
19+
column = Column("foo", String, comment="bar")
20+
table = Table("foobar", metadata, column)
21+
22+
return metadata
23+
24+
@pytest.fixture
25+
def table(self, metadata) -> Table:
26+
return metadata.tables.get("foobar")
27+
28+
@pytest.fixture
29+
def column(self, table) -> Column:
30+
return table.columns[0]
31+
32+
def test_create_table_with_column_comment(self, table):
33+
stmt = CreateTable(table)
34+
output = self.compile(stmt)
35+
36+
# output is a CREATE TABLE statement
37+
assert "foo STRING COMMENT 'bar'" in output
38+
39+
def test_alter_table_add_column_comment(self, column):
40+
stmt = SetColumnComment(column)
41+
output = self.compile(stmt)
42+
assert output == "ALTER TABLE foobar ALTER COLUMN foo COMMENT 'bar'"
43+
44+
def test_alter_table_drop_column_comment(self, column):
45+
stmt = DropColumnComment(column)
46+
output = self.compile(stmt)
47+
assert output == "ALTER TABLE foobar ALTER COLUMN foo COMMENT ''"

0 commit comments

Comments
 (0)