Simple integration test examples.

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

Author: susodapop

src/databricks/sqlalchemy/dialect.py

@@ -1,17 +1,15 @@
 from databricks import sql
 from typing import AnyStr
 
-from sqlalchemy import types
-from sqlalchemy import util
-
+from sqlalchemy import util, exc, types
 from sqlalchemy.engine import default
 
 
 class DatabricksDialect(default.DefaultDialect):
 
     # Possible attributes are defined here: https://docs.sqlalchemy.org/en/14/core/internals.html#sqlalchemy.engine.Dialect
     name: str = "databricks"
-    driver: str= "thrift"
+    driver: str = "thrift"
     default_schema_name: str = "default"
 
     @classmethod
@@ -24,7 +22,7 @@ def create_connect_args(self, url):
         kwargs = {
             "server_hostname": url.host,
             "access_token": url.password,
-            "http_path": url.query.get("http_path")
+            "http_path": url.query.get("http_path"),
         }
 
         return [], kwargs
@@ -43,3 +41,19 @@ def get_columns(self, *args, **kwargs):
     def do_rollback(self, dbapi_connection):
         # Databricks SQL Does not support transaction
         pass
+
+    def has_table(self, connection, table_name, schema=None, **kwargs) -> bool:
+        """Required for `tests.sqlalchemy.integration.test_create_table` to pass.
+        """
+        try:
+            COLUMN_NAME = 3
+            with self.get_driver_connection(
+                connection
+            )._dbapi_connection.dbapi_connection.cursor() as cur:
+                data = cur.columns(
+                    schema_name=schema or "default", table_name=table_name
+                ).fetchmany(1)
+            # the table exists as long as there's a non-zero number of columns
+            return len(data) > 0
+        except exc.NoSuchTableError:
+            return False

tests/sqlalchemy/conftest.py

@@ -0,0 +1,46 @@
+import os, datetime, pytest
+import sqlalchemy
+
+
+@pytest.fixture
+def db_engine():
+
+    host = os.getenv("host")
+    http_path = os.getenv("http_path")
+    api_token = os.getenv("api_token")
+    engine = sqlalchemy.create_engine(f"databricks+thrift://token:{api_token}@{host}?http_path={http_path}")
+    return engine
+
+def test_basic_connection(db_engine):
+    """Make sure we can connect and run basic query
+    """
+
+    curs = db_engine.execute("SELECT id FROM RANGE(100)")
+    result = curs.fetchall()
+    assert len(result) == 100
+
+def test_create_and_drop_table(db_engine):
+    """Make sure we can automatically create and drop a table defined with SQLAlchemy's MetaData object
+    """
+    
+    mdo = sqlalchemy.MetaData()
+    this_moment = datetime.datetime.utcnow().strftime("%s")
+    
+    tname = f"integration_test_table_{this_moment}"
+
+    t1 = sqlalchemy.Table(
+        tname,
+        mdo,
+        sqlalchemy.Column('f_short', sqlalchemy.types.SMALLINT),
+        sqlalchemy.Column('f_int', sqlalchemy.types.Integer),
+        sqlalchemy.Column('f_long', sqlalchemy.types.BigInteger),
+        sqlalchemy.Column('f_float', sqlalchemy.types.Float),
+        sqlalchemy.Column('f_decimal', sqlalchemy.types.DECIMAL),
+        sqlalchemy.Column('f_boolean', sqlalchemy.types.BOOLEAN)
+    )
+
+    mdo.create_all(bind=db_engine,checkfirst=True)
+
+    check_it_exists = db_engine.execute(f"DESCRIBE TABLE EXTENDED {tname}")
+    
+    mdo.drop_all(db_engine, checkfirst=True)