Drizzle Team•3mo ago

Are dynamic selects with type inference possible?

Title. I'm trying to add select statements conditionally. I'd like the function to return a properly typed query based on what I've passed. Is that possible?

.select<SelectedFields<PgColumn, PgTable>>({
  ...(opts.groupBy && { irrigationClientId: shrink.irrigationClientId }),
  shrink: sum(shrink.amount).mapWith(Number).as('shrink'),
})

.select<SelectedFields<PgColumn, PgTable>>({
  ...(opts.groupBy && { irrigationClientId: shrink.irrigationClientId }),
  shrink: sum(shrink.amount).mapWith(Number).as('shrink'),
})

This prevents errors but isn't really what I want.

1 Reply

html_extraordinaire•3mo ago

After I worked on it a little bit I came up with this.

type TotalShrinkInput = {
  irrigationYearId: number;
  irrigationClientId?: number;
};

type TotalShrinkOpts<T> = {
  tx: typeof db;
  groupBy?: T;
};

export function shrinkSubquery<T extends SelectedFields<PgColumn, PgTable>>(
  input: TotalShrinkInput,
  opts: TotalShrinkOpts<T> = { tx: db }
) {
  const andArgs = {
    irrigationYearId: eq(irrigationYear.id, input.irrigationYearId),
    irrigationClientId: eq(
      auxiliaryIssuance.irrigationClientId,
      input.irrigationClientId!
    ),
  };
  const filteredArgs = Object.entries(andArgs)
    .filter(([k, v]) => {
      switch (k as keyof typeof andArgs) {
        case 'irrigationYearId':
          return true;
        case 'irrigationClientId':
          if (input.irrigationClientId) return true;
          return false;
      }
    })
    .map(([k, v]) => v);
  const whereClause = and(...filteredArgs);

  const sq = opts.tx
    .select<T & SelectedFields<PgColumn, PgTable>>({
      ...opts.groupBy as T,
      shrink: sum(shrink.amount).mapWith(Number).as('shrink'),
    })
    .from(shrink)
    .innerJoin(irrigationYear, eq(irrigationYear.id, shrink.irrigationYearId))
    .where(whereClause);


  if (opts.groupBy) {
    const [column] = Object.values(opts.groupBy);
    sq.groupBy(column); // also not being property inferred
  }

  return sq.as('shrinkSq');
}

type TotalShrinkInput = {
  irrigationYearId: number;
  irrigationClientId?: number;
};

type TotalShrinkOpts<T> = {
  tx: typeof db;
  groupBy?: T;
};

export function shrinkSubquery<T extends SelectedFields<PgColumn, PgTable>>(
  input: TotalShrinkInput,
  opts: TotalShrinkOpts<T> = { tx: db }
) {
  const andArgs = {
    irrigationYearId: eq(irrigationYear.id, input.irrigationYearId),
    irrigationClientId: eq(
      auxiliaryIssuance.irrigationClientId,
      input.irrigationClientId!
    ),
  };
  const filteredArgs = Object.entries(andArgs)
    .filter(([k, v]) => {
      switch (k as keyof typeof andArgs) {
        case 'irrigationYearId':
          return true;
        case 'irrigationClientId':
          if (input.irrigationClientId) return true;
          return false;
      }
    })
    .map(([k, v]) => v);
  const whereClause = and(...filteredArgs);

  const sq = opts.tx
    .select<T & SelectedFields<PgColumn, PgTable>>({
      ...opts.groupBy as T,
      shrink: sum(shrink.amount).mapWith(Number).as('shrink'),
    })
    .from(shrink)
    .innerJoin(irrigationYear, eq(irrigationYear.id, shrink.irrigationYearId))
    .where(whereClause);


  if (opts.groupBy) {
    const [column] = Object.values(opts.groupBy);
    sq.groupBy(column); // also not being property inferred
  }

  return sq.as('shrinkSq');
}

It seems to know now that irrigationClientId is a field on the subquery that gets returned from the function:

const shrinkSq = shrinkSubquery(
  { irrigationYearId: input.irrigationYearId },
  { groupBy: { irrigationClientId: shrink.irrigationClientId }, tx }
);

const shrinkSq = shrinkSubquery(
  { irrigationYearId: input.irrigationYearId },
  { groupBy: { irrigationClientId: shrink.irrigationClientId }, tx }
);

But what if I just want to pass "groupByIrrigationClientId: true"? Can I get type inference from the function somehow? I guess the problem with adding <T & SelectedFields<PgColumn, PgTable>> to the select function is that you don't get type inference on the other fields. How do I fix that? lol Removing the explicit type from the select function seems to have restored type inference on the other fields now. The column in the if statement is also not being properly inferred but I'm guessing this has to do with my use of Object.values? Is this too complex? lol Or actually I guess the issue is that groupBy needs a column so I can just cast column to a PgColumn. Settled on this.

type TotalShrinkInput = {
  irrigationClientId?: IrrigationClient['id'];
  irrigationYearId: IrrigationYear['id'];
};

type TotalShrinkOpts<TSelected> = {
  select?: TSelected;
  groupBy?: (PgColumn | SQL | SQL.Aliased)[];
};

export function shrinkSubquery<TSelected extends SelectedFields>(
  input: TotalShrinkInput,
  opts: TotalShrinkOpts<TSelected>,
  tx = db
) {
  const sq = tx
    .select({
      ...(opts?.select as TSelected),
      shrink: sum(shrink.amount).mapWith(Number).as('shrink'),
    })
    .from(shrink)
    .innerJoin(irrigationYear, eq(irrigationYear.id, shrink.irrigationYearId))
    .where(
      input.irrigationClientId
        ? and(
            eq(irrigationYear.id, input.irrigationYearId),
            eq(shrink.irrigationClientId, input.irrigationClientId)
          )
        : eq(irrigationYear.id, input.irrigationYearId)
    );

  if (opts?.groupBy) sq.groupBy(...opts.groupBy);

  return sq.as('shrinkSq');
}

type TotalShrinkInput = {
  irrigationClientId?: IrrigationClient['id'];
  irrigationYearId: IrrigationYear['id'];
};

type TotalShrinkOpts<TSelected> = {
  select?: TSelected;
  groupBy?: (PgColumn | SQL | SQL.Aliased)[];
};

export function shrinkSubquery<TSelected extends SelectedFields>(
  input: TotalShrinkInput,
  opts: TotalShrinkOpts<TSelected>,
  tx = db
) {
  const sq = tx
    .select({
      ...(opts?.select as TSelected),
      shrink: sum(shrink.amount).mapWith(Number).as('shrink'),
    })
    .from(shrink)
    .innerJoin(irrigationYear, eq(irrigationYear.id, shrink.irrigationYearId))
    .where(
      input.irrigationClientId
        ? and(
            eq(irrigationYear.id, input.irrigationYearId),
            eq(shrink.irrigationClientId, input.irrigationClientId)
          )
        : eq(irrigationYear.id, input.irrigationYearId)
    );

  if (opts?.groupBy) sq.groupBy(...opts.groupBy);

  return sq.as('shrinkSq');
}

Still wondering if it's possible to get type inference from passing a boolean to the function though.

Gaming

Programming

Are dynamic selects with type inference possible?