def query(self, query_obj): """Executes the query and returns a dataframe query_obj is a dictionary representing Superset's query interface. Should return a ``superset.models.helpers.QueryResult`` """ logging.debug('query_obj: %s', query_obj) qry_start_dttm = datetime.now() status = QueryStatus.SUCCESS error_message = None df = None query_str = '' try: df = self.get_dataframe() df, query_str = self.process_dataframe(df, **query_obj) logging.debug('query_str: %s', query_str) except Exception as e: status = QueryStatus.FAILED logging.exception(e) error_message = str(e) return QueryResult(status=status, df=df, duration=datetime.now() - qry_start_dttm, query=query_str, error_message=error_message)
def query(self, query_obj): qry_start_dttm = datetime.now() sql = self.get_query_str(query_obj) status = utils.QueryStatus.SUCCESS error_message = None df = None try: df = self.database.get_df(sql, self.schema) except Exception as e: status = utils.QueryStatus.FAILED logging.exception(e) error_message = ( self.database.db_engine_spec.extract_error_message(e)) # if this is a main query with prequeries, combine them together if not query_obj['is_prequery']: query_obj['prequeries'].append(sql) sql = ';\n\n'.join(query_obj['prequeries']) sql += ';' return QueryResult(status=status, df=df, duration=datetime.now() - qry_start_dttm, query=sql, error_message=error_message)
def query(self, query_obj): qry_start_dttm = datetime.now() query_str_ext = self.get_query_str_extended(query_obj) sql = query_str_ext.sql status = utils.QueryStatus.SUCCESS error_message = None def mutator(df): labels_expected = query_str_ext.labels_expected if df is not None and not df.empty: if len(df.columns) != len(labels_expected): raise Exception(f"For {sql}, df.columns: {df.columns}" f" differs from {labels_expected}") else: df.columns = labels_expected return df try: df = self.database.get_df(sql, self.schema, mutator) except Exception as e: df = None status = utils.QueryStatus.FAILED logging.exception(f"Query {sql} on schema {self.schema} failed") db_engine_spec = self.database.db_engine_spec error_message = db_engine_spec.extract_error_message(e) return QueryResult( status=status, df=df, duration=datetime.now() - qry_start_dttm, query=sql, error_message=error_message, )
def query(self, query_obj): qry_start_dttm = datetime.now() query_str_ext = self.get_query_str_extended(query_obj) sql = query_str_ext.sql status = utils.QueryStatus.SUCCESS error_message = None df = None try: df = self.database.get_df(sql, self.schema) labels_expected = query_str_ext.labels_expected if df is not None and not df.empty: if len(df.columns) != len(labels_expected): raise Exception(f'For {sql}, df.columns: {df.columns}' f' differs from {labels_expected}') else: df.columns = labels_expected except Exception as e: status = utils.QueryStatus.FAILED logging.exception(f'Query {sql} on schema {self.schema} failed') db_engine_spec = self.database.db_engine_spec error_message = db_engine_spec.extract_error_message(e) # if this is a main query with prequeries, combine them together if not query_obj['is_prequery']: query_obj['prequeries'].append(sql) sql = ';\n\n'.join(query_obj['prequeries']) sql += ';' return QueryResult(status=status, df=df, duration=datetime.now() - qry_start_dttm, query=sql, error_message=error_message)
def query(self, query_obj): qry_start_dttm = datetime.now() query_str_ext = self.get_query_str_extended(query_obj) sql = query_str_ext.sql status = utils.QueryStatus.SUCCESS error_message = None df = None db_engine_spec = self.database.db_engine_spec try: df = self.database.get_df(sql, self.schema) if self.mutated_labels: df = df.rename(index=str, columns=self.mutated_labels) db_engine_spec.mutate_df_columns(df, sql, query_str_ext.labels_expected) except Exception as e: status = utils.QueryStatus.FAILED logging.exception(f'Query {sql} on schema {self.schema} failed') error_message = db_engine_spec.extract_error_message(e) # if this is a main query with prequeries, combine them together if not query_obj['is_prequery']: query_obj['prequeries'].append(sql) sql = ';\n\n'.join(query_obj['prequeries']) sql += ';' return QueryResult(status=status, df=df, duration=datetime.now() - qry_start_dttm, query=sql, error_message=error_message)
def query(self, query_obj: Dict) -> QueryResult: qry_start_dttm = datetime.now() query_str_ext = self.get_query_str_extended(query_obj) sql = query_str_ext.sql status = utils.QueryStatus.SUCCESS error_message = None def mutator(df): labels_expected = query_str_ext.labels_expected if df is not None and not df.empty: if len(df.columns) != len(labels_expected): raise Exception( f"For {sql}, df.columns: {df.columns}" f" differs from {labels_expected}" ) else: df.columns = labels_expected return df try: # df = self.database.get_df(sql, self.schema, mutator) database = self.database #get need mutil_tenant mutil_tenant = app.config.get("MULTI_TENANT") if mutil_tenant: # get tenant_code user_attribute = db.session.query(UserAttribute).filter_by(user_id = g.user.id).first() if user_attribute is not None: # get database by tenant_code tenant_code = user_attribute.tenant_code tenant_alias = app.config.get("TENANT_DATABASE_ALIAS", "") #search database tenant_database = db.session.query(Database).filter_by(database_name = "{0}{1}".format(tenant_alias,tenant_code)).first() if tenant_database is not None: database = tenant_database df = database.get_df(sql, self.schema, mutator) except Exception as e: df = None status = utils.QueryStatus.FAILED logging.exception(f"Query {sql} on schema {self.schema} failed") db_engine_spec = self.database.db_engine_spec error_message = db_engine_spec.extract_error_message(e) return QueryResult( status=status, df=df, duration=datetime.now() - qry_start_dttm, query=sql, error_message=error_message, )
def query(self, query_obj): qry_start_dttm = datetime.now() engine = self.database.get_sqla_engine() sql = self.get_query_str(engine, qry_start_dttm, **query_obj) status = QueryStatus.SUCCESS error_message = None df = None try: df = pd.read_sql_query(sql, con=engine) except Exception as e: status = QueryStatus.FAILED error_message = str(e) return QueryResult(status=status, df=df, duration=datetime.now() - qry_start_dttm, query=sql, error_message=error_message)
def query(self, query_obj): qry_start_dttm = datetime.now() sql = self.get_query_str(query_obj) status = QueryStatus.SUCCESS error_message = None df = None try: df = self.database.get_df(sql, self.schema) except Exception as e: status = QueryStatus.FAILED logging.exception(e) error_message = ( self.database.db_engine_spec.extract_error_message(e)) return QueryResult(status=status, df=df, duration=datetime.now() - qry_start_dttm, query=sql, error_message=error_message)
def query(self, query_obj): df = None error_message = None qry = db.session.query(Annotation) qry = qry.filter(Annotation.layer_id == query_obj['filter'][0]['val']) qry = qry.filter(Annotation.start_dttm >= query_obj['from_dttm']) qry = qry.filter(Annotation.end_dttm <= query_obj['to_dttm']) status = QueryStatus.SUCCESS try: df = pd.read_sql_query(qry.statement, db.engine) except Exception as e: status = QueryStatus.FAILED logging.exception(e) error_message = (utils.error_msg_from_exception(e)) return QueryResult(status=status, df=df, duration=0, query='', error_message=error_message)
def query(self, query_obj): qry_start_dttm = datetime.now() client = self.cluster.get_pydruid_client() query_str = self.get_query_str( client=client, query_obj=query_obj, phase=2) df = client.export_pandas() if df is None or df.size == 0: raise Exception(_('No data was returned.')) df.columns = [ DTTM_ALIAS if c == 'timestamp' else c for c in df.columns] is_timeseries = query_obj['is_timeseries'] \ if 'is_timeseries' in query_obj else True if ( not is_timeseries and DTTM_ALIAS in df.columns): del df[DTTM_ALIAS] # Reordering columns cols = [] if DTTM_ALIAS in df.columns: cols += [DTTM_ALIAS] cols += query_obj.get('groupby') or [] cols += query_obj.get('columns') or [] cols += query_obj.get('metrics') or [] cols = [col for col in cols if col in df.columns] df = df[cols] time_offset = DruidDatasource.time_offset(query_obj['granularity']) def increment_timestamp(ts): dt = utils.parse_human_datetime(ts).replace( tzinfo=DRUID_TZ) return dt + timedelta(milliseconds=time_offset) if DTTM_ALIAS in df.columns and time_offset: df[DTTM_ALIAS] = df[DTTM_ALIAS].apply(increment_timestamp) return QueryResult( df=df, query=query_str, duration=datetime.now() - qry_start_dttm)
def query(self, query_obj: QueryObjectDict) -> QueryResult: qry_start_dttm = datetime.now() query_str_ext = self.get_query_str_extended(query_obj) sql = query_str_ext.sql status = utils.QueryStatus.SUCCESS errors = None def mutator(df: pd.DataFrame) -> None: """ Some engines change the case or generate bespoke column names, either by default or due to lack of support for aliasing. This function ensures that the column names in the DataFrame correspond to what is expected by the viz components. :param df: Original DataFrame returned by the engine """ labels_expected = query_str_ext.labels_expected if df is not None and not df.empty: if len(df.columns) != len(labels_expected): raise Exception( f"For {sql}, df.columns: {df.columns}" f" differs from {labels_expected}" ) else: df.columns = labels_expected try: df = self.database.get_df(sql, self.schema, mutator) except Exception as ex: df = pd.DataFrame() status = utils.QueryStatus.FAILED logger.warning(f"Query {sql} on schema {self.schema} failed", exc_info=True) db_engine_spec = self.database.db_engine_spec errors = db_engine_spec.extract_errors(ex) return QueryResult( status=status, df=df, duration=datetime.now() - qry_start_dttm, query=sql, errors=errors, )
def query(self, query_obj: Dict[str, Any]) -> QueryResult: error_message = None qry = db.session.query(Annotation) qry = qry.filter(Annotation.layer_id == query_obj["filter"][0]["val"]) if query_obj["from_dttm"]: qry = qry.filter(Annotation.start_dttm >= query_obj["from_dttm"]) if query_obj["to_dttm"]: qry = qry.filter(Annotation.end_dttm <= query_obj["to_dttm"]) status = utils.QueryStatus.SUCCESS try: df = pd.read_sql_query(qry.statement, db.engine) except Exception as e: df = pd.DataFrame() status = utils.QueryStatus.FAILED logger.exception(e) error_message = utils.error_msg_from_exception(e) return QueryResult( status=status, df=df, duration=0, query="", error_message=error_message )
def query(self, query_obj): client = self.cluster.get_client() equery = {} # Aggregations equery['aggregations'] = {} for m in self.metrics: if m.metric_name in query_obj.get('metrics'): equery['aggregations'][m.metric_name] = m.json_obj data = client.search(index=self.index, body=equery) print('-=' * 20) print('query is: {}'.format(equery)) data = data['hits']['hits'] data = [k['_source'] for k in data] print('-=' * 20) query_str = self.query_str() qry_start_dttm = datetime.now() df = pd.DataFrame(data) print('-=' * 20) print(df) return QueryResult(df=df, query=query_str, duration=datetime.now() - qry_start_dttm)