Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: editor awareness of default schemas #1422

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 12 additions & 4 deletions querybook/server/datasources/search.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,16 +130,24 @@ def vector_search_tables(


@register("/suggest/<int:metastore_id>/tables/", methods=["GET"])
def suggest_tables(metastore_id, prefix, limit=10):
def suggest_tables(metastore_id, prefix, limit=10, active_schema='default'):
api_assert(limit is None or limit <= 100, "Requesting too many tables")
verify_metastore_permission(metastore_id)

possible_schema = prefix.split(".")[0]
omit_schema = not active_schema.startswith(possible_schema)

query = construct_suggest_table_query(prefix, limit, metastore_id)
options = get_matching_suggestions(query, ES_CONFIG["tables"]["index_name"])
texts = [
"{}.{}".format(
option.get("_source", {}).get("schema", ""),
option.get("_source", {}).get("name", ""),
(
option.get("_source", {}).get("name", "")
if omit_schema and active_schema == option.get("_source", {}).get("schema", "")
Comment on lines +137 to +145
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure I'm following what this active_schema and this omit_schema is trying to do here. could you elaborate a bit?

else
"{}.{}".format(
option.get("_source", {}).get("schema", ""),
option.get("_source", {}).get("name", ""),
)
)
for option in options
]
Expand Down
14 changes: 14 additions & 0 deletions querybook/server/models/admin.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,19 @@ class QueryEngine(CRUDMixin, Base):
),
)

def get_default_schema(self):
default_schema = "default"
connection_string = self.executor_params.get("connection_string")
try:
connection_url = sql.engine.make_url(connection_string)
except Exception:
connection_url = None

if connection_url:
default_schema = connection_url.database or default_schema
Comment on lines +110 to +117
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not all the connection strings have the database, I think it will be better to have a seprate executor param for the default schema.


return default_schema

def to_dict(self):
# IMPORTANT: do not expose executor params unless it is for admin
return {
Expand All @@ -115,6 +128,7 @@ def to_dict(self):
"metastore_id": self.metastore_id,
"feature_params": self.get_feature_params(),
"executor": self.executor,
"default_schema": self.get_default_schema(),
}

def to_dict_admin(self):
Expand Down
2 changes: 2 additions & 0 deletions querybook/webapp/components/QueryEditor/BoundQueryEditor.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ export const BoundQueryEditor = React.forwardRef<
| 'functionDocumentationByNameByLanguage'
| 'metastoreId'
| 'language'
| 'defaultSchema'
> & {
engine?: IQueryEngine;
cellId?: number;
Expand Down Expand Up @@ -109,6 +110,7 @@ export const BoundQueryEditor = React.forwardRef<
}
metastoreId={engine?.metastore_id}
language={engine?.language}
defaultSchema={engine?.default_schema}
/>
);

Expand Down
3 changes: 3 additions & 0 deletions querybook/webapp/components/QueryEditor/QueryEditor.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ export interface IQueryEditorProps extends IStyledQueryEditorProps {
lineWrapping?: boolean;
readOnly?: boolean;
language?: string;
defaultSchema?: string;
theme?: string;
functionDocumentationByNameByLanguage?: FunctionDocumentationCollection;
metastoreId?: number;
Expand Down Expand Up @@ -107,6 +108,7 @@ export const QueryEditor: React.FC<
lineWrapping = false,
readOnly,
language = 'hive',
defaultSchema = 'default',
theme = 'default',
functionDocumentationByNameByLanguage = {},
metastoreId,
Expand Down Expand Up @@ -136,6 +138,7 @@ export const QueryEditor: React.FC<
const { codeAnalysis, codeAnalysisRef } = useCodeAnalysis({
language,
query: value,
defaultSchema
});
const autoCompleterRef = useAutoComplete(
metastoreId,
Expand Down
1 change: 1 addition & 0 deletions querybook/webapp/const/queryEngine.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ export interface IQueryEngine {
name: string;
language: string;
description: string;
default_schema: string;

metastore_id: number;
executor: string;
Expand Down
9 changes: 5 additions & 4 deletions querybook/webapp/hooks/queryEditor/useCodeAnalysis.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,10 @@ import { analyzeCode } from 'lib/web-worker';
interface IUseCodeAnalysisParams {
language: string;
query: string;
defaultSchema: string;
}

export function useCodeAnalysis({ language, query }: IUseCodeAnalysisParams) {
export function useCodeAnalysis({ language, query, defaultSchema }: IUseCodeAnalysisParams) {
/**
* the ref version is used to pass into functions in codemirror
* this is to prevent unnecessary codemirror refreshes
Expand All @@ -19,13 +20,13 @@ export function useCodeAnalysis({ language, query }: IUseCodeAnalysisParams) {
const debouncedQuery = useDebounce(query, 500);

useEffect(() => {
analyzeCode(debouncedQuery, 'autocomplete', language).then(
analyzeCode(debouncedQuery, 'autocomplete', language, defaultSchema).then(
(codeAnalysis) => {
codeAnalysisRef.current = codeAnalysis;
setCodeAnalysis(codeAnalysis);
}
);
}, [debouncedQuery, language]);
}, [debouncedQuery, language, defaultSchema]);

return { codeAnalysisRef, codeAnalysis };
return { codeAnalysis, codeAnalysisRef };
}
5 changes: 4 additions & 1 deletion querybook/webapp/lib/sql-helper/sql-autocompleter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -240,9 +240,12 @@ export class SqlAutoCompleter {
return [];
}

let omit_schema = this.codeAnalysis.lineage.activeSchema || 'default';
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

const. also || default is not necessarily needed, as the API has the default value default for the 3rd parameter. could be just below

suggest(
            metastoreId,
            prefix,
            this.codeAnalysis.lineage.activeSchema)


const { data: names } = await SearchTableResource.suggest(
metastoreId,
prefix
prefix,
omit_schema
);
return names;
}
Expand Down
11 changes: 9 additions & 2 deletions querybook/webapp/lib/sql-helper/sql-lexer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ const contextSensitiveKeyWord = {
table: 'table',
update: 'table',
join: 'table',
on: 'column',
set: 'column',

desc: 'table',
Expand Down Expand Up @@ -183,6 +184,7 @@ export interface ILinterWarning extends IRange {
export interface ILineage {
references: Record<number, TableToken[]>;
aliases: Record<number, Record<string, TableToken>>;
activeSchema: string;
}

export interface ICodeAnalysis {
Expand Down Expand Up @@ -646,10 +648,13 @@ export function findWithStatementPlaceholder(statement: IToken[]) {
return placeholders;
}

export function findTableReferenceAndAlias(statements: IToken[][]) {
let defaultSchema = 'default';
export function findTableReferenceAndAlias(
statements: IToken[][],
defaultSchema: string = 'default'
) {
const references: Record<number, TableToken[]> = {};
const aliases: Record<number, Record<string, TableToken>> = {};
let activeSchema = defaultSchema;

statements.forEach((statement, statementNum) => {
if (statement.length === 0) {
Expand Down Expand Up @@ -679,6 +684,7 @@ export function findTableReferenceAndAlias(statements: IToken[][]) {
const secondToken = statement[tokenCounter++];
if (secondToken && secondToken.type === 'VARIABLE') {
defaultSchema = secondToken.text;
activeSchema = defaultSchema;
}
} else {
const placeholders: Set<string> = new Set(
Expand Down Expand Up @@ -807,6 +813,7 @@ export function findTableReferenceAndAlias(statements: IToken[][]) {
return {
references,
aliases,
activeSchema,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks like activeSchema is just the same as defaultSchema, you dont need to define a new variable, but just return the defaultSchema.

};
}

Expand Down
4 changes: 3 additions & 1 deletion querybook/webapp/lib/web-worker/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@ let sqlEditorWorker = null;
export function analyzeCode(
code: string,
mode = 'autocomplete',
language = 'hive'
language = 'hive',
defaultSchema = 'default'
): Promise<ICodeAnalysis> {
if (!sqlEditorWorker) {
sqlEditorWorker = new Worker(
Expand Down Expand Up @@ -40,6 +41,7 @@ export function analyzeCode(
mode,
id,
language,
defaultSchema
});
});
}
4 changes: 2 additions & 2 deletions querybook/webapp/lib/web-worker/sql-editor.worker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,12 @@ const context: Worker = self as any;
context.addEventListener(
'message',
(e) => {
const { id, mode, code, language } = e.data;
const { id, mode, code, language, defaultSchema } = e.data;
const tokens = tokenize(code, { language });
const statements = simpleParse(tokens);

const codeAnalysis: ICodeAnalysis = {
lineage: findTableReferenceAndAlias(statements),
lineage: findTableReferenceAndAlias(statements, defaultSchema),
};

if (mode === 'autocomplete') {
Expand Down
4 changes: 2 additions & 2 deletions querybook/webapp/resource/search.ts
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,9 @@ export const SearchTableResource = {
count: number;
}>('/search/tables/vector/', { ...params }),

suggest: (metastoreId: number, prefix: string) =>
suggest: (metastoreId: number, prefix: string, active_schema: string = 'default') =>
ds.fetch<string[]>(`/suggest/${metastoreId}/tables/`, {
prefix,
prefix, active_schema: active_schema
}),
};

Expand Down