Merge branch 'dbt_1.1' into dbt_1.2

pingcap · Aug 22, 2022 · 05a1f38 · 05a1f38
2 parents d64b591 + fcdd5d1
commit 05a1f38
Show file tree

Hide file tree

Showing 55 changed files with 275 additions and 1,731 deletions.
diff --git a/.github/workflows/fmt.yml b/.github/workflows/fmt.yml
@@ -0,0 +1,10 @@
+name: Lint
+
+on: [push, pull_request]
+
+jobs:
+  fmt:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: psf/black@stable
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
@@ -35,10 +35,6 @@ 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:
@@ -70,10 +66,6 @@ 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:
@@ -105,10 +97,6 @@ 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:
@@ -140,8 +128,4 @@ 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
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -21,4 +21,19 @@ We love to accept any kinds of contributions, such as:
 5. Open a pull request to upstream
 6. Request a review from the reviewers
 
+## How to formate your code
+
+dbt-tidb use black to fmt Python. To fmt your code, you need to:
+
+First, install black
+```
+pip3 install black
+```
+Then fmt
+```
+black .
+```
+
+You can also integrate with IntelliJ IDEA or other editors. See [here](https://black.readthedocs.io/en/stable/integrations/editors.html) for more details.
+
 
diff --git a/README.md b/README.md
@@ -14,7 +14,6 @@ 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)
@@ -47,8 +46,6 @@ $ pip install dbt-tidb
 |       ✅        |       ✅       |        ✅        |      Custom data tests      |
 |       ✅        |       ✅       |        ✅        |        Docs generate        |
 |       ❌        |       ❌       |        ✅        |          Snapshots          |
-|       ✅        |       ✅       |        ✅        |      Connection retry       |
-|       ✅        |       ✅       |        ✅        |            Grant            |
 
 Note: 
 
@@ -58,31 +55,6 @@ 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.
@@ -100,7 +72,6 @@ your_profile_name:
       schema: database_name
       username: tidb_username
       password: tidb_password
-      retries: 2
 ```
 
 | Option           | Description                                           | Required? | Example                        |
@@ -111,7 +82,6 @@ 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
 

diff --git a/dbt/adapters/tidb/__init__.py b/dbt/adapters/tidb/__init__.py
@@ -9,6 +9,5 @@
 
 
 Plugin = AdapterPlugin(
-    adapter=TiDBAdapter,
-    credentials=TiDBCredentials,
-    include_path=tidb.PACKAGE_PATH)
+    adapter=TiDBAdapter, credentials=TiDBCredentials, include_path=tidb.PACKAGE_PATH
+)
diff --git a/dbt/adapters/tidb/column.py b/dbt/adapters/tidb/column.py
@@ -3,7 +3,7 @@
 
 from dbt.adapters.base.column import Column
 
-Self = TypeVar('Self', bound='TiDBColumn')
+Self = TypeVar("Self", bound="TiDBColumn")
 
 
 @dataclass
@@ -18,7 +18,7 @@ class TiDBColumn(Column):
 
     @property
     def quoted(self) -> str:
-        return '`{}`'.format(self.column)
+        return "`{}`".format(self.column)
 
     def __repr__(self) -> str:
         return "<TiDBColumn {} ({})>".format(self.name, self.data_type)
diff --git a/dbt/adapters/tidb/connections.py b/dbt/adapters/tidb/connections.py
@@ -23,7 +23,6 @@ class TiDBCredentials(Credentials):
     username: Optional[str] = None
     password: Optional[str] = None
     charset: Optional[str] = None
-    retries: int = 1
 
     _ALIASES = {
         "UID": "username",
@@ -39,10 +38,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
-        ):
+        if self.database is not None and self.database != self.schema:
             raise dbt.exceptions.RuntimeException(
                 f"    schema: {self.schema} \n"
                 f"    database: {self.database} \n"
@@ -76,8 +72,8 @@ class TiDBConnectionManager(SQLConnectionManager):
 
     @classmethod
     def open(cls, connection):
-        if connection.state == 'open':
-            logger.debug('Connection is already open, skipping open.')
+        if connection.state == "open":
+            logger.debug("Connection is already open, skipping open.")
             return connection
 
         credentials = cls.get_credentials(connection.credentials)
@@ -91,22 +87,35 @@ def open(cls, connection):
         if credentials.port:
             kwargs["port"] = credentials.port
 
-        def connect():
-            handle = mysql.connector.connect(**kwargs)
-            return handle
-
-        # we just retry for any error now
-        retryable_exceptions = [
-                 mysql.connector.Error
-        ]
-
-        return cls.retry_connection(
-            connection,
-            connect=connect,
-            logger=logger,
-            retry_limit=credentials.retries,
-            retryable_exceptions=retryable_exceptions,
-        )
+        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
+
+                connection.handle = mysql.connector.connect(**kwargs)
+                connection.state = "open"
+            except mysql.connector.Error as e:
+
+                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
 
     @classmethod
     def get_credentials(cls, credentials):
@@ -121,7 +130,7 @@ def exception_handler(self, sql):
             yield
 
         except mysql.connector.DatabaseError as e:
-            logger.debug('TiDB error: {}'.format(str(e)))
+            logger.debug("TiDB error: {}".format(str(e)))
 
             try:
                 self.rollback_if_open()
@@ -154,7 +163,5 @@ def get_response(cls, cursor) -> AdapterResponse:
         # There's no real way to get the status from the mysql-connector-python driver.
         # So just return the default value.
         return AdapterResponse(
-            _message="{} {}".format(code, num_rows),
-            rows_affected=num_rows,
-            code=code
+            _message="{} {}".format(code, num_rows), rows_affected=num_rows, code=code
         )