def test_cloudsql_hook_delete_connection_on_exception( self, get_connections, run, get_connection, delete_connection): connection = Connection() connection.parse_from_uri( "gcpcloudsql://*****:*****@127.0.0.1:3200/testdb?database_type=mysql&" "project_id=example-project&location=europe-west1&instance=testdb&" "use_proxy=False") get_connection.return_value = connection db_connection = Connection() db_connection.host = "127.0.0.1" db_connection.set_extra( json.dumps({ "project_id": "example-project", "location": "europe-west1", "instance": "testdb", "database_type": "mysql" })) get_connections.return_value = [db_connection] run.side_effect = Exception("Exception when running a query") operator = CloudSqlQueryOperator(sql=['SELECT * FROM TABLE'], task_id='task_id') with self.assertRaises(Exception) as cm: operator.execute(None) err = cm.exception self.assertEqual("Exception when running a query", str(err)) delete_connection.assert_called_once_with()
def add_gcp_connection(ds, **kwargs): """"Add a airflow connection for GCP""" new_conn = Connection( conn_id='bigquery', conn_type='google_cloud_platform', ) scopes = [ "https://www.googleapis.com/auth/bigquery", ] conn_extra = { "extra__google_cloud_platform__scope": ",".join(scopes), "extra__google_cloud_platform__project": os.environ['GOOGLE_CLOUD_PROJECT'], "extra__google_cloud_platform__key_path": os.environ['AIRFLOW_CONN_GOOGLE_CLOUD_PLATFORM'] } conn_extra_json = json.dumps(conn_extra) new_conn.set_extra(conn_extra_json) session = settings.Session() if not (session.query(Connection).filter( Connection.conn_id == new_conn.conn_id).first()): session.add(new_conn) session.commit() else: msg = '\n\tA connection with `conn_id`={conn_id} already exists\n' msg = msg.format(conn_id=new_conn.conn_id) print(msg)
def add_gcp_connection(ds, **kwargs): """"Add a airflow connection for GCP""" new_conn = Connection( conn_id=CONN_PARAMS_DICT["gcp_conn_id"], conn_type="google_cloud_platform", ) scopes = [ "https://www.googleapis.com/auth/cloud-platform", ] conn_extra = { "extra__google_cloud_platform__scope": ",".join(scopes), "extra__google_cloud_platform__project": CONN_PARAMS_DICT.get("gcp_project"), "extra__google_cloud_platform__keyfile_dict": get_secret( project_name=CONN_PARAMS_DICT.get("gcp_project"), secret_name=CONN_PARAMS_DICT.get("secret_name"), ), } conn_extra_json = json.dumps(conn_extra) new_conn.set_extra(conn_extra_json) session = settings.Session() if not (session.query(Connection).filter(Connection.conn_id == new_conn.conn_id).first()): session.add(new_conn) session.commit() msg = "\n\tA connection with `conn_id`={conn_id} is newly created\n" msg = msg.format(conn_id=new_conn.conn_id) print(msg) else: msg = "\n\tA connection with `conn_id`={conn_id} already exists\n" msg = msg.format(conn_id=new_conn.conn_id) print(msg)
def test_cloudsql_hook_delete_connection_on_exception( self, get_connections, run, get_connection, delete_connection): connection = Connection() connection.parse_from_uri( "gcpcloudsql://*****:*****@127.0.0.1:3200/testdb?database_type=mysql&" "project_id=example-project&location=europe-west1&instance=testdb&" "use_proxy=False") get_connection.return_value = connection db_connection = Connection() db_connection.host = "127.0.0.1" db_connection.set_extra(json.dumps({"project_id": "example-project", "location": "europe-west1", "instance": "testdb", "database_type": "mysql"})) get_connections.return_value = [db_connection] run.side_effect = Exception("Exception when running a query") operator = CloudSqlQueryOperator( sql=['SELECT * FROM TABLE'], task_id='task_id' ) with self.assertRaises(Exception) as cm: operator.execute(None) err = cm.exception self.assertEqual("Exception when running a query", str(err)) delete_connection.assert_called_once_with()
def add_gcp_connection(ds, **kwargs): """"Add a airflow connection for GCP""" new_conn = Connection( conn_id="my_gcp_connection", # TODO: parameterize conn_type="google_cloud_platform", ) scopes = [ "https://www.googleapis.com/auth/cloud-platform", ] conn_extra = { "extra__google_cloud_platform__scope": ",".join(scopes), "extra__google_cloud_platform__project": "wam-bam-258119", # TODO: parameterize "extra__google_cloud_platform__key_path": "service_account.json", # TODO: parameterize } conn_extra_json = json.dumps(conn_extra) new_conn.set_extra(conn_extra_json) session = settings.Session() if not ( session.query(Connection).filter(Connection.conn_id == new_conn.conn_id).first() ): session.add(new_conn) session.commit() else: msg = "\n\tA connection with `conn_id`={conn_id} already exists\n" msg = msg.format(conn_id=new_conn.conn_id) print(msg)
def set_connection( self, conn_id: str, conn_type: str, host: Optional[str] = None, schema: Optional[str] = None, login: Optional[str] = None, password: Optional[str] = None, port: Optional[int] = None, extra: Optional[Union[str, dict]] = None, ): assert str(settings.engine.url) == self.sql_alchemy_conn session = settings.Session() try: new_conn = Connection(conn_id=conn_id, conn_type=conn_type, host=host, login=login, password=password, schema=schema, port=port) if extra is not None: new_conn.set_extra( extra if isinstance(extra, str) else json.dumps(extra)) session.add(new_conn) session.commit() finally: session.close()
def add_gcp_connection(ds, **kwargs): """"Add a airflow connection for GCP""" json_data = json.loads(open('/usr/local/airflow/gcp.json').read()) new_conn = Connection( conn_id='AirflowGCPKey', conn_type='google_cloud_platform', ) scopes = [ "https://www.googleapis.com/auth/pubsub", "https://www.googleapis.com/auth/datastore", "https://www.googleapis.com/auth/bigquery", "https://www.googleapis.com/auth/devstorage.read_write", "https://www.googleapis.com/auth/logging.write", "https://www.googleapis.com/auth/cloud-platform", ] conn_extra = { "extra__google_cloud_platform__scope": ",".join(scopes), "extra__google_cloud_platform__project": "<insert your project platform>", "extra__google_cloud_platform__key_path": '/usr/local/airflow/gcp.json', "extra__google_cloud_platform__key_json": json_data } conn_extra_json = json.dumps(conn_extra) new_conn.set_extra(conn_extra_json) session = settings.Session() if not (session.query(Connection).filter( Connection.conn_id == new_conn.conn_id).first()): session.add(new_conn) session.commit() else: msg = '\n\tA connection with `conn_id`={conn_id} already exists\n' msg = msg.format(conn_id=new_conn.conn_id) print(msg)
def add_gcp_connection(**kwargs): new_conn = Connection( conn_id="google_cloud_default", #conn_id="airflowgcp_connection", conn_type='google_cloud_platform', ) extra_field = { "extra__google_cloud_platform__scope": "https://www.googleapis.com/auth/cloud-platform", "extra__google_cloud_platform__project": "airflowgcp", "extra__google_cloud_platform__key_path": '/usr/local/airflow/dags/gcpLicense.json' } session = settings.Session() #checking if connection exist if session.query(Connection).filter( Connection.conn_id == new_conn.conn_id).first(): my_connection = session.query(Connection).filter( Connection.conn_id == new_conn.conn_id).one() my_connection.set_extra(json.dumps(extra_field)) session.add(my_connection) session.commit() else: #if it doesn't exit create one new_conn.set_extra(json.dumps(extra_field)) session.add(new_conn) session.commit()
def add_gcp_connection(**kwargs): new_conn = Connection( conn_id="google_cloud_default", conn_type='google_cloud_platform', ) extra_field = { "extra__google_cloud_platform__scope": "https://www.googleapis.com/auth/cloud-platform", "extra__google_cloud_platform__project": "task-ts", "extra__google_cloud_platform__key_path": os.environ["GCP_KEY_PATH"] } session = settings.Session() #checking if connection exist if session.query(Connection).filter( Connection.conn_id == new_conn.conn_id).first(): my_connection = session.query(Connection).filter( Connection.conn_id == new_conn.conn_id).one() my_connection.set_extra(json.dumps(extra_field)) session.add(my_connection) session.commit() else: #if it doesn't exit create one new_conn.set_extra(json.dumps(extra_field)) session.add(new_conn) session.commit()
def add_connection( conn_id, uri=None, conn_type=None, host=None, login=None, password=None, port=None, schema=None, extra=None, ): if uri: new_conn = Connection(conn_id=conn_id, uri=uri) else: new_conn = Connection( conn_id=conn_id, conn_type=conn_type, host=host, login=login, password=password, schema=schema, port=port, ) if extra is not None: new_conn.set_extra(extra) session = settings.Session() if not (session.query(Connection).filter( Connection.conn_id == new_conn.conn_id).first()): session.add(new_conn) session.commit() msg = "\n\tSuccessfully added `conn_id`={conn_id} : {uri}\n" msg = msg.format( conn_id=new_conn.conn_id, uri=uri or urlunparse(( conn_type, "{login}:{password}@{host}:{port}".format( login=login or "", password=password or "", host=host or "", port=port or "", ), schema or "", "", "", "", )), ) logger.info(msg) logger.info("extra=%s" % str(extra)) else: msg = "\n\tA connection with `conn_id`={conn_id} already exists\n" raise DatabandRuntimeError(msg.format(conn_id=new_conn.conn_id))
def connections_add(args): """Adds new connection""" # Check that the conn_id and conn_uri args were passed to the command: missing_args = [] invalid_args = [] if args.conn_uri: for arg in alternative_conn_specs: if getattr(args, arg) is not None: invalid_args.append(arg) elif not args.conn_type: missing_args.append('conn-uri or conn-type') if missing_args: msg = ('The following args are required to add a connection:' + ' {missing!r}'.format(missing=missing_args)) raise SystemExit(msg) if invalid_args: msg = ('The following args are not compatible with the ' + 'add flag and --conn-uri flag: {invalid!r}') msg = msg.format(invalid=invalid_args) raise SystemExit(msg) if args.conn_uri: new_conn = Connection(conn_id=args.conn_id, uri=args.conn_uri) else: new_conn = Connection(conn_id=args.conn_id, conn_type=args.conn_type, host=args.conn_host, login=args.conn_login, password=args.conn_password, schema=args.conn_schema, port=args.conn_port) if args.conn_extra is not None: new_conn.set_extra(args.conn_extra) with create_session() as session: if not (session.query(Connection) .filter(Connection.conn_id == new_conn.conn_id).first()): session.add(new_conn) msg = '\n\tSuccessfully added `conn_id`={conn_id} : {uri}\n' msg = msg.format(conn_id=new_conn.conn_id, uri=args.conn_uri or urlunparse((args.conn_type, '{login}:{password}@{host}:{port}' .format(login=args.conn_login or '', password='******' if args.conn_password else '', host=args.conn_host or '', port=args.conn_port or ''), args.conn_schema or '', '', '', ''))) print(msg) else: msg = '\n\tA connection with `conn_id`={conn_id} already exists\n' msg = msg.format(conn_id=new_conn.conn_id) print(msg)
def connections_add(args): """Adds new connection""" # Check that the conn_id and conn_uri args were passed to the command: missing_args = [] invalid_args = [] if args.conn_uri: if not _valid_uri(args.conn_uri): raise SystemExit(f'The URI provided to --conn-uri is invalid: {args.conn_uri}') for arg in alternative_conn_specs: if getattr(args, arg) is not None: invalid_args.append(arg) elif not args.conn_type: missing_args.append('conn-uri or conn-type') if missing_args: raise SystemExit(f'The following args are required to add a connection: {missing_args!r}') if invalid_args: raise SystemExit( f'The following args are not compatible with the ' f'add flag and --conn-uri flag: {invalid_args!r}' ) if args.conn_uri: new_conn = Connection(conn_id=args.conn_id, description=args.conn_description, uri=args.conn_uri) else: new_conn = Connection( conn_id=args.conn_id, conn_type=args.conn_type, description=args.conn_description, host=args.conn_host, login=args.conn_login, password=args.conn_password, schema=args.conn_schema, port=args.conn_port, ) if args.conn_extra is not None: new_conn.set_extra(args.conn_extra) with create_session() as session: if not session.query(Connection).filter(Connection.conn_id == new_conn.conn_id).first(): session.add(new_conn) msg = 'Successfully added `conn_id`={conn_id} : {uri}' msg = msg.format( conn_id=new_conn.conn_id, uri=args.conn_uri or urlunparse( ( args.conn_type, '{login}:{password}@{host}:{port}'.format( login=args.conn_login or '', password='******' if args.conn_password else '', host=args.conn_host or '', port=args.conn_port or '', ), args.conn_schema or '', '', '', '', ) ), ) print(msg) else: msg = f'A connection with `conn_id`={new_conn.conn_id} already exists.' raise SystemExit(msg)
def connections_add(args): """Adds new connection""" has_uri = bool(args.conn_uri) has_json = bool(args.conn_json) has_type = bool(args.conn_type) if not has_type and not (has_json or has_uri): raise SystemExit( 'Must supply either conn-uri or conn-json if not supplying conn-type' ) if has_json and has_uri: raise SystemExit('Cannot supply both conn-uri and conn-json') if has_type and not (args.conn_type in _get_connection_types()): warnings.warn( f'The type provided to --conn-type is invalid: {args.conn_type}') warnings.warn(f'Supported --conn-types are:{_get_connection_types()}.' 'Hence overriding the conn-type with generic') args.conn_type = 'generic' if has_uri or has_json: invalid_args = [] if has_uri and not _valid_uri(args.conn_uri): raise SystemExit( f'The URI provided to --conn-uri is invalid: {args.conn_uri}') for arg in alternative_conn_specs: if getattr(args, arg) is not None: invalid_args.append(arg) if has_json and args.conn_extra: invalid_args.append("--conn-extra") if invalid_args: raise SystemExit( "The following args are not compatible with " f"the --conn-{'uri' if has_uri else 'json'} flag: {invalid_args!r}" ) if args.conn_uri: new_conn = Connection(conn_id=args.conn_id, description=args.conn_description, uri=args.conn_uri) if args.conn_extra is not None: new_conn.set_extra(args.conn_extra) elif args.conn_json: new_conn = Connection.from_json(conn_id=args.conn_id, value=args.conn_json) if not new_conn.conn_type: raise SystemExit('conn-json is invalid; must supply conn-type') else: new_conn = Connection( conn_id=args.conn_id, conn_type=args.conn_type, description=args.conn_description, host=args.conn_host, login=args.conn_login, password=args.conn_password, schema=args.conn_schema, port=args.conn_port, ) if args.conn_extra is not None: new_conn.set_extra(args.conn_extra) with create_session() as session: if not session.query(Connection).filter( Connection.conn_id == new_conn.conn_id).first(): session.add(new_conn) msg = 'Successfully added `conn_id`={conn_id} : {uri}' msg = msg.format( conn_id=new_conn.conn_id, uri=args.conn_uri or urlunparse(( new_conn.conn_type, f"{new_conn.login or ''}:{'******' if new_conn.password else ''}" f"@{new_conn.host or ''}:{new_conn.port or ''}", new_conn.schema or '', '', '', '', )), ) print(msg) else: msg = f'A connection with `conn_id`={new_conn.conn_id} already exists.' raise SystemExit(msg)