databricks-sqlalchemy 2.0.3__py3-none-any.whl → 2.0.5__py3-none-any.whl

databricks/sqlalchemy/test_local/__init__.py

@@ -1,5 +0,0 @@
-"""
-This module contains tests entirely maintained by Databricks.
-
-These tests do not rely on SQLAlchemy's custom test runner.
-"""

databricks/sqlalchemy/test_local/conftest.py

@@ -1,44 +0,0 @@
-import os
-import pytest
-
-
-@pytest.fixture(scope="session")
-def host():
-    return os.getenv("DATABRICKS_SERVER_HOSTNAME")
-
-
-@pytest.fixture(scope="session")
-def http_path():
-    return os.getenv("DATABRICKS_HTTP_PATH")
-
-
-@pytest.fixture(scope="session")
-def access_token():
-    return os.getenv("DATABRICKS_TOKEN")
-
-
-@pytest.fixture(scope="session")
-def ingestion_user():
-    return os.getenv("DATABRICKS_USER")
-
-
-@pytest.fixture(scope="session")
-def catalog():
-    return os.getenv("DATABRICKS_CATALOG")
-
-
-@pytest.fixture(scope="session")
-def schema():
-    return os.getenv("DATABRICKS_SCHEMA", "default")
-
-
-@pytest.fixture(scope="session", autouse=True)
-def connection_details(host, http_path, access_token, ingestion_user, catalog, schema):
-    return {
-        "host": host,
-        "http_path": http_path,
-        "access_token": access_token,
-        "ingestion_user": ingestion_user,
-        "catalog": catalog,
-        "schema": schema,
-    }

databricks/sqlalchemy/test_local/e2e/test_basic.py

@@ -1,543 +0,0 @@
-import datetime
-import decimal
-from typing import Tuple, Union, List
-from unittest import skipIf
-
-import pytest
-from sqlalchemy import (
-    Column,
-    MetaData,
-    Table,
-    Text,
-    create_engine,
-    insert,
-    select,
-    text,
-)
-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, Integer, String
-
-try:
-    from sqlalchemy.orm import declarative_base
-except ImportError:
-    from sqlalchemy.ext.declarative import declarative_base
-
-
-USER_AGENT_TOKEN = "PySQL e2e Tests"
-
-
-def sqlalchemy_1_3():
-    import sqlalchemy
-
-    return sqlalchemy.__version__.startswith("1.3")
-
-
-def version_agnostic_select(object_to_select, *args, **kwargs):
-    """
-    SQLAlchemy==1.3.x requires arguments to select() to be a Python list
-
-    https://docs.sqlalchemy.org/en/20/changelog/migration_14.html#orm-query-is-internally-unified-with-select-update-delete-2-0-style-execution-available
-    """
-
-    if sqlalchemy_1_3():
-        return select([object_to_select], *args, **kwargs)
-    else:
-        return select(object_to_select, *args, **kwargs)
-
-
-def version_agnostic_connect_arguments(connection_details) -> Tuple[str, dict]:
-    HOST = connection_details["host"]
-    HTTP_PATH = connection_details["http_path"]
-    ACCESS_TOKEN = connection_details["access_token"]
-    CATALOG = connection_details["catalog"]
-    SCHEMA = connection_details["schema"]
-
-    ua_connect_args = {"_user_agent_entry": USER_AGENT_TOKEN}
-
-    if sqlalchemy_1_3():
-        conn_string = f"databricks://token:{ACCESS_TOKEN}@{HOST}"
-        connect_args = {
-            **ua_connect_args,
-            "http_path": HTTP_PATH,
-            "server_hostname": HOST,
-            "catalog": CATALOG,
-            "schema": SCHEMA,
-        }
-
-        return conn_string, connect_args
-    else:
-        return (
-            f"databricks://token:{ACCESS_TOKEN}@{HOST}?http_path={HTTP_PATH}&catalog={CATALOG}&schema={SCHEMA}",
-            ua_connect_args,
-        )
-
-
-@pytest.fixture
-def db_engine(connection_details) -> Engine:
-    conn_string, connect_args = version_agnostic_connect_arguments(connection_details)
-    return create_engine(conn_string, connect_args=connect_args)
-
-
-def run_query(db_engine: Engine, query: Union[str, Text]):
-    if not isinstance(query, Text):
-        _query = text(query)  # type: ignore
-    else:
-        _query = query  # type: ignore
-    with db_engine.begin() as conn:
-        return conn.execute(_query).fetchall()
-
-
-@pytest.fixture
-def samples_engine(connection_details) -> Engine:
-    details = connection_details.copy()
-    details["catalog"] = "samples"
-    details["schema"] = "nyctaxi"
-    conn_string, connect_args = version_agnostic_connect_arguments(details)
-    return create_engine(conn_string, connect_args=connect_args)
-
-
-@pytest.fixture()
-def base(db_engine):
-    return declarative_base()
-
-
-@pytest.fixture()
-def session(db_engine):
-    return Session(db_engine)
-
-
-@pytest.fixture()
-def metadata_obj(db_engine):
-    return MetaData()
-
-
-def test_can_connect(db_engine):
-    simple_query = "SELECT 1"
-    result = run_query(db_engine, simple_query)
-    assert len(result) == 1
-
-
-def test_connect_args(db_engine):
-    """Verify that extra connect args passed to sqlalchemy.create_engine are passed to DBAPI
-
-    This will most commonly happen when partners supply a user agent entry
-    """
-
-    conn = db_engine.connect()
-    connection_headers = conn.connection.thrift_backend._transport._headers
-    user_agent = connection_headers["User-Agent"]
-
-    expected = f"(sqlalchemy + {USER_AGENT_TOKEN})"
-    assert expected in user_agent
-
-
-@pytest.mark.skipif(sqlalchemy_1_3(), reason="Pandas requires SQLAlchemy >= 1.4")
-@pytest.mark.skip(
-    reason="DBR is currently limited to 256 parameters per call to .execute(). Test cannot pass."
-)
-def test_pandas_upload(db_engine, metadata_obj):
-    import pandas as pd
-
-    SCHEMA = "default"
-    try:
-        df = pd.read_excel(
-            "src/databricks/sqlalchemy/test_local/e2e/demo_data/MOCK_DATA.xlsx"
-        )
-        df.to_sql(
-            "mock_data",
-            db_engine,
-            schema=SCHEMA,
-            index=False,
-            method="multi",
-            if_exists="replace",
-        )
-
-        df_after = pd.read_sql_table("mock_data", db_engine, schema=SCHEMA)
-        assert len(df) == len(df_after)
-    except Exception as e:
-        raise e
-    finally:
-        db_engine.execute("DROP TABLE mock_data")
-
-
-def test_create_table_not_null(db_engine, metadata_obj: MetaData):
-    table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
-
-    SampleTable = Table(
-        table_name,
-        metadata_obj,
-        Column("name", String(255)),
-        Column("episodes", Integer),
-        Column("some_bool", BOOLEAN, nullable=False),
-    )
-
-    metadata_obj.create_all(db_engine)
-
-    columns = db_engine.dialect.get_columns(
-        connection=db_engine.connect(), table_name=table_name
-    )
-
-    name_column_description = columns[0]
-    some_bool_column_description = columns[2]
-
-    assert name_column_description.get("nullable") is True
-    assert some_bool_column_description.get("nullable") is False
-
-    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") == None
-
-    metadata_obj.drop_all(db_engine)
-
-
-def test_bulk_insert_with_core(db_engine, metadata_obj, session):
-    import random
-
-    # Maximum number of parameter is 256. 256/4 == 64
-    num_to_insert = 64
-
-    table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
-
-    names = ["Bim", "Miki", "Sarah", "Ira"]
-
-    SampleTable = Table(
-        table_name, metadata_obj, Column("name", String(255)), Column("number", Integer)
-    )
-
-    rows = [
-        {"name": names[i % 3], "number": random.choice(range(64))}
-        for i in range(num_to_insert)
-    ]
-
-    metadata_obj.create_all(db_engine)
-    with db_engine.begin() as conn:
-        conn.execute(insert(SampleTable).values(rows))
-
-    with db_engine.begin() as conn:
-        rows = conn.execute(version_agnostic_select(SampleTable)).fetchall()
-
-    assert len(rows) == num_to_insert
-
-
-def test_create_insert_drop_table_core(base, db_engine, metadata_obj: MetaData):
-    """ """
-
-    SampleTable = Table(
-        "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s")),
-        metadata_obj,
-        Column("name", String(255)),
-        Column("episodes", Integer),
-        Column("some_bool", BOOLEAN),
-        Column("dollars", DECIMAL(10, 2)),
-    )
-
-    metadata_obj.create_all(db_engine)
-
-    insert_stmt = insert(SampleTable).values(
-        name="Bim Adewunmi", episodes=6, some_bool=True, dollars=decimal.Decimal(125)
-    )
-
-    with db_engine.connect() as conn:
-        conn.execute(insert_stmt)
-
-    select_stmt = version_agnostic_select(SampleTable)
-    with db_engine.begin() as conn:
-        resp = conn.execute(select_stmt)
-
-    result = resp.fetchall()
-
-    assert len(result) == 1
-
-    metadata_obj.drop_all(db_engine)
-
-
-# ORM tests are made following this tutorial
-# https://docs.sqlalchemy.org/en/14/orm/quickstart.html
-
-
-@skipIf(False, "Unity catalog must be supported")
-def test_create_insert_drop_table_orm(db_engine):
-    """ORM classes built on the declarative base class must have a primary key.
-    This is restricted to Unity Catalog.
-    """
-
-    class Base(DeclarativeBase):
-        pass
-
-    class SampleObject(Base):
-        __tablename__ = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
-
-        name: Mapped[str] = mapped_column(String(255), primary_key=True)
-        episodes: Mapped[int] = mapped_column(Integer)
-        some_bool: Mapped[bool] = mapped_column(BOOLEAN)
-
-    Base.metadata.create_all(db_engine)
-
-    sample_object_1 = SampleObject(name="Bim Adewunmi", episodes=6, some_bool=True)
-    sample_object_2 = SampleObject(name="Miki Meek", episodes=12, some_bool=False)
-
-    session = Session(db_engine)
-    session.add(sample_object_1)
-    session.add(sample_object_2)
-    session.flush()
-
-    stmt = version_agnostic_select(SampleObject).where(
-        SampleObject.name.in_(["Bim Adewunmi", "Miki Meek"])
-    )
-
-    if sqlalchemy_1_3():
-        output = [i for i in session.execute(stmt)]
-    else:
-        output = [i for i in session.scalars(stmt)]
-
-    assert len(output) == 2
-
-    Base.metadata.drop_all(db_engine)
-
-
-def test_dialect_type_mappings(db_engine, metadata_obj: MetaData):
-    """Confirms that we get back the same time we declared in a model and inserted using Core"""
-
-    class Base(DeclarativeBase):
-        pass
-
-    SampleTable = Table(
-        "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s")),
-        metadata_obj,
-        Column("string_example", String(255)),
-        Column("integer_example", Integer),
-        Column("boolean_example", BOOLEAN),
-        Column("decimal_example", DECIMAL(10, 2)),
-        Column("date_example", Date),
-    )
-
-    string_example = ""
-    integer_example = 100
-    boolean_example = True
-    decimal_example = decimal.Decimal(125)
-    date_example = datetime.date(2013, 1, 1)
-
-    metadata_obj.create_all(db_engine)
-
-    insert_stmt = insert(SampleTable).values(
-        string_example=string_example,
-        integer_example=integer_example,
-        boolean_example=boolean_example,
-        decimal_example=decimal_example,
-        date_example=date_example,
-    )
-
-    with db_engine.connect() as conn:
-        conn.execute(insert_stmt)
-
-    select_stmt = version_agnostic_select(SampleTable)
-    with db_engine.begin() as conn:
-        resp = conn.execute(select_stmt)
-
-    result = resp.fetchall()
-    this_row = result[0]
-
-    assert this_row.string_example == string_example
-    assert this_row.integer_example == integer_example
-    assert this_row.boolean_example == boolean_example
-    assert this_row.decimal_example == decimal_example
-    assert this_row.date_example == date_example
-
-    metadata_obj.drop_all(db_engine)
-
-
-def test_inspector_smoke_test(samples_engine: Engine):
-    """It does not appear that 3L namespace is supported here"""
-
-    schema, table = "nyctaxi", "trips"
-
-    try:
-        inspector = Inspector.from_engine(samples_engine)
-    except Exception as e:
-        assert False, f"Could not build inspector: {e}"
-
-    # Expect six columns
-    columns = inspector.get_columns(table, schema=schema)
-
-    # Expect zero views, but the method should return
-    views = inspector.get_view_names(schema=schema)
-
-    assert (
-        len(columns) == 6
-    ), "Dialect did not find the expected number of columns in samples.nyctaxi.trips"
-    assert len(views) == 0, "Views could not be fetched"
-
-
-@pytest.mark.skip(reason="engine.table_names has been removed in sqlalchemy verison 2")
-def test_get_table_names_smoke_test(samples_engine: Engine):
-    with samples_engine.connect() as conn:
-        _names = samples_engine.table_names(schema="nyctaxi", connection=conn)  # type: ignore
-        _names is not None, "get_table_names did not succeed"
-
-
-def test_has_table_across_schemas(
-    db_engine: Engine, samples_engine: Engine, catalog: str, schema: str
-):
-    """For this test to pass these conditions must be met:
-    - Table samples.nyctaxi.trips must exist
-    - Table samples.tpch.customer must exist
-    - The `catalog` and `schema` environment variables must be set and valid
-    """
-
-    with samples_engine.connect() as conn:
-        # 1) Check for table within schema declared at engine creation time
-        assert samples_engine.dialect.has_table(connection=conn, table_name="trips")
-
-        # 2) Check for table within another schema in the same catalog
-        assert samples_engine.dialect.has_table(
-            connection=conn, table_name="customer", schema="tpch"
-        )
-
-        # 3) Check for a table within a different catalog
-        # Create a table in a different catalog
-        with db_engine.connect() as conn:
-            conn.execute(text("CREATE TABLE test_has_table (numbers_are_cool INT);"))
-
-            try:
-                # Verify that this table is not found in the samples catalog
-                assert not samples_engine.dialect.has_table(
-                    connection=conn, table_name="test_has_table"
-                )
-                # Verify that this table is found in a separate catalog
-                assert samples_engine.dialect.has_table(
-                    connection=conn,
-                    table_name="test_has_table",
-                    schema=schema,
-                    catalog=catalog,
-                )
-            finally:
-                conn.execute(text("DROP TABLE test_has_table;"))
-
-
-def test_user_agent_adjustment(db_engine):
-    # If .connect() is called multiple times on an engine, don't keep pre-pending the user agent
-    # https://github.com/databricks/databricks-sql-python/issues/192
-    c1 = db_engine.connect()
-    c2 = db_engine.connect()
-
-    def get_conn_user_agent(conn):
-        return conn.connection.dbapi_connection.thrift_backend._transport._headers.get(
-            "User-Agent"
-        )
-
-    ua1 = get_conn_user_agent(c1)
-    ua2 = get_conn_user_agent(c2)
-    same_ua = ua1 == ua2
-
-    c1.close()
-    c2.close()
-
-    assert same_ua, f"User agents didn't match \n {ua1} \n {ua2}"
-
-
-@pytest.fixture
-def sample_table(metadata_obj: MetaData, db_engine: Engine):
-    """This fixture creates a sample table and cleans it up after the test is complete."""
-    from databricks.sqlalchemy._parse import GET_COLUMNS_TYPE_MAP
-
-    table_name = "PySQLTest_{}".format(datetime.datetime.utcnow().strftime("%s"))
-
-    args: List[Column] = [
-        Column(colname, coltype) for colname, coltype in GET_COLUMNS_TYPE_MAP.items()
-    ]
-
-    SampleTable = Table(table_name, metadata_obj, *args)
-
-    metadata_obj.create_all(db_engine)
-
-    yield table_name
-
-    metadata_obj.drop_all(db_engine)
-
-
-def test_get_columns(db_engine, sample_table: str):
-    """Created after PECO-1297 and Github Issue #295 to verify that get_columsn behaves like it should for all known SQLAlchemy types"""
-
-    inspector = Inspector.from_engine(db_engine)
-
-    # this raises an exception if `parse_column_info_from_tgetcolumnsresponse` fails a lookup
-    columns = inspector.get_columns(sample_table)
-
-    assert True
-
-
-class TestCommentReflection:
-    @pytest.fixture(scope="class")
-    def engine(self, connection_details: dict):
-        HOST = connection_details["host"]
-        HTTP_PATH = connection_details["http_path"]
-        ACCESS_TOKEN = connection_details["access_token"]
-        CATALOG = connection_details["catalog"]
-        SCHEMA = connection_details["schema"]
-
-        connection_string = f"databricks://token:{ACCESS_TOKEN}@{HOST}?http_path={HTTP_PATH}&catalog={CATALOG}&schema={SCHEMA}"
-        connect_args = {"_user_agent_entry": USER_AGENT_TOKEN}
-
-        engine = create_engine(connection_string, connect_args=connect_args)
-        return engine
-
-    @pytest.fixture
-    def inspector(self, engine: Engine) -> Inspector:
-        return Inspector.from_engine(engine)
-
-    @pytest.fixture(scope="class")
-    def table(self, engine):
-        md = MetaData()
-        tbl = Table(
-            "foo",
-            md,
-            Column("bar", String, comment="column comment"),
-            comment="table comment",
-        )
-        md.create_all(bind=engine)
-
-        yield tbl
-
-        md.drop_all(bind=engine)
-
-    def test_table_comment_reflection(self, inspector: Inspector, table: Table):
-        comment = inspector.get_table_comment(table.name)
-        assert comment == {"text": "table comment"}
-
-    def test_column_comment(self, inspector: Inspector, table: Table):
-        result = inspector.get_columns(table.name)[0].get("comment")
-        assert result == "column comment"

databricks/sqlalchemy/test_local/test_ddl.py

@@ -1,96 +0,0 @@
-import pytest
-from sqlalchemy import Column, MetaData, String, Table, create_engine
-from sqlalchemy.schema import (
-    CreateTable,
-    DropColumnComment,
-    DropTableComment,
-    SetColumnComment,
-    SetTableComment,
-)
-
-
-class DDLTestBase:
-    engine = create_engine(
-        "databricks://token:****@****?http_path=****&catalog=****&schema=****"
-    )
-
-    def compile(self, stmt):
-        return str(stmt.compile(bind=self.engine))
-
-
-class TestColumnCommentDDL(DDLTestBase):
-    @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 ''"
-
-
-class TestTableCommentDDL(DDLTestBase):
-    @pytest.fixture
-    def metadata(self) -> MetaData:
-        """Assemble a metadata object with one table containing one column."""
-        metadata = MetaData()
-
-        col1 = Column("foo", String)
-        col2 = Column("foo", String)
-        tbl_w_comment = Table("martin", metadata, col1, comment="foobar")
-        tbl_wo_comment = Table("prs", metadata, col2)
-
-        return metadata
-
-    @pytest.fixture
-    def table_with_comment(self, metadata) -> Table:
-        return metadata.tables.get("martin")
-
-    @pytest.fixture
-    def table_without_comment(self, metadata) -> Table:
-        return metadata.tables.get("prs")
-
-    def test_create_table_with_comment(self, table_with_comment):
-        stmt = CreateTable(table_with_comment)
-        output = self.compile(stmt)
-        assert "USING DELTA" in output
-        assert "COMMENT 'foobar'" in output
-
-    def test_alter_table_add_comment(self, table_without_comment: Table):
-        table_without_comment.comment = "wireless mechanical keyboard"
-        stmt = SetTableComment(table_without_comment)
-        output = self.compile(stmt)
-
-        assert output == "COMMENT ON TABLE prs IS 'wireless mechanical keyboard'"
-
-    def test_alter_table_drop_comment(self, table_with_comment):
-        """The syntax for COMMENT ON is here: https://docs.databricks.com/en/sql/language-manual/sql-ref-syntax-ddl-comment.html"""
-        stmt = DropTableComment(table_with_comment)
-        output = self.compile(stmt)
-        assert output == "COMMENT ON TABLE martin IS NULL"