Esempio n. 1
0
    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)
Esempio n. 3
0
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()
Esempio n. 5
0
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()
Esempio n. 7
0
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()
Esempio n. 9
0
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()
Esempio n. 10
0
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))
Esempio n. 11
0
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)
Esempio n. 13
0
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)