Skip to content

Commit

Permalink
Merge pull request #2 from singlestore-labs/add_join_support
Browse files Browse the repository at this point in the history 
Add join support
  • Loading branch information
@demenskyi
demenskyi authored Aug 29, 2024
2 parents 437b885 + 195f764 commit 37bffb6
Show file tree
Hide file tree
Showing 9 changed files with 216 additions and 135 deletions.
2 changes: 1 addition & 1 deletion examples/estore/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -155,7 +155,7 @@ async function main() {

console.log("Executing database methods...");
console.log('Creating "users" table...');
const newUsersTable = await db.createTable<Database["tables"]["users"]>({
const newUsersTable = await db.createTable<"users", Database["tables"]["users"]>({
name: "users",
columns: {
id: { type: "bigint", autoIncrement: true, primaryKey: true },
Expand Down
12 changes: 8 additions & 4 deletions packages/client/src/database.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ export interface DatabaseType {
*/
export interface DatabaseSchema<TType extends DatabaseType> {
name: string;
tables?: { [K in keyof TType["tables"]]: Omit<TableSchema<TType["tables"][K]>, "name"> };
tables?: { [K in keyof TType["tables"]]: Omit<TableSchema<any, TType["tables"][K]>, "name"> };
}

/**
Expand Down Expand Up @@ -255,14 +255,15 @@ export class Database<TDatabaseType extends DatabaseType = DatabaseType, TAi ext
*
* @param {TName} name - The name of the table to retrieve.
*
* @returns {Table<TType extends TableType ? TType : TDatabaseType["tables"][TName] extends TableType ? TDatabaseType["tables"][TName] : TableType, TDatabaseType, TAi>} A `Table` instance representing the specified table.
* @returns {Table<TName, TType extends TableType ? TType : TDatabaseType["tables"][TName] extends TableType ? TDatabaseType["tables"][TName] : TableType, TDatabaseType, TAi>} A `Table` instance representing the specified table.
*/
table<
TType,
TName extends DatabaseTableName<TDatabaseType> | (string & {}) = DatabaseTableName<TDatabaseType> | (string & {}),
>(
name: TName,
): Table<
TName,
TType extends TableType
? TType
: TDatabaseType["tables"][TName] extends TableType
Expand All @@ -272,6 +273,7 @@ export class Database<TDatabaseType extends DatabaseType = DatabaseType, TAi ext
TAi
> {
return new Table<
TName,
TType extends TableType
? TType
: TDatabaseType["tables"][TName] extends TableType
Expand Down Expand Up @@ -307,8 +309,10 @@ export class Database<TDatabaseType extends DatabaseType = DatabaseType, TAi ext
*
* @returns {Promise<Table<TType, TDatabaseType, TAi>>} A promise that resolves to the created `Table` instance.
*/
createTable<TType extends TableType = TableType>(schema: TableSchema<TType>): Promise<Table<TType, TDatabaseType, TAi>> {
return Table.create<TType, TDatabaseType, TAi>(this._connection, this.name, schema, this._ai);
createTable<TName extends string = string, TType extends TableType = TableType>(
schema: TableSchema<TName, TType>,
): Promise<Table<TName, TType, TDatabaseType, TAi>> {
return Table.create<TName, TType, TDatabaseType, TAi>(this._connection, this.name, schema, this._ai);
}

/**
Expand Down
1 change: 1 addition & 0 deletions packages/client/src/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import { Workspace, type ConnectWorkspaceConfig, type WorkspaceType } from "./wo

export type * from "./types";
export { escape } from "mysql2";
export { QueryBuilder } from "./query/builder";

/**
* Configuration object for initializing a `SingleStoreClient` instance.
Expand Down
88 changes: 57 additions & 31 deletions packages/client/src/query/builder.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,12 @@ import { escape } from "mysql2";
import type { DatabaseType } from "../database";
import type { TableType } from "../table";

export type SelectClause<TColumn> = (keyof TColumn | (string & {}))[];
export type SelectClause<
TTableName extends string,
TTableType extends TableType,
TDatabaseType extends DatabaseType,
_TTableColumns = TTableType["columns"],
> = ((string & {}) | keyof _TTableColumns)[];

export type WhereOperator<TColumnValue> = TColumnValue extends string
? {
Expand All @@ -26,44 +31,65 @@ export type WhereOperator<TColumnValue> = TColumnValue extends string
}
: never;

export type WhereClause<TColumn> = {
[K in keyof TColumn]?: WhereOperator<TColumn[K]> | TColumn[K];
} & {
OR?: WhereClause<TColumn>[];
NOT?: WhereClause<TColumn>;
export type WhereClause<
TTableName extends string,
TTableType extends TableType,
TDatabaseType extends DatabaseType,
_TTableColumns = TTableType["columns"],
> = { [K in keyof _TTableColumns]?: WhereOperator<_TTableColumns[K]> | _TTableColumns[K] } & {
OR?: WhereClause<TTableName, TTableType, TDatabaseType>[];
NOT?: WhereClause<TTableName, TTableType, TDatabaseType>;
};

export type GroupByClause<TColumn> = (keyof TColumn)[];
export type GroupByClause<
TTableName extends string,
TTableType extends TableType,
TDatabaseType extends DatabaseType,
_TTableColumns = TTableType["columns"],
> = ((string & {}) | keyof _TTableColumns)[];

export type OrderByDirection = "asc" | "desc";

export type OrderByClause<TColumn> = {
[K in keyof TColumn]?: OrderByDirection;
};

export interface QueryBuilderParams<TTableType extends TableType, TDatabaseType extends DatabaseType = DatabaseType> {
select?: SelectClause<TTableType["columns"]>;
where?: WhereClause<TTableType["columns"]>;
groupBy?: GroupByClause<TTableType["columns"]>;
orderBy?: OrderByClause<TTableType["columns"]>;
export type OrderByClause<
TTableName extends string,
TTableType extends TableType,
TDatabaseType extends DatabaseType,
_TTableColumns = TTableType["columns"],
> = { [K in string & {}]: OrderByDirection } & { [K in keyof _TTableColumns]?: OrderByDirection };

export interface QueryBuilderParams<
TTableName extends string,
TTableType extends TableType,
TDatabaseType extends DatabaseType,
> {
select?: SelectClause<TTableName, TTableType, TDatabaseType>;
where?: WhereClause<TTableName, TTableType, TDatabaseType>;
groupBy?: GroupByClause<TTableName, TTableType, TDatabaseType>;
orderBy?: OrderByClause<TTableName, TTableType, TDatabaseType>;
limit?: number;
offset?: number;
}

export type ExtractQuerySelectedColumn<TColumn, TParams extends QueryBuilderParams<any, any> | undefined> =
TParams extends QueryBuilderParams<any, any>
? TParams["select"] extends (keyof TColumn)[]
? Pick<TColumn, TParams["select"][number]>
: TColumn
: TColumn;
export type AnyQueryBuilderParams = QueryBuilderParams<any, any, any>;

export type ExtractQuerySelectedColumn<
TTableName extends string,
TDatabaseType extends DatabaseType,
TParams extends AnyQueryBuilderParams | undefined,
_Table extends TDatabaseType["tables"][TTableName] = TDatabaseType["tables"][TTableName],
> = TParams extends AnyQueryBuilderParams
? TParams["select"] extends (keyof _Table["columns"])[]
? Pick<_Table["columns"], TParams["select"][number]>
: _Table["columns"]
: _Table["columns"];

export class QueryBuilder<TTableType extends TableType, TDatabaseType extends DatabaseType = DatabaseType> {
export class QueryBuilder<TName extends string, TTableType extends TableType, TDatabaseType extends DatabaseType> {
constructor(
private _databaseName: string,
private _tableName: string,
private _tableName: TName,
) {}

buildSelectClause(select?: SelectClause<any>) {
buildSelectClause(select?: SelectClause<any, any, any>) {
const columns = select ? select : ["*"];
return `SELECT ${columns.join(", ")}`;
}
Expand Down Expand Up @@ -97,7 +123,7 @@ export class QueryBuilder<TTableType extends TableType, TDatabaseType extends Da
}
}

buildWhereClause(conditions?: WhereClause<any>): string {
buildWhereClause(conditions?: WhereClause<any, any, any>): string {
if (!conditions || !Object.keys(conditions).length) return "";

const clauses: string[] = [];
Expand All @@ -119,11 +145,11 @@ export class QueryBuilder<TTableType extends TableType, TDatabaseType extends Da
return `WHERE ${clauses.join(" AND ")}`;
}

buildGroupByClause(columns?: GroupByClause<any>): string {
buildGroupByClause(columns?: GroupByClause<any, any, any>): string {
return columns?.length ? `GROUP BY ${columns.join(", ")}` : "";
}

buildOrderByClause(clauses?: OrderByClause<any>): string {
buildOrderByClause(clauses?: OrderByClause<any, any, any>): string {
if (!clauses) return "";

const condition = Object.entries(clauses)
Expand All @@ -145,7 +171,7 @@ export class QueryBuilder<TTableType extends TableType, TDatabaseType extends Da
return typeof offset === "number" ? `OFFSET ${offset}` : "";
}

buildClauses(params?: QueryBuilderParams<TTableType, TDatabaseType>) {
buildClauses<TParams extends QueryBuilderParams<TName, TTableType, TDatabaseType>>(params?: TParams) {
return {
select: this.buildSelectClause(params?.select),
from: this.buildFromClause(),
Expand All @@ -157,7 +183,7 @@ export class QueryBuilder<TTableType extends TableType, TDatabaseType extends Da
};
}

buildQuery(params?: QueryBuilderParams<TTableType, TDatabaseType>) {
return Object.values(this.buildClauses(params)).join(" ");
buildQuery<TParams extends QueryBuilderParams<TName, TTableType, TDatabaseType>>(params?: TParams) {
return Object.values(this.buildClauses(params)).join(" ").trim();
}
}
60 changes: 30 additions & 30 deletions packages/client/src/table.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,16 +18,17 @@ export interface TableType {
/**
* Interface representing the schema of a table, including its columns, primary keys, full-text keys, and additional clauses.
*
* @typeParam TName - A type extending `string` that defines the name of the table.
* @typeParam TType - A type extending `TableType` that defines the structure of the table.
*
* @property {string} name - The name of the table.
* @property {TName} name - The name of the table.
* @property {Object} columns - An object where each key is a column name and each value is the schema of that column, excluding the name.
* @property {string[]} [primaryKeys] - An optional array of column names that form the primary key.
* @property {string[]} [fulltextKeys] - An optional array of column names that form full-text keys.
* @property {string[]} [clauses] - An optional array of additional SQL clauses for the table definition.
*/
export interface TableSchema<TType extends TableType> {
name: string;
export interface TableSchema<TName extends string, TType extends TableType> {
name: TName;
columns: { [K in keyof TType["columns"]]: Omit<ColumnSchema, "name"> };
primaryKeys?: string[];
fulltextKeys?: string[];
Expand Down Expand Up @@ -87,6 +88,7 @@ type VectorScoreKey = "v_score";
* @property {VectorScoreKey} vScoreKey - The key used for vector scoring in vector search queries, defaulting to `"v_score"`.
*/
export class Table<
TName extends string = string,
TType extends TableType = TableType,
TDatabaseType extends DatabaseType = DatabaseType,
TAi extends AnyAI | undefined = undefined,
Expand All @@ -97,7 +99,7 @@ export class Table<
constructor(
private _connection: Connection,
public databaseName: string,
public name: string,
public name: TName,
private _ai?: TAi,
) {
this._path = [databaseName, name].join(".");
Expand Down Expand Up @@ -144,7 +146,7 @@ export class Table<
*
* @returns {string} An SQL string representing the table definition.
*/
static schemaToClauses(schema: TableSchema<any>): string {
static schemaToClauses(schema: TableSchema<any, any>): string {
const clauses: string[] = [
...Object.entries(schema.columns).map(([name, schema]) => {
return Column.schemaToClauses({ ...schema, name });
Expand All @@ -160,6 +162,7 @@ export class Table<
/**
* Creates a new table in the database with the specified schema.
*
* @typeParam TName - The name of the table, which extends `string`.
* @typeParam TType - The type of the table, which extends `TableType`.
* @typeParam TDatabaseType - The type of the database, which extends `DatabaseType`.
* @typeParam TAi - The type of AI functionalities integrated with the table, which can be undefined.
Expand All @@ -169,24 +172,25 @@ export class Table<
* @param {TableSchema<TType>} schema - The schema defining the structure of the table.
* @param {TAi} [ai] - Optional AI functionalities to associate with the table.
*
* @returns {Promise<Table<TDatabaseType, TType, TAi>>} A promise that resolves to the created `Table` instance.
* @returns {Promise<Table<TName, TType, TDatabaseType, TAi>>} A promise that resolves to the created `Table` instance.
*/
static async create<
TName extends string = string,
TType extends TableType = TableType,
TDatabaseType extends DatabaseType = DatabaseType,
TAi extends AnyAI | undefined = undefined,
>(
connection: Connection,
databaseName: string,
schema: TableSchema<TType>,
schema: TableSchema<TName, TType>,
ai?: TAi,
): Promise<Table<TType, TDatabaseType, TAi>> {
): Promise<Table<TName, TType, TDatabaseType, TAi>> {
const clauses = Table.schemaToClauses(schema);
await connection.client.execute<ResultSetHeader>(`\
CREATE TABLE IF NOT EXISTS ${databaseName}.${schema.name} (${clauses})
`);

return new Table<TType, TDatabaseType, TAi>(connection, databaseName, schema.name, ai);
return new Table<TName, TType, TDatabaseType, TAi>(connection, databaseName, schema.name, ai);
}

/**
Expand Down Expand Up @@ -331,13 +335,11 @@ export class Table<
*
* @param {TParams} params - The arguments defining the query, including selected columns, filters, and other options.
*
* @returns {Promise<(ExtractQuerySelectedColumn<TType["columns"], TParams> & RowDataPacket)[]>} A promise that resolves to an array of selected rows.
* @returns {Promise<(ExtractQuerySelectedColumn<TName, TDatabaseType, TParams> & RowDataPacket)[]>} A promise that resolves to an array of selected rows.
*/
async find<TParams extends QueryBuilderParams<TType, TDatabaseType> | undefined>(
params?: TParams,
): Promise<(ExtractQuerySelectedColumn<TType["columns"], TParams> & RowDataPacket)[]> {
type SelectedColumn = ExtractQuerySelectedColumn<TType["columns"], TParams>;
const queryBuilder = new QueryBuilder<TType, TDatabaseType>(this.databaseName, this.name);
async find<TParams extends QueryBuilderParams<TName, TType, TDatabaseType>>(params?: TParams) {
type SelectedColumn = ExtractQuerySelectedColumn<TName, TDatabaseType, TParams>;
const queryBuilder = new QueryBuilder<TName, TType, TDatabaseType>(this.databaseName, this.name);
const query = queryBuilder.buildQuery(params);
const [rows] = await this._connection.client.execute<(SelectedColumn & RowDataPacket)[]>(query);
return rows;
Expand All @@ -347,11 +349,14 @@ export class Table<
* Updates rows in the table based on the specified values and filters.
*
* @param {Partial<TType["columns"]>} values - The values to update in the table.
* @param {WhereClause<TType["columns"]>} where - The where clause to apply to the update query.
* @param {WhereClause<TName, TDatabaseType>} where - The where clause to apply to the update query.
*
* @returns {Promise<[ResultSetHeader, FieldPacket[]]>} A promise that resolves when the update is complete.
*/
update(values: Partial<TType["columns"]>, where: WhereClause<TType["columns"]>): Promise<[ResultSetHeader, FieldPacket[]]> {
update(
values: Partial<TType["columns"]>,
where: WhereClause<TName, TType, TDatabaseType>,
): Promise<[ResultSetHeader, FieldPacket[]]> {
const _where = new QueryBuilder(this.databaseName, this.name).buildWhereClause(where);

const columnAssignments = Object.keys(values)
Expand All @@ -365,11 +370,11 @@ export class Table<
/**
* Deletes rows from the table based on the specified filters. If no filters are provided, the table is truncated.
*
* @param {WhereClause<TType["columns"]>} [where] - The where clause to apply to the delete query.
* @param {WhereClause<TName, TDatabaseType>} [where] - The where clause to apply to the delete query.
*
* @returns {Promise<[ResultSetHeader, FieldPacket[]]>} A promise that resolves when the delete operation is complete.
*/
delete(where?: WhereClause<TType["columns"]>): Promise<[ResultSetHeader, FieldPacket[]]> {
delete(where?: WhereClause<TName, TType, TDatabaseType>): Promise<[ResultSetHeader, FieldPacket[]]> {
if (!where) return this.truncate();
const _where = new QueryBuilder(this.databaseName, this.name).buildWhereClause(where);
const query = `DELETE FROM ${this._path} ${_where}`;
Expand All @@ -395,7 +400,7 @@ export class Table<
* @param {TQueryParams} [queryParams] - Optional query builder parameters to refine the search, such as filters,
* groupings, orderings, limits, and offsets.
*
* @returns {Promise<(ExtractQuerySelectedColumn<TType["columns"], TQueryParams> & { [K in VectorScoreKey]: number } & RowDataPacket)[]>}
* @returns {Promise<(ExtractQuerySelectedColumn<TName, TDatabaseType, TQueryParams> & { v_score: number } & RowDataPacket)[]>}
* A promise that resolves to an array of rows matching the vector search criteria, each row including
* the selected columns and a vector similarity score.
*/
Expand All @@ -405,17 +410,12 @@ export class Table<
vectorColumn: TableColumnName<TType>;
embeddingParams?: TAi extends AnyAI ? Parameters<TAi["embeddings"]["create"]>[1] : never;
},
TQueryParams extends QueryBuilderParams<TType, TDatabaseType>,
>(
params: TParams,
queryParams?: TQueryParams,
): Promise<
(ExtractQuerySelectedColumn<TType["columns"], TQueryParams> & { [K in VectorScoreKey]: number } & RowDataPacket)[]
> {
type SelectedColumn = ExtractQuerySelectedColumn<TType["columns"], TQueryParams>;
TQueryParams extends QueryBuilderParams<TName, TType, TDatabaseType>,
>(params: TParams, queryParams?: TQueryParams) {
type SelectedColumn = ExtractQuerySelectedColumn<TName, TDatabaseType, TQueryParams>;
type ResultColumn = SelectedColumn & { [K in VectorScoreKey]: number };

const clauses = new QueryBuilder<TType, TDatabaseType>(this.databaseName, this.name).buildClauses(queryParams);
const clauses = new QueryBuilder<TName, TType, TDatabaseType>(this.databaseName, this.name).buildClauses(queryParams);
const promptEmbedding = (await this.ai.embeddings.create(params.prompt, params.embeddingParams))[0] || [];
let orderByClause = `ORDER BY ${this.vScoreKey} DESC`;

Expand Down Expand Up @@ -464,7 +464,7 @@ export class Table<
async createChatCompletion<
TParams extends Parameters<this["vectorSearch"]>[0] &
(TAi extends AnyAI ? Parameters<TAi["chatCompletions"]["create"]>[0] : never) & { template?: string },
TQueryParams extends QueryBuilderParams<TType, TDatabaseType>,
TQueryParams extends QueryBuilderParams<TName, TType, TDatabaseType>,
>(params: TParams, queryParams?: TQueryParams): Promise<CreateChatCompletionResult<TParams["stream"]>> {
const { prompt, systemRole, template, vectorColumn, embeddingParams, ...createChatCompletionParams } = params;

Expand Down
Loading

0 comments on commit 37bffb6

Please sign in to comment.