Skip to content

Commit

Permalink
feat: add placeholder to connection
Browse files Browse the repository at this point in the history
desc: Modify the odbc connection to add default placeholder when creating a connection in the UI by specifying it in the extra options.
  • Loading branch information
rdb1988 committed Dec 1, 2023
1 parent 0e157b3 commit fa5b034
Show file tree
Hide file tree
Showing 5 changed files with 16 additions and 8 deletions.
12 changes: 8 additions & 4 deletions airflow/providers/common/sql/hooks/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@ class DbApiHook(BaseHook):
# Override with db-specific query to check connection
_test_connection_sql = "select 1"
# Override with the db-specific value used for placeholders
placeholder: str = "%s"
_placeholder: str = "%s"

def __init__(self, *args, schema: str | None = None, log_sql: bool = True, **kwargs):
super().__init__()
Expand All @@ -164,6 +164,10 @@ def __init__(self, *args, schema: str | None = None, log_sql: bool = True, **kwa
self.log_sql = log_sql
self.descriptions: list[Sequence[Sequence] | None] = []

@property
def placeholder(self) -> str:
return self._placeholder

def get_conn(self):
"""Return a connection object."""
db = self.get_connection(getattr(self, cast(str, self.conn_name_attr)))
Expand Down Expand Up @@ -463,8 +467,8 @@ def get_cursor(self):
"""Return a cursor."""
return self.get_conn().cursor()

@classmethod
def _generate_insert_sql(cls, table, values, target_fields, replace, **kwargs) -> str:

def _generate_insert_sql(self, table, values, target_fields, replace, **kwargs) -> str:
"""
Generate the INSERT SQL statement.
Expand All @@ -477,7 +481,7 @@ def _generate_insert_sql(cls, table, values, target_fields, replace, **kwargs) -
:return: The generated INSERT or REPLACE SQL statement
"""
placeholders = [
cls.placeholder,
self.placeholder,
] * len(values)

if target_fields:
Expand Down
6 changes: 5 additions & 1 deletion airflow/providers/odbc/hooks/odbc.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,7 @@ def odbc_connection_string(self):
if self.connection.port:
conn_str += f"PORT={self.connection.port};"

extra_exclude = {"driver", "dsn", "connect_kwargs", "sqlalchemy_scheme"}
extra_exclude = {"driver", "dsn", "connect_kwargs", "sqlalchemy_scheme", "placeholder"}
extra_params = {
k: v for k, v in self.connection.extra_dejson.items() if k.lower() not in extra_exclude
}
Expand Down Expand Up @@ -198,6 +198,10 @@ def get_conn(self) -> pyodbc.Connection:
conn = pyodbc.connect(self.odbc_connection_string, **self.connect_kwargs)
return conn

@property
def placeholder(self):
return self.connection.extra_dejson.get("placeholder") or self._placeholder

def get_uri(self) -> str:
"""URI invoked in :meth:`~airflow.providers.common.sql.hooks.sql.DbApiHook.get_sqlalchemy_engine`."""
quoted_conn_str = quote_plus(self.odbc_connection_string)
Expand Down
2 changes: 1 addition & 1 deletion airflow/providers/presto/hooks/presto.py
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ class PrestoHook(DbApiHook):
default_conn_name = "presto_default"
conn_type = "presto"
hook_name = "Presto"
placeholder = "?"
_placeholder = "?"

def get_conn(self) -> Connection:
"""Returns a connection object."""
Expand Down
2 changes: 1 addition & 1 deletion airflow/providers/sqlite/hooks/sqlite.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ class SqliteHook(DbApiHook):
default_conn_name = "sqlite_default"
conn_type = "sqlite"
hook_name = "Sqlite"
placeholder = "?"
_placeholder = "?"

def get_conn(self) -> sqlite3.dbapi2.Connection:
"""Returns a sqlite connection object."""
Expand Down
2 changes: 1 addition & 1 deletion airflow/providers/trino/hooks/trino.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ class TrinoHook(DbApiHook):
conn_type = "trino"
hook_name = "Trino"
query_id = ""
placeholder = "?"
_placeholder = "?"
_test_connection_sql = "select 1"

def get_conn(self) -> Connection:
Expand Down

0 comments on commit fa5b034

Please sign in to comment.