Skip to content

Commit

Permalink
upgrade to support dbt-core v1.6.4 (#52)
Browse files Browse the repository at this point in the history
Actually this should work for dbt-core>=1.4.0 but I didn't test against old
verions of dbt-core.

Tested against Python 3.8 and 3.11, TiDB 6.5.5.

See dbt-labs/dbt-core#6624

    XxxException -> DbtXxxError
    raise_compiler_error -> CompilationError
  • Loading branch information
Dieken committed Dec 13, 2023
1 parent 2a48718 commit 14e657d
Show file tree
Hide file tree
Showing 8 changed files with 33 additions and 28 deletions.
2 changes: 1 addition & 1 deletion dbt/adapters/tidb/__version__.py
Original file line number Diff line number Diff line change
@@ -1 +1 @@
version = "1.3.0"
version = "1.6.4"
8 changes: 4 additions & 4 deletions dbt/adapters/tidb/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ def __init__(self, **kwargs):
def __post_init__(self):
# mysql classifies database and schema as the same thing
if self.database is not None and self.database != self.schema:
raise dbt.exceptions.RuntimeException(
raise dbt.exceptions.DbtRuntimeError(
f" schema: {self.schema} \n"
f" database: {self.database} \n"
f"On TiDB, database must be omitted or have the same value as"
Expand Down Expand Up @@ -126,19 +126,19 @@ def exception_handler(self, sql):
logger.debug("Failed to release connection!")
pass

raise dbt.exceptions.DatabaseException(str(e).strip()) from e
raise dbt.exceptions.DbtDatabaseError(str(e).strip()) from e

except Exception as e:
logger.debug("Error running SQL: {}", sql)
logger.debug("Rolling back transaction.")
self.rollback_if_open()
if isinstance(e, dbt.exceptions.RuntimeException):
if isinstance(e, dbt.exceptions.DbtRuntimeError):
# during a sql query, an internal to dbt exception was raised.
# this sounds a lot like a signal handler and probably has
# useful information, so raise it without modification.
raise

raise dbt.exceptions.RuntimeException(e) from e
raise dbt.exceptions.DbtRuntimeError(e) from e

@classmethod
def get_response(cls, cursor) -> AdapterResponse:
Expand Down
14 changes: 7 additions & 7 deletions dbt/adapters/tidb/impl.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ def list_relations_without_caching(
kwargs = {"schema_relation": schema_relation}
try:
results = self.execute_macro(LIST_RELATIONS_MACRO_NAME, kwargs=kwargs)
except dbt.exceptions.RuntimeException as e:
except dbt.exceptions.DbtRuntimeError as e:
errmsg = getattr(e, "msg", "")
if f"TiDB database '{schema_relation}' not found" in errmsg:
return []
Expand All @@ -61,7 +61,7 @@ def list_relations_without_caching(
relations = []
for row in results:
if len(row) != 4:
raise dbt.exceptions.RuntimeException(
raise dbt.exceptions.DbtRuntimeError(
"Invalid value from "
f'"tidb__list_relations_without_caching({kwargs})", '
f"got {len(row)} values, expected 4"
Expand Down Expand Up @@ -94,7 +94,7 @@ def _get_columns_for_catalog(
def get_relation(
self, database: str, schema: str, identifier: str
) -> Optional[BaseRelation]:
if not self.Relation.include_policy.database:
if not self.Relation.get_default_include_policy().database:
database = None

return super().get_relation(database, schema, identifier)
Expand All @@ -121,7 +121,7 @@ def get_catalog(self, manifest):
schema_map = self._get_catalog_schemas(manifest)

if len(schema_map) > 1:
dbt.exceptions.raise_compiler_error(
raise dbt.exceptions.CompilationError(
f"Expected only one database in get_catalog, found "
f"{list(schema_map)}"
)
Expand Down Expand Up @@ -150,7 +150,7 @@ def _get_one_catalog(
manifest,
) -> agate.Table:
if len(schemas) != 1:
dbt.exceptions.raise_compiler_error(
raise dbt.exceptions.CompilationError(
f"Expected only one schema in tidb _get_one_catalog, found "
f"{schemas}"
)
Expand All @@ -160,7 +160,7 @@ def _get_one_catalog(

columns: List[Dict[str, Any]] = []
for relation in self.list_relations(database, schema):
logger.debug("Getting table schema for relation {}", relation)
logger.debug("Getting table schema for relation {}", str(relation))
columns.extend(self._get_columns_for_catalog(relation))
return agate.Table.from_object(columns, column_types=DEFAULT_TYPE_TESTER)

Expand Down Expand Up @@ -205,7 +205,7 @@ def string_add_sql(
elif location == "prepend":
return f"concat({value}, '{add_to}')"
else:
raise dbt.exceptions.RuntimeException(
raise dbt.exceptions.DbtRuntimeError(
f'Got an unexpected location value of "{location}"'
)

Expand Down
14 changes: 8 additions & 6 deletions dbt/adapters/tidb/relation.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from dataclasses import dataclass
from dataclasses import dataclass, field

from dbt.adapters.base.relation import BaseRelation, Policy
from dbt.exceptions import RuntimeException
from dbt.exceptions import DbtRuntimeError


@dataclass
Expand All @@ -20,17 +20,19 @@ class TiDBIncludePolicy(Policy):

@dataclass(frozen=True, eq=False, repr=False)
class TiDBRelation(BaseRelation):
quote_policy: TiDBQuotePolicy = TiDBQuotePolicy()
include_policy: TiDBIncludePolicy = TiDBIncludePolicy()
quote_policy: TiDBQuotePolicy = field(default_factory=lambda: TiDBQuotePolicy())
include_policy: TiDBIncludePolicy = field(
default_factory=lambda: TiDBIncludePolicy()
)
quote_character: str = "`"

def __post_init__(self):
if self.database != self.schema and self.database:
raise RuntimeException(f"Cannot set database {self.database} in tidb!")
raise DbtRuntimeError(f"Cannot set database {self.database} in tidb!")

def render(self):
if self.include_policy.database and self.include_policy.schema:
raise RuntimeException(
raise DbtRuntimeError(
"Got a tidb relation with schema and database set to "
"include, but only one can be set"
)
Expand Down
8 changes: 4 additions & 4 deletions requirements_dev.txt
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
dbt-core~=1.3.1
mysql-connector-python>=8.0.0,<8.1
dbt-core~=1.6.4
mysql-connector-python>=8.0.0,<8.2
pytest~=7.0
markupsafe==2.0.1
dbt-tests-adapter~=1.3.1
pytest-dotenv
dbt-tests-adapter~=1.6.4
pytest-dotenv
13 changes: 7 additions & 6 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,9 @@
import os
import sys

if sys.version_info < (3, 8) or sys.version_info >= (3, 11):
if sys.version_info < (3, 8):
print("Error: dbt-tidb does not support this version of Python.")
print("Please install Python 3.8 or higher but less than 3.11.")
print("Please install Python 3.8 or higher.")
sys.exit(1)

# require version of setuptools that supports find_namespace_packages
Expand All @@ -27,8 +27,8 @@
long_description = f.read()

package_name = "dbt-tidb"
package_version = "1.3.0"
dbt_core_version = "1.3.1"
package_version = "1.6.4"
dbt_core_version = "1.6.4"
description = """The TiDB adapter plugin for dbt"""

setup(
Expand All @@ -44,7 +44,7 @@
include_package_data=True,
install_requires=[
"dbt-core~={}".format(dbt_core_version),
"mysql-connector-python>=8.0.0,<8.1",
"mysql-connector-python>=8.0.0,<8.2",
],
zip_safe=False,
classifiers=[
Expand All @@ -56,6 +56,7 @@
"Programming Language :: Python :: 3.8",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
],
python_requires=">=3.8,<=3.10",
python_requires=">=3.8",
)
1 change: 1 addition & 0 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

pytest_plugins = ["dbt.tests.fixtures.project"]


# The profile dictionary, used to write out profiles.yml
# dbt will supply a unique schema per test, so we do not specify 'schema' here
@pytest.fixture(scope="class")
Expand Down
1 change: 1 addition & 0 deletions tests/functional/adapter/tidb5_1/grant/test_grant.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from dbt.tests.adapter.grants.test_seed_grants import BaseSeedGrants
from dbt.tests.adapter.grants.test_snapshot_grants import BaseSnapshotGrants


# need to export DBT_TEST_USER_1,DBT_TEST_USER_2,DBT_TEST_USER_3
class TestModelGrantsTiDB(BaseModelGrants):
pass
Expand Down

0 comments on commit 14e657d

Please sign in to comment.