Skip to content

Commit

Permalink
add db util to prepand/append comma delimited SQL comment from key/va…
Browse files Browse the repository at this point in the history
…lue pairs (DataDog#17145)

* add db util to prepand/append sql comment from dict

* add changelog

* fix for py2

* use quote from six

* do not quote

* change to kwargs
  • Loading branch information
lu-zhengda authored Mar 12, 2024
1 parent 79fbbbc commit 457a95a
Show file tree
Hide file tree
Showing 3 changed files with 82 additions and 0 deletions.
1 change: 1 addition & 0 deletions datadog_checks_base/changelog.d/17145.added
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add db util to prepand/append comma delimited SQL comment from key/value pairs
47 changes: 47 additions & 0 deletions datadog_checks_base/datadog_checks/base/utils/db/sql_commenter.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
# (C) Datadog, Inc. 2024-present
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)
def generate_sql_comment(**kwargs):
'''
Generate a SQL comment from a dictionary of key-value pairs.
Returns an empty string if the input dictionary is empty or if all values are empty.
Example:
>>> generate_sql_comment(foo='bar', baz='qux')
"/* baz='qux',foo='bar' */"
'''
if not kwargs:
return ''

comment = ','.join(
"{}='{}'".format(key, value) for key, value in kwargs.items() if value is not None and value != ''
)

return '/* {} */'.format(comment)


def add_sql_comment(sql, prepand=True, **kwargs):
'''
Prepand or append a SQL comment to a SQL statement from a dictionary of key-value pairs.
Returns the original SQL statement if the input dictionary is empty or if all values are empty.
Example:
>>> add_sql_comment('SELECT * FROM table', foo='bar', baz='qux')
"/* baz='qux',foo='bar' */ SELECT * FROM table"
>>> add_sql_comment('SELECT * FROM table', False, foo='bar', baz='qux')
"SELECT * FROM table /* baz='qux',foo='bar' */"
'''
comment = generate_sql_comment(**kwargs)
if not comment:
return sql

sql = sql.strip()
if not sql:
return sql

if prepand:
sql = '{} {}'.format(comment, sql)
else:
if sql[-1] == ';':
sql = '{} {};'.format(sql[:-1], comment)
else:
sql = '{} {}'.format(sql, comment)
return sql
34 changes: 34 additions & 0 deletions datadog_checks_base/tests/base/utils/db/test_sql_commenter.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
# (C) Datadog, Inc. 2024-present
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)
import pytest

from datadog_checks.base.utils.db.sql_commenter import add_sql_comment, generate_sql_comment


class TestSQLCommenter:
def test_generate_sql_comment(self):
assert generate_sql_comment(foo='bar', baz='qux') == "/* foo='bar',baz='qux' */"

def test_generate_sql_comment_empty(self):
assert generate_sql_comment() == ""

def test_generate_sql_comment_none(self):
assert generate_sql_comment(foo=None, baz='qux') == "/* baz='qux' */"

@pytest.mark.parametrize(
'sql,prepand,expected',
[
pytest.param('SELECT * FROM table', True, "/* foo='bar',baz='qux' */ SELECT * FROM table", id='prepand'),
pytest.param('SELECT * FROM table', False, "SELECT * FROM table /* foo='bar',baz='qux' */", id='append'),
pytest.param(
'SELECT * FROM table;', False, "SELECT * FROM table /* foo='bar',baz='qux' */;", id='append_semicolon'
),
pytest.param('', True, '', id='empty_sql'),
],
)
def test_add_sql_comment(self, sql, prepand, expected):
assert add_sql_comment(sql, prepand, foo='bar', baz='qux') == expected

def test_add_sql_comment_empty(self):
assert add_sql_comment('SELECT * FROM table') == 'SELECT * FROM table'

0 comments on commit 457a95a

Please sign in to comment.