Python sqlalchemy.sql() Examples

def _create_join_conds(left_sel, right_sel, on):
    left_cols  = left_sel.columns  #lift_inner_cols(left_sel)
    right_cols = right_sel.columns #lift_inner_cols(right_sel)

    if callable(on):
        # callable, like with sql_on arg
        conds = [on(left_cols, right_cols)]
    else:
        # dict-like of form {left: right}
        conds = []
        for l, r in on.items():
            col_expr = left_cols[l] == right_cols[r]
            conds.append(col_expr)
            
    return sql.and_(*conds)
    

# Head ------------------------------------------------------------------------ 

def test_contain_percents_character_query_with_parameter(self, engine, conn):
        query = sqlalchemy.sql.text(
            """
            SELECT date_parse('20191030', '%Y%m%d'), :word
            """
        )
        result = engine.execute(query, word="cat")
        self.assertEqual(result.fetchall(), [(datetime(2019, 10, 30), "cat")])

        query = sqlalchemy.sql.text(
            """
            SELECT col_string FROM one_row_complex
            WHERE col_string LIKE 'a%' OR col_string LIKE :param
            """
        )
        result = engine.execute(query, param="b%")
        self.assertEqual(result.fetchall(), [("a string",)]) 

def add_note(db, server, user, text):
    id_query = select([sqlalchemy.sql.expression.func.max(table.c.note_id).label("maxid")]) \
        .where(table.c.user == user.lower())
    max_id = db.execute(id_query).scalar()

    if max_id is None:
        note_id = 1
    else:
        note_id = max_id + 1

    query = table.insert().values(
        note_id=note_id,
        connection=server,
        user=user.lower(),
        text=text,
        deleted=False,
        added=datetime.today()
    )
    db.execute(query)
    db.commit() 

def generate_default_port_by_dialect(dialect: str) -> str:
    """
    In case no port was chosen, a default port will be chosen according to the SQL db type
    :param dialect: sql db type
    :return: default port needed for connection
    """
    if dialect == "MySQL":
        return "3306"
    elif dialect == "PostgreSQL":
        return "5432"
    elif dialect == "Oracle":
        return "1521"
    elif dialect == "Microsoft SQL Server":
        return "1433"
    else:
        # set default to mysql
        return "3306" 

def _validate_join_arg_on(on, sql_on = None):
    # handle sql on case
    if sql_on is not None:
        if on is not None:
            raise ValueError("Cannot specify both on and sql_on")

        return sql_on

    # handle general cases
    if on is None:
        raise NotImplementedError("on arg currently cannot be None (default) for SQL")
    elif isinstance(on, str):
        on = {on: on}
    elif isinstance(on, (list, tuple)):
        on = dict(zip(on, on))

    if not isinstance(on, Mapping):
        raise TypeError("on must be a Mapping (e.g. dict)")

    return on 

def _anti_join(left, right = None, on = None, *args, sql_on = None):
    _raise_if_args(args)

    left_sel = left.last_op.alias()
    right_sel = right.last_op.alias()

    # handle arguments ----
    on  = _validate_join_arg_on(on, sql_on)
    
    # create join conditions ----
    bool_clause = _create_join_conds(left_sel, right_sel, on)

    # create inner join ----
    not_exists = ~sql.exists([1], from_obj = right_sel).where(bool_clause)
    sel = sql.select(left_sel.columns, from_obj = left_sel).where(not_exists)
    return left.append_op(sel) 

def col_expr_requires_cte(call, sel, is_mutate = False):
    """Return whether a variable assignment needs a CTE"""

    call_vars = set(call.op_vars(attr_calls = False))

    columns = lift_inner_cols(sel)
    sel_labs = set(k for k,v in columns.items() if isinstance(v, sql.elements.Label))

    # I use the acronym fwg sol (frog soul) to remember sql clause eval order
    # from, where, group by, select, order by, limit
    # group clause evaluated before select clause, so not issue for mutate
    group_needs_cte = not is_mutate and len(sel._group_by_clause)
    
    return (   group_needs_cte
            or len(sel._order_by_clause)
            or not sel_labs.isdisjoint(call_vars)
            ) 

def test_contain_percents_character_query_with_parameter(self, engine, conn):
        query = sqlalchemy.sql.text(
            """
            SELECT date_parse('20191030', '%Y%m%d'), :word
            """
        )
        result = engine.execute(query, word="cat")
        self.assertEqual(result.fetchall(), [(datetime(2019, 10, 30), "cat")])

        query = sqlalchemy.sql.text(
            """
            SELECT col_string FROM one_row_complex
            WHERE col_string LIKE 'a%' OR col_string LIKE :param
            """
        )
        result = engine.execute(query, param="b%")
        self.assertEqual(result.fetchall(), [("a string",)]) 

def _group_by(__data, *args, add = False, **kwargs):
    if kwargs:
        data = mutate(__data, **kwargs)
    else:
        data = __data

    cols = data.last_op.columns

    # put kwarg grouping vars last, so similar order to function call
    groups =  tuple(simple_varname(arg) for arg in args) + tuple(kwargs)
    if None in groups:
        raise NotImplementedError("Complex expressions not supported in sql group_by")

    unmatched = set(groups) - set(cols.keys())
    if unmatched:
        raise KeyError("group_by specifies columns missing from table: %s" %unmatched)

    if add:
        groups = ordered_union(data.group_by, groups)

    return data.copy(group_by = groups) 

def add_note(db, server, user, text):
    id_query = select([sqlalchemy.sql.expression.func.max(table.c.note_id).label("maxid")]) \
        .where(table.c.user == user.lower())
    max_id = db.execute(id_query).scalar()

    if max_id is None:
        note_id = 1
    else:
        note_id = max_id + 1

    query = table.insert().values(
        note_id=note_id,
        connection=server,
        user=user.lower(),
        text=text,
        deleted=False,
        added=datetime.today()
    )
    db.execute(query)
    db.commit() 

def _mutate(__data, **kwargs):
    # Cases
    #  - work with group by
    #  - window functions
    # TODO: verify it can follow a renaming select

    # track labeled columns in set
    sel = __data.last_op

    # evaluate each call
    for colname, func in kwargs.items():
        # keep set of columns labeled (aliased) in this select statement
        # need to use inner cols, since sel.columns uses ColumnClause, not Label
        labs = set(k for k,v in lift_inner_cols(sel).items() if isinstance(v, sql.elements.Label))
        new_call = __data.shape_call(func, verb_name = "Mutate", arg_name = colname)

        sel = _mutate_select(sel, colname, new_call, labs, __data)

    return __data.append_op(sel) 

def lift_inner_cols(tbl):
    cols = list(tbl.inner_columns)
    data = {col.key: col for col in cols}

    return sql.base.ImmutableColumnCollection(data, cols) 

def exit(self, node):
        col_expr = node(self.columns)

        if not isinstance(col_expr, sql.elements.ClauseElement):
            return col_expr

        over_clauses = [x for x in sa_get_over_clauses(col_expr) if isinstance(x, CustomOverClause)]

        # put groupings and orderings onto custom over clauses
        for over in over_clauses:
            # TODO: shouldn't mutate these over clauses
            group_by = sql.elements.ClauseList(
                    *[self.columns[name] for name in self.group_by]
                    )
            order_by = sql.elements.ClauseList(
                    *_create_order_by_clause(self.columns, *self.order_by)
                    )

            over.set_over(group_by, order_by)

        if len(over_clauses) and self.window_cte is not None:
            label = col_expr.label(None)
            self.windows.append(label)

            # optionally put into CTE, and return its resulting column
            self.window_cte.append_column(label)
            win_col = self.window_cte.c.values()[-1]
            return win_col
                
        return col_expr 

def _repr_grouped_df_html_(self):
    return "<div><p>(grouped data frame)</p>" + self._selected_obj._repr_html_() + "</div>"



# Main Funcs 
# =============================================================================

# sql raw -------------- 

def use_simple_names():
    get_col_name = lambda el, *args, **kwargs: str(el.element.name)
    try:
        yield compiles(sql.compiler._CompileLabel)(get_col_name)
    except:
        pass
    finally:
        deregister(sql.compiler._CompileLabel) 

def _collect(__data, as_df = True):
    # TODO: maybe remove as_df options, always return dataframe
    # normally can just pass the sql objects to execute, but for some reason
    # psycopg2 completes about incomplete template.
    # see https://stackoverflow.com/a/47193568/1144523
    query = __data.last_op
    compiled = query.compile(
        dialect = __data.source.dialect,
        compile_kwargs = {"literal_binds": True}
    )
    if as_df:
        return pd.read_sql(compiled, __data.source)

    return __data.source.execute(compiled).fetchall() 

def _filter(__data, *args):
    # TODO: aggregate funcs
    # Note: currently always produces 2 additional select statements,
    #       1 for window/aggs, and 1 for the where clause
    sel = __data.last_op.alias()                   # original select
    win_sel = sql.select([sel], from_obj = sel)    # first cte

    conds = []
    windows = []
    for ii, arg in enumerate(args):
        if isinstance(arg, Call):
            new_call = __data.shape_call(arg, verb_name = "Filter", arg_name = ii)
            #var_cols = new_call.op_vars(attr_calls = False)

            col_expr, win_cols = __data.track_call_windows(
                    new_call,
                    sel.columns,
                    window_cte = win_sel
                    )

            conds.append(col_expr)
        else:
            conds.append(arg)

    bool_clause = sql.and_(*conds)

    # move non-window functions to refer to win_sel clause (not the innermost) ---
    win_alias = win_sel.alias()
    bool_clause = sql.util.ClauseAdapter(win_alias).traverse(bool_clause)

    
    # create second cte ----
    orig_cols = [win_alias.columns[k] for k in sel.columns.keys()]
    filt_sel = sql.select(orig_cols, from_obj = win_alias, whereclause = bool_clause)
    return __data.append_op(filt_sel) 

def _case_when(__data, cases):
    # TODO: will need listener to enter case statements, to handle when they use windows
    if isinstance(cases, Call):
        cases = cases(__data)

    whens = []
    case_items = list(cases.items())
    n_items = len(case_items)

    else_val = None
    for ii, (expr, val) in enumerate(case_items):
        # handle where val is a column expr
        if callable(val):
            val = val(__data)

        # handle when expressions
        if ii+1 == n_items and expr is True:
            else_val = val
        elif callable(expr):
            whens.append((expr(__data), val))
        else:
            whens.append((expr, val))

    return sql.case(whens, else_ = else_val)
        

# Join ------------------------------------------------------------------------ 

def _joined_cols(left_cols, right_cols, on_keys, full = False):
    """Return labeled columns, according to selection rules for joins.

    Rules:
        1. For join keys, keep left table's column
        2. When keys have the same labels, add suffix
    """
    # TODO: remove sets, so uses stable ordering
    # when left and right cols have same name, suffix with _x / _y
    keep_right = set(right_cols.keys()) - set(on_keys.values())
    shared_labs = set(left_cols.keys()).intersection(keep_right)

    right_cols_no_keys = {k: right_cols[k] for k in keep_right}

    # for an outer join, have key columns coalesce values
    if full:
        left_cols = {**left_cols}
        for lk, rk in on_keys.items():
            col = sql.functions.coalesce(left_cols[lk], right_cols[rk])
            left_cols[lk] = col.label(lk)

    # create labels ----
    l_labs = _relabeled_cols(left_cols, shared_labs, "_x")
    r_labs = _relabeled_cols(right_cols_no_keys, shared_labs, "_y")

    return l_labs + r_labs 

def _semi_join(left, right = None, on = None, *args, sql_on = None):
    _raise_if_args(args)

    left_sel = left.last_op.alias()
    right_sel = right.last_op.alias()

    # handle arguments ----
    on  = _validate_join_arg_on(on, sql_on)
    
    # create join conditions ----
    bool_clause = _create_join_conds(left_sel, right_sel, on)

    # create inner join ----
    exists_clause = sql.select(
            [sql.literal(1)],
            from_obj = right_sel,
            whereclause = bool_clause
            )

    # only keep left hand select's columns ----
    sel = sql.select(
            left_sel.columns,
            from_obj = left_sel,
            whereclause = sql.exists(exists_clause)
            )

    return left.append_op(sel) 

def _distinct(__data, *args, _keep_all = False, **kwargs):
    if (args or kwargs) and _keep_all:
        raise NotImplementedError("Distinct with variables specified in sql requires _keep_all = False")
    
    inner_sel = mutate(__data, **kwargs).last_op if kwargs else __data.last_op

    # TODO: this is copied from the df distinct version
    # cols dict below is used as ordered set
    cols = {simple_varname(x): True for x in args}
    cols.update(kwargs)

    if None in cols:
        raise KeyError("positional arguments must be simple column, "
                        "e.g. _.colname or _['colname']"
                        )

    # use all columns by default
    if not cols:
        cols = list(inner_sel.columns.keys())

    if not len(inner_sel._order_by_clause):
        # select distinct has to include any columns in the order by clause,
        # so can only safely modify existing statement when there's no order by
        sel_cols = lift_inner_cols(inner_sel)
        distinct_cols = [sel_cols[k] for k in cols]
        sel = inner_sel.with_only_columns(distinct_cols).distinct()
    else:
        # fallback to cte
        cte = inner_sel.alias()
        distinct_cols = [cte.columns[k] for k in cols]
        sel = sql.select(distinct_cols, from_obj = cte).distinct()

    return __data.append_op(sel)

    
# if_else --------------------------------------------------------------------- 

def load_ddl():
    for script in ('triggers.sql', 'rpmvercmp.sql'):
        with open(os.path.join(get_config('directories.datadir'), script)) as ddl_script:
            ddl = DDL(ddl_script.read())
        listen(Base.metadata, 'after_create', ddl.execute_if(dialect='postgresql')) 

def test_contain_percents_character_query(self, engine, conn):
        query = sqlalchemy.sql.text(
            """
            SELECT date_parse('20191030', '%Y%m%d')
            """
        )
        result = engine.execute(query)
        self.assertEqual(result.fetchall(), [(datetime(2019, 10, 30),)]) 

def test_query_with_parameter(self, engine, conn):
        query = sqlalchemy.sql.text(
            """
            SELECT :word
            """
        )
        result = engine.execute(query, word="cat")
        self.assertEqual(result.fetchall(), [("cat",)]) 

def sql_query_execute(client: Client, args: dict, *_) -> Tuple[str, Dict[str, Any], List[Dict[str, Any]]]:
    """
    Executes the sql query with the connection that was configured in the client
    :param client: the client object with the db connection
    :param args: demisto.args() including the sql query
    :return: Demisto outputs
    """
    try:
        sql_query = str(args.get('query'))
        limit = int(args.get('limit', 50))
        skip = int(args.get('skip', 0))
        bind_variables_names = args.get('bind_variables_names', "")
        bind_variables_values = args.get('bind_variables_values', "")
        bind_variables = generate_bind_vars(bind_variables_names, bind_variables_values)

        result, headers = client.sql_query_execute_request(sql_query, bind_variables)
        # converting an sqlalchemy object to a table
        converted_table = [dict(row) for row in result]
        # converting b'' and datetime objects to readable ones
        table = [{str(key): str(value) for key, value in dictionary.items()} for dictionary in converted_table]
        table = table[skip:skip + limit]
        human_readable = tableToMarkdown(name="Query result:", t=table, headers=headers,
                                         removeNull=True)
        context = {
            'Result': table,
            'Query': sql_query,
            'InstanceName': f'{client.dialect}_{client.dbname}'
        }
        entry_context: Dict = {'GenericSQL(val.Query && val.Query === obj.Query)': {'GenericSQL': context}}
        return human_readable, entry_context, table

    except Exception as err:
        # In case there is no query executed and only an action e.g - insert, delete, update
        # the result will raise an exception when we try to read the data from it
        if str(err) == "This result object does not return rows. It has been closed automatically.":
            human_readable = "Command executed"
            return human_readable, {}, []
        raise err 

def main():
    params = demisto.params()
    dialect = params.get('dialect')
    port = params.get('port')
    if port is None:
        port = generate_default_port_by_dialect(dialect)
    user = params.get("credentials").get("identifier")
    password = params.get("credentials").get("password")
    host = params.get('host')
    database = params.get('dbname')
    ssl_connect = params.get('ssl_connect')
    connect_parameters = params.get('connect_parameters')
    try:
        command = demisto.command()
        LOG(f'Command being called in SQL is: {command}')
        client = Client(dialect=dialect, host=host, username=user, password=password,
                        port=port, database=database, connect_parameters=connect_parameters, ssl_connect=ssl_connect)
        commands: Dict[str, Callable[[Client, Dict[str, str], str], Tuple[str, Dict[Any, Any], List[Any]]]] = {
            'test-module': test_module,
            'query': sql_query_execute,
            'sql-command': sql_query_execute
        }
        if command in commands:
            return_outputs(*commands[command](client, demisto.args(), command))
        else:
            raise NotImplementedError(f'{command} is not an existing Generic SQL command')
        client.connection.close()
    except Exception as err:
        return_error(f'Unexpected error: {str(err)} \nquery: {demisto.args().get("query")} \n{traceback.format_exc()}') 

def test_contain_percents_character_query(self, engine, conn):
        query = sqlalchemy.sql.text(
            """
            SELECT date_parse('20191030', '%Y%m%d')
            """
        )
        result = engine.execute(query)
        self.assertEqual(result.fetchall(), [(datetime(2019, 10, 30),)]) 

def test_query_with_parameter(self, engine, conn):
        query = sqlalchemy.sql.text(
            """
            SELECT :word
            """
        )
        result = engine.execute(query, word="cat")
        self.assertEqual(result.fetchall(), [("cat",)]) 

def test_drop_dup_entries_in_file_conn(self):
        table_name = "__test_tmp_table__"
        tmp_db_file = self.create_tempfiles([['name', '']], ext='.sql')[0]
        in_file_engine = session.EngineFacade(
            'sqlite:///%s' % tmp_db_file).get_engine()
        meta = MetaData()
        meta.bind = in_file_engine
        test_table, values = self._populate_db_for_drop_duplicate_entries(
            in_file_engine, meta, table_name)
        utils.drop_old_duplicate_entries_from_table(
            in_file_engine, table_name, False, 'b', 'c') 

def getIntegerLiteral(self, iValue):
        return sqlalchemy.sql.expression.literal(iValue, sqlalchemy.types.Integer);