support dbt v1.2 (#6)

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>

.github/workflows/main.yml

@@ -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:
@@ -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:
@@ -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:
@@ -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

README.md

@@ -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)
@@ -46,6 +47,8 @@ $ pip install dbt-tidb
 |       ✅        |       ✅       |        ✅        |      Custom data tests      |
 |       ✅        |       ✅       |        ✅        |        Docs generate        |
 |       ❌        |       ❌       |        ✅        |          Snapshots          |
+|       ✅        |       ✅       |        ✅        |      Connection retry       |
+|       ✅        |       ✅       |        ✅        |            Grant            |
 
 Note: 
 
@@ -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.
@@ -72,6 +100,7 @@ your_profile_name:
       schema: database_name
       username: tidb_username
       password: tidb_password
+      retries: 2
 ```
 
 | Option           | Description                                           | Required? | Example                        |
@@ -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
 
@@ -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
 

dbt/adapters/tidb/connections.py

@@ -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",
@@ -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):

dbt/include/tidb/macros/apply_grants.sql

@@ -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 %}

dbt/include/tidb/macros/materializations/incremental/incremental.sql

@@ -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:
@@ -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) }}

dbt/include/tidb/macros/materializations/snapshot/snapshot.sql

@@ -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) %}
@@ -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) }}

dbt/include/tidb/macros/materializations/snapshot/strategies.sql

@@ -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 %}

dbt/include/tidb/macros/utils/bool_or.sql

@@ -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 %}

dbt/include/tidb/macros/utils/cast_bool_to_text.sql

@@ -0,0 +1,8 @@
+{% macro tidb__cast_bool_to_text(field) %}
+
+    case {{ field }}
+        when true then 'true'
+        when false then 'false'
+    end
+
+{% endmacro %}

dbt/include/tidb/macros/utils/date_trunc.sql

@@ -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 %}

dbt/include/tidb/macros/utils/dateadd.sql

@@ -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 %}

dbt/include/tidb/macros/utils/datediff.sql

@@ -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 %}

dbt/include/tidb/macros/utils/hash.sql

@@ -0,0 +1,5 @@
+{% macro tidb__hash(field) -%}
+
+    md5(cast({{ field }} as CHAR))
+
+{%- endmacro %}