[AIRFLOW-858] - Configurable database name for DB operators

airflow/hooks/mssql_hook.py

@@ -18,14 +18,18 @@
 
 
 class MsSqlHook(DbApiHook):
-    '''
+    """
     Interact with Microsoft SQL Server.
-    '''
+    """
 
     conn_name_attr = 'mssql_conn_id'
     default_conn_name = 'mssql_default'
     supports_autocommit = True
 
+    def __init__(self, *args, **kwargs):
+        super(MsSqlHook, self).__init__(*args, **kwargs)
+        self.schema = kwargs.pop("schema", None)
+
     def get_conn(self):
         """
         Returns a mssql connection object
@@ -35,7 +39,7 @@ def get_conn(self):
             server=conn.host,
             user=conn.login,
             password=conn.password,
-            database=conn.schema,
+            database=self.schema or conn.schema,
             port=conn.port)
         return conn
 

airflow/hooks/mysql_hook.py

@@ -19,36 +19,39 @@
 
 
 class MySqlHook(DbApiHook):
-    '''
+    """
     Interact with MySQL.
 
     You can specify charset in the extra field of your connection
     as ``{"charset": "utf8"}``. Also you can choose cursor as
     ``{"cursor": "SSCursor"}``. Refer to the MySQLdb.cursors for more details.
-    '''
+    """
 
     conn_name_attr = 'mysql_conn_id'
     default_conn_name = 'mysql_default'
     supports_autocommit = True
 
+    def __init__(self, *args, **kwargs):
+        super(MySqlHook, self).__init__(*args, **kwargs)
+        self.schema = kwargs.pop("schema", None)
+
     def get_conn(self):
         """
         Returns a mysql connection object
         """
         conn = self.get_connection(self.mysql_conn_id)
         conn_config = {
             "user": conn.login,
-            "passwd": conn.password or ''
+            "passwd": conn.password or '',
+            "host": conn.host or 'localhost',
+            "db": self.schema or conn.schema or ''
         }
 
-        conn_config["host"] = conn.host or 'localhost'
         if not conn.port:
             conn_config["port"] = 3306
         else:
             conn_config["port"] = int(conn.port)
 
-        conn_config["db"] = conn.schema or ''
-
         if conn.extra_dejson.get('charset', False):
             conn_config["charset"] = conn.extra_dejson["charset"]
             if (conn_config["charset"]).lower() == 'utf8' or\

airflow/hooks/postgres_hook.py

@@ -19,22 +19,26 @@
 
 
 class PostgresHook(DbApiHook):
-    '''
+    """
     Interact with Postgres.
     You can specify ssl parameters in the extra field of your connection
     as ``{"sslmode": "require", "sslcert": "/path/to/cert.pem", etc}``.
-    '''
+    """
     conn_name_attr = 'postgres_conn_id'
     default_conn_name = 'postgres_default'
     supports_autocommit = True
 
+    def __init__(self, *args, **kwargs):
+        super(PostgresHook, self).__init__(*args, **kwargs)
+        self.schema = kwargs.pop("schema", None)
+
     def get_conn(self):
         conn = self.get_connection(self.postgres_conn_id)
         conn_args = dict(
             host=conn.host,
             user=conn.login,
             password=conn.password,
-            dbname=conn.schema,
+            dbname=self.schema or conn.schema,
             port=conn.port)
         # check for ssl parameters in conn.extra
         for arg_name, arg_val in conn.extra_dejson.items():

airflow/operators/mssql_operator.py

@@ -27,6 +27,8 @@ class MsSqlOperator(BaseOperator):
     :param sql: the sql code to be executed
     :type sql: string or string pointing to a template file.
     File must have a '.sql' extensions.
+    :param database: name of database which overwrite defined one in connection
+    :type database: string
     """
 
     template_fields = ('sql',)
@@ -36,14 +38,17 @@ class MsSqlOperator(BaseOperator):
     @apply_defaults
     def __init__(
             self, sql, mssql_conn_id='mssql_default', parameters=None,
-            autocommit=False, *args, **kwargs):
+            autocommit=False, database=None, *args, **kwargs):
         super(MsSqlOperator, self).__init__(*args, **kwargs)
         self.mssql_conn_id = mssql_conn_id
         self.sql = sql
         self.parameters = parameters
         self.autocommit = autocommit
+        self.database = database
 
     def execute(self, context):
         logging.info('Executing: ' + str(self.sql))
-        hook = MsSqlHook(mssql_conn_id=self.mssql_conn_id)
-        hook.run(self.sql, autocommit=self.autocommit, parameters=self.parameters)
+        hook = MsSqlHook(mssql_conn_id=self.mssql_conn_id,
+                         schema=self.database)
+        hook.run(self.sql, autocommit=self.autocommit,
+                 parameters=self.parameters)

airflow/operators/mysql_operator.py

@@ -29,6 +29,8 @@ class MySqlOperator(BaseOperator):
     :type sql: Can receive a str representing a sql statement,
         a list of str (sql statements), or reference to a template file.
         Template reference are recognized by str ending in '.sql'
+    :param database: name of database which overwrite defined one in connection
+    :type database: string
     """
 
     template_fields = ('sql',)
@@ -38,16 +40,18 @@ class MySqlOperator(BaseOperator):
     @apply_defaults
     def __init__(
             self, sql, mysql_conn_id='mysql_default', parameters=None,
-            autocommit=False, *args, **kwargs):
+            autocommit=False, database=None, *args, **kwargs):
         super(MySqlOperator, self).__init__(*args, **kwargs)
         self.mysql_conn_id = mysql_conn_id
         self.sql = sql
         self.autocommit = autocommit
         self.parameters = parameters
+        self.database = database
 
     def execute(self, context):
         logging.info('Executing: ' + str(self.sql))
-        hook = MySqlHook(mysql_conn_id=self.mysql_conn_id)
+        hook = MySqlHook(mysql_conn_id=self.mysql_conn_id,
+                         schema=self.database)
         hook.run(
             self.sql,
             autocommit=self.autocommit,

airflow/operators/postgres_operator.py

@@ -29,6 +29,8 @@ class PostgresOperator(BaseOperator):
     :type sql: Can receive a str representing a sql statement,
         a list of str (sql statements), or reference to a template file.
         Template reference are recognized by str ending in '.sql'
+    :param database: name of database which overwrite defined one in connection
+    :type database: string
     """
 
     template_fields = ('sql',)
@@ -40,14 +42,17 @@ def __init__(
             self, sql,
             postgres_conn_id='postgres_default', autocommit=False,
             parameters=None,
+            database=None,
             *args, **kwargs):
         super(PostgresOperator, self).__init__(*args, **kwargs)
         self.sql = sql
         self.postgres_conn_id = postgres_conn_id
         self.autocommit = autocommit
         self.parameters = parameters
+        self.database = database
 
     def execute(self, context):
         logging.info('Executing: ' + str(self.sql))
-        self.hook = PostgresHook(postgres_conn_id=self.postgres_conn_id)
+        self.hook = PostgresHook(postgres_conn_id=self.postgres_conn_id,
+                                 schema=self.database)
         self.hook.run(self.sql, self.autocommit, parameters=self.parameters)

tests/operators/operators.py

@@ -114,6 +114,27 @@ def test_sql_sensor(self):
             dag=self.dag)
         t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
 
+    def test_overwrite_schema(self):
+        """
+        Verifies option to overwrite connection schema
+        """
+        import airflow.operators.mysql_operator
+
+        sql = "SELECT 1;"
+        t = operators.mysql_operator.MySqlOperator(
+            task_id='test_mysql_operator_test_schema_overwrite',
+            sql=sql,
+            dag=self.dag,
+            database="foobar",
+        )
+
+        from _mysql_exceptions import OperationalError
+        try:
+            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
+                  ignore_ti_state=True)
+        except OperationalError as e:
+            assert "Unknown database 'foobar'" in str(e)
+
 
 @skipUnlessImported('airflow.operators.postgres_operator', 'PostgresOperator')
 class PostgresTest(unittest.TestCase):
@@ -193,6 +214,28 @@ def test_vacuum(self):
             autocommit=True)
         t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
 
+    def test_overwrite_schema(self):
+        """
+        Verifies option to overwrite connection schema
+        """
+        import airflow.operators.postgres_operator
+
+        sql = "SELECT 1;"
+        t = operators.postgres_operator.PostgresOperator(
+            task_id='postgres_operator_test_schema_overwrite',
+            sql=sql,
+            dag=self.dag,
+            autocommit=True,
+            database="foobar",
+        )
+
+        from psycopg2._psycopg import OperationalError
+        try:
+            t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
+                  ignore_ti_state=True)
+        except OperationalError as e:
+            assert 'database "foobar" does not exist' in str(e)
+
 
 @skipUnlessImported('airflow.operators.hive_operator', 'HiveOperator')
 @skipUnlessImported('airflow.operators.postgres_operator', 'PostgresOperator')