fix sqlalchemy

Author: hantmac

cc_superset/engine.py

@@ -35,6 +35,7 @@ class DatabendEngineSpec(BaseEngineSpec, BasicParametersMixin):
     """
 
     engine = "databend"
+    engine_name = 'Databend Sqlalchemy'
 
     time_secondary_columns = True
     time_groupby_inline = True

databend_sqlalchemy/__init__.py

@@ -3,7 +3,5 @@
 
 from databend_sqlalchemy.entry_points import validate_entrypoints
 
-driver_name = 'clickhousedb'
-
-VERSION = (0, 1, 8)
+VERSION = (0, 2, 1)
 __version__ = '.'.join(str(x) for x in VERSION)

databend_sqlalchemy/databend_dialect.py

@@ -235,16 +235,6 @@ def connect(self, *cargs, **cparams):
         # inherits the docstring from interfaces.Dialect.connect
         return self.dbapi.connect(*cargs, **cparams)
 
-    def create_connect_args(self, url):
-        kwargs = {
-            'db_url': 'databend://%s:%d/' % (url.host, url.port or 8000),
-            'username': url.username,
-            'password': url.password,
-            'database': url.database,
-        }
-        kwargs.update(url.query)
-        return ([url.database or 'default'], kwargs)
-
     def _get_default_schema_name(self, connection):
         return connection.scalar("select currentDatabase()")
 

databend_sqlalchemy/entry_points.py

@@ -12,7 +12,7 @@ def validate_entrypoints():
     try:
         dist = pkg_resources.get_distribution('databend-sqlalchemy')
     except pkg_resources.DistributionNotFound:
-        print ('\nDatabend Sqlalchemy package not found in this Python installation')
+        print('\nDatabend Sqlalchemy package not found in this Python installation')
         return -1
     entry_map = dist.get_entry_map()
     print()
@@ -24,14 +24,14 @@ def validate_entrypoints():
             try:
                 expected_eps.remove(name)
             except KeyError:
-                print (f'\nUnexpected entry point {name} found')
+                print(f'\nUnexpected entry point {name} found')
                 return -1
     if expected_eps:
         print()
         for name in expected_eps:
-            print (f'Did not find expected ep {name}')
+            print(f'Did not find expected ep {name}')
         return -1
-    print ('\nEntrypoints correctly installed')
+    print('\nEntrypoints correctly installed')
     return 0
 
 

setup.py

@@ -37,8 +37,7 @@ def read_version():
     packages=find_packages('.', exclude=['tests*']),
     python_requires='>=3.4, <4',
     install_requires=[
-        'databend_py==0.3.5',
-        'mysql.connector',
+        'databend_py',
         'sqlalchemy',
     ],
 
@@ -84,5 +83,6 @@ def read_version():
 
     entry_points={
         'sqlalchemy.dialects': ['databend = databend_sqlalchemy.databend_dialect:DatabendDialect'],
+        'superset.db_engine_specs': ['databend = databend_sqlalchemy.cc_superset.engine:DatabendEngineSpec']
     },
 )

tests/integration/test_sqlalchemy_integration.py

@@ -3,7 +3,6 @@
 
 import pytest
 from sqlalchemy import create_engine, text
-from sqlalchemy.dialects import registry
 from sqlalchemy.engine.base import Connection, Engine
 from sqlalchemy.exc import OperationalError
 
@@ -12,13 +11,12 @@ class TestDatabendDialect:
     def test_set_params(
             self, username: str, password: str, database_name: str, host_port_name: str
     ):
-        print("sss")
-        registry.register("databend", __name__, "databend_sqlalchemy.databend_dialect.DatabendDialect")
         engine = create_engine(
             f"databend://{username}:{password}@{host_port_name}/{database_name}"
         )
         connection = engine.connect()
         result = connection.execute(text("SELECT 1"))
+        print(result)
         assert len(result.fetchall()) == 1
         engine.dispose()
 

tests/unit/test_databend_dialect.py

@@ -30,15 +30,15 @@ def test_create_dialect(self, dialect: DatabendDialect):
         assert isinstance(dialect.type_compiler, DatabendTypeCompiler)
         assert dialect.context == {}
 
-    def test_create_connect_args_service_account(self, dialect: DatabendDialect):
-        u = url.make_url(
-            "databend://user:pass@localhost:8000/testdb?secure=false"
-        )
-
-        result_list, result_dict = dialect.create_connect_args(u)
-        assert result_dict["db_url"] == "databend://localhost:8000/"
-        assert result_dict["database"] == "testdb"
-        assert result_dict["username"] == "user"
+    # def test_create_connect_args_service_account(self, dialect: DatabendDialect):
+    #     u = url.make_url(
+    #         "databend://user:pass@localhost:8000/testdb?secure=false"
+    #     )
+    #
+    #     result_list, result_dict = dialect.create_connect_args(u)
+    #     assert result_dict["db_url"] == "databend://localhost:8000/"
+    #     assert result_dict["database"] == "testdb"
+    #     assert result_dict["username"] == "user"
 
     def test_do_execute(
             self, dialect: DatabendDialect, cursor: mock.Mock(spec=MockCursor)