[GitHub] [airflow] xinbinhuang commented on a change in pull request #12677: Refactor SQL/BigQuery/Qubole Check operators

[GitBox <gi...@apache.org>]

xinbinhuang commented on a change in pull request #12677:
URL: https://github.com/apache/airflow/pull/12677#discussion_r551071034



##########
File path: airflow/operators/sql.py
##########
@@ -38,7 +41,53 @@
 }
 
 
-class SQLCheckOperator(BaseOperator):
+class BaseSQLOperator(BaseOperator):
+    """
+    This is a base class for generic SQL Operator to get a DB Hook
+
+    The provided method is .get_db_hook(). The default behavior will try to
+    retrieve the DB hook based on connection type.
+    You can custom the behavior by overriding the .get_db_hook() method.
+    """
+
+    @apply_defaults
+    def __init__(self, *, conn_id: Optional[str] = None, database: Optional[str] = None, **kwargs):
+        super().__init__(**kwargs)
+        self.conn_id = conn_id
+        self.database = database
+
+    @cached_property
+    def _hook(self):
+        """Get DB Hook based on connection type"""
+        self.log.debug("Get connection for %s", self.conn_id)
+        conn = BaseHook.get_connection(self.conn_id)
+
+        if conn.conn_type not in ALLOWED_CONN_TYPE:

Review comment:
       I think the reason for `ALLOWED_CONN_TYPE` is to only whitelist a subset of all connection types that are `DbApiHook` compatible. Maybe we can do something like this instead?
   
   
   ```python
   hook = conn.get_hook()
   if not isinstance(hook, DbApiHook):
       raise AirflowException('The connection type associated hook is not a subclass of `DbApiHook`') 
   ```
   
   However, there are edge cases that need to be resolved in future PRs to make this work.
   - Currently, bigQuery seems to be not discoverable because it overlaps with the GoogleBaseHook: https://github.com/apache/airflow/blob/5f81fc73c8ea3fc1c3b08080f439fa123926f250/airflow/providers/google/cloud/hooks/bigquery.py#L72
   - CloudSQLHook/CloudSQLDatabaseHook is not a subclass of `DbApiHook`
   https://github.com/apache/airflow/blob/5f81fc73c8ea3fc1c3b08080f439fa123926f250/airflow/providers/google/cloud/hooks/cloud_sql.py#L666
   - ... not sure if there are other edge cases




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org