Skip to content

Commit

Permalink
support dbt v1.2 (#6)
Browse files Browse the repository at this point in the history
Support dbt v1.2
1. Lift + shift for cross-db macros
2. Support connection retry
3. Support grant
4. Doc Update

Co-authored-by: Shi Yuhang <1136742008.com>
  • Loading branch information
shiyuhang0 committed Aug 22, 2022
1 parent 89fbbda commit ad580d7
Show file tree
Hide file tree
Showing 46 changed files with 1,657 additions and 48 deletions.
16 changes: 16 additions & 0 deletions .github/workflows/main.yml
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,10 @@ jobs:
- name: Run tests
run: |
mysql -P4000 -uroot -h127.0.0.1 < tests/functional/adapter/tidb/grant/create_user.sql
export DBT_TEST_USER_1=user1
export DBT_TEST_USER_2=user2
export DBT_TEST_USER_3=user3
PYTHONPATH=. pytest tests/functional/adapter/tidb
tidb_5_3:
Expand Down Expand Up @@ -66,6 +70,10 @@ jobs:
- name: Run tests
run: |
mysql -P4000 -uroot -h127.0.0.1 < tests/functional/adapter/tidb/grant/create_user.sql
export DBT_TEST_USER_1=user1
export DBT_TEST_USER_2=user2
export DBT_TEST_USER_3=user3
PYTHONPATH=. pytest tests/functional/adapter/tidb
tidb_5_1:
Expand Down Expand Up @@ -97,6 +105,10 @@ jobs:
- name: Run tests
run: |
mysql -P4000 -uroot -h127.0.0.1 < tests/functional/adapter/tidb5_1/grant/create_user.sql
export DBT_TEST_USER_1=user1
export DBT_TEST_USER_2=user2
export DBT_TEST_USER_3=user3
PYTHONPATH=. pytest tests/functional/adapter/tidb5_1
tidb_4_0:
Expand Down Expand Up @@ -128,4 +140,8 @@ jobs:
- name: Run tests
run: |
mysql -P4000 -uroot -h127.0.0.1 < tests/functional/adapter/tidb4_0/grant/create_user.sql
export DBT_TEST_USER_1=user1
export DBT_TEST_USER_2=user2
export DBT_TEST_USER_3=user3
PYTHONPATH=. pytest tests/functional/adapter/tidb4_0
32 changes: 31 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ Thanks to them for their excellent work.
## Table of Contents
* [Installation](#installation)
* [Supported features](#supported-features)
* [Supported functions](#supported-functions)
* [Profile Configuration](#profile-configuration)
* [Database User Privileges](#database-user-privileges)
* [Running Tests](#running-tests)
Expand Down Expand Up @@ -46,6 +47,8 @@ $ pip install dbt-tidb
|||| Custom data tests |
|||| Docs generate |
|||| Snapshots |
|||| Connection retry |
|||| Grant |

Note:

Expand All @@ -55,6 +58,31 @@ Note:
* TiDB 4.X does not support using SQL func in `CREATE VIEW`, avoid it in your SQL code.
You can find more detail [here](https://github.com/pingcap/tidb/pull/27252).

## Supported functions

cross-db macros are moved from dbt-utils into dbt-core, so you can use the following functions directly, see [dbt-util](https://github.com/dbt-labs/dbt-utils) on how to use them.
- bool_or
- cast_bool_to_text
- dateadd
- datediff
- date_trunc
- hash
- safe_cast
- split_part
- last_day
- cast_bool_to_text
- concat
- escape_single_quotes
- except
- intersect
- length
- position
- replace
- right
- listagg (not support yet)

> pay attention that datediff is a little different from dbt-util that it will round down rather than round up.
## Profile Configuration

TiDB targets should be set up using the following configuration in your `profiles.yml` file.
Expand All @@ -72,6 +100,7 @@ your_profile_name:
schema: database_name
username: tidb_username
password: tidb_password
retries: 2
```

| Option | Description | Required? | Example |
Expand All @@ -82,6 +111,7 @@ your_profile_name:
| schema | Specify the schema (database) to build models into | Required | `analytics` |
| username | The username to use to connect to the server | Required | `dbt_admin` |
| password | The password to use for authenticating to the server | Required | `correct-horse-battery-staple` |
| retries | The retry times for connection to TiDB (1 in default) | Optional | `2` |

## Database User Privileges

Expand All @@ -102,7 +132,7 @@ You can find some help [here](https://docs.pingcap.com/tidb/v4.0/privilege-manag

## Running Tests

See [test/README.md](test/README.md) for details on running the integration tests.
See [tests/README.md](tests/README.md) for details on running the integration tests.

## Example

Expand Down
40 changes: 13 additions & 27 deletions dbt/adapters/tidb/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ class TiDBCredentials(Credentials):
username: Optional[str] = None
password: Optional[str] = None
charset: Optional[str] = None
retries: int = 1

_ALIASES = {
"UID": "username",
Expand Down Expand Up @@ -87,35 +88,20 @@ def open(cls, connection):
if credentials.port:
kwargs["port"] = credentials.port

try:
connection.handle = mysql.connector.connect(**kwargs)
connection.state = "open"
except mysql.connector.Error:

try:
logger.debug(
"Failed connection without supplying the `database`. "
"Trying again with `database` included."
)

# Try again with the database included
kwargs["database"] = credentials.schema
def connect():
handle = mysql.connector.connect(**kwargs)
return handle

connection.handle = mysql.connector.connect(**kwargs)
connection.state = "open"
except mysql.connector.Error as e:
# we just retry for any error now
retryable_exceptions = [mysql.connector.Error]

logger.debug(
"Got an error when attempting to open a tidb "
"connection: '{}'".format(e)
)

connection.handle = None
connection.state = "fail"

raise dbt.exceptions.FailedToConnectException(str(e))

return connection
return cls.retry_connection(
connection,
connect=connect,
logger=logger,
retry_limit=credentials.retries,
retryable_exceptions=retryable_exceptions,
)

@classmethod
def get_credentials(cls, credentials):
Expand Down
29 changes: 29 additions & 0 deletions dbt/include/tidb/macros/apply_grants.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
-- support Select/Insert/Delete/Update now
{% macro default__get_show_grant_sql(relation) %}

select case(Table_priv) when null then null else 'select' end as privilege_type, `User` as grantee from mysql.tables_priv where `DB` = '{{relation.schema}}' and `Table_name` = '{{relation.identifier}}' and Table_priv like '%Select%'
union ALL
select case(Table_priv) when null then null else 'insert' end as privilege_type, `User` as grantee from mysql.tables_priv where `DB` = '{{relation.schema}}' and `Table_name` = '{{relation.identifier}}' and Table_priv like '%Insert%'
union ALL
select case(Table_priv) when null then null else 'update' end as privilege_type, `User` as grantee from mysql.tables_priv where `DB` = '{{relation.schema}}' and `Table_name` = '{{relation.identifier}}' and Table_priv like '%Update%'
union ALL
select case(Table_priv) when null then null else 'delete' end as privilege_type, `User` as grantee from mysql.tables_priv where `DB` = '{{relation.schema}}' and `Table_name` = '{{relation.identifier}}' and Table_priv like '%Delete%'

{% endmacro %}

{%- macro tidb__get_grant_sql(relation, privilege, grantees) -%}
grant {{ privilege }} on {{ relation }} to {{ '\"' + grantees|join('\", \"') + '\"' }}
{%- endmacro -%}

{%- macro tidb__get_revoke_sql(relation, privilege, grantees) -%}
revoke {{ privilege }} on {{ relation }} from {{ '\"' + grantees|join('\", \"') + '\"' }}
{%- endmacro -%}

-- tidb-dbt does not support multi=true now, so we need to split every grant/revoke statement
{% macro tidb__call_dcl_statements(dcl_statement_list) %}
{% for dcl_statement in dcl_statement_list %}
{% call statement('grant_or_revoke') %}
{{ dcl_statement }}
{% endcall %}
{% endfor %}
{% endmacro %}
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@
{% set existing_relation = load_relation(this) %}
{% set tmp_relation = make_temp_relation(this) %}

-- grab current tables grants config for comparision later on
{% set grant_config = config.get('grants') %}

{{ run_hooks(pre_hooks, inside_transaction=False) }}

-- `BEGIN` happens here:
Expand Down Expand Up @@ -44,6 +47,8 @@
{{ build_sql }}
{% endcall %}

{% set should_revoke = should_revoke(existing_relation, full_refresh_mode) %}
{% do apply_grants(target_relation, grant_config, should_revoke=should_revoke) %}
{% do persist_docs(target_relation, model) %}

{{ run_hooks(post_hooks, inside_transaction=True) }}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@

{%- set strategy_name = config.get('strategy') -%}
{%- set unique_key = config.get('unique_key') %}
-- grab current tables grants config for comparision later on
{%- set grant_config = config.get('grants') -%}

{% if not adapter.check_schema_exists(model.database, model.schema) %}
{% do create_schema(model.database, model.schema) %}
Expand Down Expand Up @@ -98,6 +100,9 @@

{% endif %}

{% set should_revoke = should_revoke(target_relation_exists, full_refresh_mode=False) %}
{% do apply_grants(target_relation, grant_config, should_revoke=should_revoke) %}

{% do persist_docs(target_relation, model) %}

{{ run_hooks(post_hooks, inside_transaction=True) }}
Expand Down
41 changes: 28 additions & 13 deletions dbt/include/tidb/macros/materializations/snapshot/strategies.sql
Original file line number Diff line number Diff line change
Expand Up @@ -6,26 +6,41 @@
{%- endfor -%}))
{%- endmacro %}

{% macro snapshot_check_all_get_existing_columns(node, target_exists) -%}
{%- set query_columns = get_columns_in_query(node['compiled_sql']) -%}
-- copy from dbt-core v1.2, just alter database=None in adapter.get_relation
{% macro snapshot_check_all_get_existing_columns(node, target_exists, check_cols_config) -%}
{%- if not target_exists -%}
{# no table yet -> return whatever the query does #}
{{ return([false, query_columns]) }}
{#-- no table yet -> return whatever the query does --#}
{{ return((false, query_columns)) }}
{%- endif -%}
{# handle any schema changes #}
{%- set target_table = node.get('alias', node.get('name')) -%}
{%- set target_relation = adapter.get_relation(database=None, schema=node.schema, identifier=target_table) -%}
{%- set existing_cols = get_columns_in_query('select * from ' ~ target_relation) -%}
{%- set ns = namespace() -%} {# handle for-loop scoping with a namespace #}

{#-- handle any schema changes --#}
{%- set target_relation = adapter.get_relation(database=None, schema=node.schema, identifier=node.alias) -%}

{% if check_cols_config == 'all' %}
{%- set query_columns = get_columns_in_query(node['compiled_sql']) -%}

{% elif check_cols_config is iterable and (check_cols_config | length) > 0 %}
{#-- query for proper casing/quoting, to support comparison below --#}
{%- set select_check_cols_from_target -%}
select {{ check_cols_config | join(', ') }} from ({{ node['compiled_sql'] }}) subq
{%- endset -%}
{% set query_columns = get_columns_in_query(select_check_cols_from_target) %}

{% else %}
{% do exceptions.raise_compiler_error("Invalid value for 'check_cols': " ~ check_cols_config) %}
{% endif %}

{%- set existing_cols = adapter.get_columns_in_relation(target_relation) | map(attribute = 'name') | list -%}
{%- set ns = namespace() -%} {#-- handle for-loop scoping with a namespace --#}
{%- set ns.column_added = false -%}

{%- set intersection = [] -%}
{%- for col in query_columns -%}
{%- if col in existing_cols -%}
{%- do intersection.append(col) -%}
{%- do intersection.append(adapter.quote(col)) -%}
{%- else -%}
{% set ns.column_added = true %}
{%- endif -%}
{%- endif -%}
{%- endfor -%}
{{ return([ns.column_added, intersection]) }}
{%- endmacro %}
{{ return((ns.column_added, intersection)) }}
{%- endmacro %}
7 changes: 7 additions & 0 deletions dbt/include/tidb/macros/utils/bool_or.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
-- bool_or is an agg function which will return true once any expression is true
-- use max to replace it
{% macro tidb__bool_or(expression) -%}

max({{ expression }})

{%- endmacro %}
8 changes: 8 additions & 0 deletions dbt/include/tidb/macros/utils/cast_bool_to_text.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
{% macro tidb__cast_bool_to_text(field) %}

case {{ field }}
when true then 'true'
when false then 'false'
end

{% endmacro %}
31 changes: 31 additions & 0 deletions dbt/include/tidb/macros/utils/date_trunc.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
-- date_trunc can truncate date with given datepart(year,month,quarter,day is supported)
{% macro tidb__date_trunc(datepart, date) -%}

{%- if datepart =='day' -%}

DATE_FORMAT({{date}}, '%Y-%m-%d')

{%- elif datepart == 'month' -%}

DATE_FORMAT({{date}}, '%Y-%m-01')

{%- elif datepart == 'quarter' -%}

case QUARTER({{date}})
when 1 then DATE_FORMAT({{date}}, '%Y-01-01')
when 2 then DATE_FORMAT({{date}}, '%Y-04-01')
when 2 then DATE_FORMAT({{date}}, '%Y-07-01')
when 2 then DATE_FORMAT({{date}}, '%Y-10-01')
end

{%- elif datepart == 'year' -%}

DATE_FORMAT({{date}}, '%Y-01-01')

{%- else -%}

{{ exceptions.raise_compiler_error("macro date_trunc not implemented for datepart ~ '" ~ datepart ~ "' ~ on TiDB") }}

{%- endif -%}

{%- endmacro %}
9 changes: 9 additions & 0 deletions dbt/include/tidb/macros/utils/dateadd.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
-- add interval to given from_date_or_timestamp with datepart(day,month,hour...)
{% macro tidb__dateadd(datepart, interval, from_date_or_timestamp) %}

DATE_ADD(
{{ from_date_or_timestamp }},
interval {{ interval }} {{ datepart }}
)

{% endmacro %}
14 changes: 14 additions & 0 deletions dbt/include/tidb/macros/utils/datediff.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
-- the behavior is a little different from default_datediff that it will round down rather than round up
-- and millisecond is not supported
{% macro tidb__datediff(first_date, second_date, datepart) -%}
{%- if datepart =='millisecond' -%}

{{ exceptions.raise_compiler_error("macro datediff not implemented for datepart ~ '" ~ datepart ~ "' ~ on TiDB") }}

{%- else -%}

TIMESTAMPDIFF({{datepart}},{{first_date}},{{second_date}})

{%- endif -%}

{%- endmacro %}
5 changes: 5 additions & 0 deletions dbt/include/tidb/macros/utils/hash.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
{% macro tidb__hash(field) -%}

md5(cast({{ field }} as CHAR))

{%- endmacro %}

0 comments on commit ad580d7

Please sign in to comment.