Merge branch 'sqlalchemy-staging' into table-comment

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: Jesse Whitehouse

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Release History
 
+# 3.1.0 (TBD)
+
+- Fix: `server_hostname` URIs that included `https://` would raise an exception
+
 ## 3.0.1 (2023-12-01)
 
 - Other: updated docstring comment about default parameterization approach (#287)

src/databricks/__init__.py

@@ -1,4 +1,7 @@
-# https://packaging.python.org/guides/packaging-namespace-packages/#pkgutil-style-namespace-packages
-# This file should only contain the following line. Otherwise other sub-packages databricks.* namespace
-# may not be importable.
+# See: https://packaging.python.org/guides/packaging-namespace-packages/#pkgutil-style-namespace-packages
+#
+# This file must only contain the following line, or other packages in the databricks.* namespace
+# may not be importable. The contents of this file must be byte-for-byte equivalent across all packages.
+# If they are not, parallel package installation may lead to clobbered and invalid files.
+# Also see https://github.com/databricks/databricks-sdk-py/issues/343.
 __path__ = __import__("pkgutil").extend_path(__path__, __name__)

src/databricks/sql/thrift_backend.py

@@ -141,9 +141,11 @@ def __init__(
         if kwargs.get("_connection_uri"):
             uri = kwargs.get("_connection_uri")
         elif server_hostname and http_path:
-            uri = "https://{host}:{port}/{path}".format(
-                host=server_hostname, port=port, path=http_path.lstrip("/")
+            uri = "{host}:{port}/{path}".format(
+                host=server_hostname.rstrip("/"), port=port, path=http_path.lstrip("/")
             )
+            if not uri.startswith("https://"):
+                uri = "https://" + uri
         else:
             raise ValueError("No valid connection settings.")
 

src/databricks/sqlalchemy/_ddl.py

@@ -45,17 +45,40 @@ def visit_identity_column(self, identity, **kw):
         )
         return text
 
+    def visit_set_column_comment(self, create, **kw):
+        return "ALTER TABLE %s ALTER COLUMN %s COMMENT %s" % (
+            self.preparer.format_table(create.element.table),
+            self.preparer.format_column(create.element),
+            self.sql_compiler.render_literal_value(
+                create.element.comment, sqltypes.String()
+            ),
+        )
+
+    def visit_drop_column_comment(self, create, **kw):
+        return "ALTER TABLE %s ALTER COLUMN %s COMMENT ''" % (
+            self.preparer.format_table(create.element.table),
+            self.preparer.format_column(create.element),
+        )
+
     def get_column_specification(self, column, **kwargs):
-        """Currently we override this method only to emit a log message if a user attempts to set
-        autoincrement=True on a column. See comments in test_suite.py. We may implement implicit
-        IDENTITY using this feature in the future, similar to the Microsoft SQL Server dialect.
+        """
+        Emit a log message if a user attempts to set autoincrement=True on a column.
+        See comments in test_suite.py. We may implement implicit IDENTITY using this
+        feature in the future, similar to the Microsoft SQL Server dialect.
         """
         if column is column.table._autoincrement_column or column.autoincrement is True:
             logger.warning(
                 "Databricks dialect ignores SQLAlchemy's autoincrement semantics. Use explicit Identity() instead."
             )
 
-        return super().get_column_specification(column, **kwargs)
+        colspec = super().get_column_specification(column, **kwargs)
+        if column.comment is not None:
+            literal = self.sql_compiler.render_literal_value(
+                column.comment, sqltypes.STRINGTYPE
+            )
+            colspec += " COMMENT " + literal
+
+        return colspec
 
 
 class DatabricksStatementCompiler(compiler.SQLCompiler):

src/databricks/sqlalchemy/_parse.py

@@ -354,6 +354,7 @@ def parse_column_info_from_tgetcolumnsresponse(thrift_resp_row) -> ReflectedColu
         "type": final_col_type,
         "nullable": bool(thrift_resp_row.NULLABLE),
         "default": thrift_resp_row.COLUMN_DEF,
+        "comment": thrift_resp_row.REMARKS or None,
     }
 
     # TODO: figure out how to return sqlalchemy.interfaces in a way that mypy respects

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

@@ -18,6 +18,7 @@
 from sqlalchemy.engine import Engine
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column
+from sqlalchemy.schema import DropColumnComment, SetColumnComment
 from sqlalchemy.types import BOOLEAN, DECIMAL, Date, DateTime, Integer, String
 
 try:
@@ -188,6 +189,41 @@ def test_create_table_not_null(db_engine, metadata_obj: MetaData):
     metadata_obj.drop_all(db_engine)
 
 
+def test_column_comment(db_engine, metadata_obj: MetaData):
+    table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
+
+    column = Column("name", String(255), comment="some comment")
+    SampleTable = Table(table_name, metadata_obj, column)
+
+    metadata_obj.create_all(db_engine)
+    connection = db_engine.connect()
+
+    columns = db_engine.dialect.get_columns(
+        connection=connection, table_name=table_name
+    )
+
+    assert columns[0].get("comment") == "some comment"
+
+    column.comment = "other comment"
+    connection.execute(SetColumnComment(column))
+
+    columns = db_engine.dialect.get_columns(
+        connection=connection, table_name=table_name
+    )
+
+    assert columns[0].get("comment") == "other comment"
+
+    connection.execute(DropColumnComment(column))
+
+    columns = db_engine.dialect.get_columns(
+        connection=connection, table_name=table_name
+    )
+
+    assert columns[0].get("comment") == ""
+
+    metadata_obj.drop_all(db_engine)
+
+
 def test_bulk_insert_with_core(db_engine, metadata_obj, session):
     import random
 

src/databricks/sqlalchemy/test_local/test_ddl.py

@@ -0,0 +1,47 @@
+import pytest
+from sqlalchemy import Column, MetaData, String, Table, create_engine
+from sqlalchemy.schema import CreateTable, DropColumnComment, SetColumnComment
+
+
+class TestTableCommentDDL:
+    engine = create_engine(
+        "databricks://token:****@****?http_path=****&catalog=****&schema=****"
+    )
+
+    def compile(self, stmt):
+        return str(stmt.compile(bind=self.engine))
+
+    @pytest.fixture
+    def metadata(self) -> MetaData:
+        """Assemble a metadata object with one table containing one column."""
+        metadata = MetaData()
+
+        column = Column("foo", String, comment="bar")
+        table = Table("foobar", metadata, column)
+
+        return metadata
+
+    @pytest.fixture
+    def table(self, metadata) -> Table:
+        return metadata.tables.get("foobar")
+
+    @pytest.fixture
+    def column(self, table) -> Column:
+        return table.columns[0]
+
+    def test_create_table_with_column_comment(self, table):
+        stmt = CreateTable(table)
+        output = self.compile(stmt)
+
+        # output is a CREATE TABLE statement
+        assert "foo STRING COMMENT 'bar'" in output
+
+    def test_alter_table_add_column_comment(self, column):
+        stmt = SetColumnComment(column)
+        output = self.compile(stmt)
+        assert output == "ALTER TABLE foobar ALTER COLUMN foo COMMENT 'bar'"
+
+    def test_alter_table_drop_column_comment(self, column):
+        stmt = DropColumnComment(column)
+        output = self.compile(stmt)
+        assert output == "ALTER TABLE foobar ALTER COLUMN foo COMMENT ''"

tests/unit/test_init_file.py

@@ -0,0 +1,19 @@
+import hashlib
+
+
+class TestInitFile:
+    """
+    Micro test to confirm the contents of `databricks/__init__.py` does not change.
+
+    Also see https://github.com/databricks/databricks-sdk-py/issues/343#issuecomment-1866029118.
+    """
+
+    def test_init_file_contents(self):
+        with open("src/databricks/__init__.py") as f:
+            init_file_contents = f.read()
+
+        # This hash is the expected hash of the contents of `src/databricks/__init__.py`.
+        # It must not change, or else parallel package installation may lead to clobbered and invalid files.
+        expected_sha1 = "2772edbf52e517542acf8c039479c4b57b6ca2cd"
+        actual_sha1 = hashlib.sha1(init_file_contents.encode("utf-8")).hexdigest()
+        assert expected_sha1 == actual_sha1

tests/unit/test_thrift_backend.py

@@ -212,6 +212,18 @@ def test_port_and_host_are_respected(self, t_http_client_class):
         self.assertEqual(t_http_client_class.call_args[1]["uri_or_host"],
                          "https://hostname:123/path_value")
 
+    @patch("databricks.sql.auth.thrift_http_client.THttpClient")
+    def test_host_with_https_does_not_duplicate(self, t_http_client_class):
+        ThriftBackend("https://hostname", 123, "path_value", [], auth_provider=AuthProvider())
+        self.assertEqual(t_http_client_class.call_args[1]["uri_or_host"],
+                         "https://hostname:123/path_value")
+        
+    @patch("databricks.sql.auth.thrift_http_client.THttpClient")
+    def test_host_with_trailing_backslash_does_not_duplicate(self, t_http_client_class):
+        ThriftBackend("https://hostname/", 123, "path_value", [], auth_provider=AuthProvider())
+        self.assertEqual(t_http_client_class.call_args[1]["uri_or_host"],
+                         "https://hostname:123/path_value")        
+
     @patch("databricks.sql.auth.thrift_http_client.THttpClient")
     def test_socket_timeout_is_propagated(self, t_http_client_class):
         ThriftBackend("hostname", 123, "path_value", [], auth_provider=AuthProvider(), _socket_timeout=129)