Python Language Generation

.nvmrc

@@ -1 +1 @@
-v20
+v20

src/server/routes/generators/python.ts

@@ -0,0 +1,35 @@
+import type { FastifyInstance } from 'fastify'
+import { PostgresMeta } from '../../../lib/index.js'
+import { DEFAULT_POOL_CONFIG } from '../../constants.js'
+import { extractRequestForLogging } from '../../utils.js'
+import { apply as applyPyTemplate } from '../../templates/python.js'
+import { getGeneratorMetadata } from '../../../lib/generators.js'
+
+export default async (fastify: FastifyInstance) => {
+  fastify.get<{
+    Headers: { pg: string }
+    Querystring: {
+      excluded_schemas?: string
+      included_schemas?: string
+    }
+  }>('/', async (request, reply) => {
+    const connectionString = request.headers.pg
+    const excludedSchemas =
+      request.query.excluded_schemas?.split(',').map((schema) => schema.trim()) ?? []
+    const includedSchemas =
+      request.query.included_schemas?.split(',').map((schema) => schema.trim()) ?? []
+
+    const pgMeta: PostgresMeta = new PostgresMeta({ ...DEFAULT_POOL_CONFIG, connectionString })
+    const { data: generatorMeta, error: generatorMetaError } = await getGeneratorMetadata(pgMeta, {
+      includedSchemas,
+      excludedSchemas,
+    })
+    if (generatorMetaError) {
+      request.log.error({ error: generatorMetaError, request: extractRequestForLogging(request) })
+      reply.code(500)
+      return { error: generatorMetaError.message }
+    }
+
+    return applyPyTemplate(generatorMeta)
+  })
+}

src/server/routes/index.ts

@@ -21,6 +21,7 @@ import ViewsRoute from './views.js'
 import TypeScriptTypeGenRoute from './generators/typescript.js'
 import GoTypeGenRoute from './generators/go.js'
 import SwiftTypeGenRoute from './generators/swift.js'
+import PythonTypeGenRoute from './generators/python.js'
 import { PG_CONNECTION, CRYPTO_KEY } from '../constants.js'
 
 export default async (fastify: FastifyInstance) => {
@@ -67,4 +68,5 @@ export default async (fastify: FastifyInstance) => {
   fastify.register(TypeScriptTypeGenRoute, { prefix: '/generators/typescript' })
   fastify.register(GoTypeGenRoute, { prefix: '/generators/go' })
   fastify.register(SwiftTypeGenRoute, { prefix: '/generators/swift' })
+  fastify.register(PythonTypeGenRoute, { prefix: '/generators/python' })
 }

src/server/server.ts

@@ -17,6 +17,7 @@ import {
 import { apply as applyTypescriptTemplate } from './templates/typescript.js'
 import { apply as applyGoTemplate } from './templates/go.js'
 import { apply as applySwiftTemplate } from './templates/swift.js'
+import { apply as applyPythonTemplate } from './templates/python.js'
 
 const logger = pino({
   formatters: {
@@ -140,6 +141,8 @@ async function getTypeOutput(): Promise<string | null> {
       })
     case 'go':
       return applyGoTemplate(config)
+    case 'python':
+      return applyPythonTemplate(config)
     default:
       throw new Error(`Unsupported language for GENERATE_TYPES: ${GENERATE_TYPES}`)
   }

src/server/templates/python.ts

@@ -0,0 +1,304 @@
+import type {
+  PostgresColumn,
+  PostgresMaterializedView,
+  PostgresSchema,
+  PostgresTable,
+  PostgresType,
+  PostgresView,
+} from '../../lib/index.js'
+import type { GeneratorMetadata } from '../../lib/generators.js'
+
+type Operation = 'Select' | 'Insert' | 'Update'
+
+export const apply = ({
+  schemas,
+  tables,
+  views,
+  materializedViews,
+  columns,
+  types,
+}: GeneratorMetadata): string => {
+  const columnsByTableId = columns
+    .sort(({ name: a }, { name: b }) => a.localeCompare(b))
+    .reduce(
+      (acc, curr) => {
+        acc[curr.table_id] ??= []
+        acc[curr.table_id].push(curr)
+        return acc
+      },
+      {} as Record<string, PostgresColumn[]>
+    )
+
+  const compositeTypes = types.filter((type) => type.attributes.length > 0)
+
+  let output = `
+from pydantic import BaseModel, Json
+from typing import Any, Annotated
+import datetime
+
+${tables
+  .filter((table) => schemas.some((schema) => schema.name === table.schema))
+  .flatMap((table) =>
+    generateTableStructsForOperations(
+      schemas.find((schema) => schema.name === table.schema)!,
+      table,
+      columnsByTableId[table.id],
+      types,
+      ['Select', 'Insert', 'Update']
+    )
+  )
+  .join('\n\n')}
+
+${views
+  .filter((view) => schemas.some((schema) => schema.name === view.schema))
+  .flatMap((view) =>
+    generateTableStructsForOperations(
+      schemas.find((schema) => schema.name === view.schema)!,
+      view,
+      columnsByTableId[view.id],
+      types,
+      ['Select']
+    )
+  )
+  .join('\n\n')}
+
+${materializedViews
+  .filter((materializedView) => schemas.some((schema) => schema.name === materializedView.schema))
+  .flatMap((materializedView) =>
+    generateTableStructsForOperations(
+      schemas.find((schema) => schema.name === materializedView.schema)!,
+      materializedView,
+      columnsByTableId[materializedView.id],
+      types,
+      ['Select']
+    )
+  )
+  .join('\n\n')}
+
+${compositeTypes
+  .filter((compositeType) => schemas.some((schema) => schema.name === compositeType.schema))
+  .map((compositeType) =>
+    generateCompositeTypeStruct(
+      schemas.find((schema) => schema.name === compositeType.schema)!,
+      compositeType,
+      types
+    )
+  )
+  .join('\n\n')}
+`.trim()
+
+  return output
+}
+
+/**
+ * Converts a Postgres name to PascalCase.
+ *
+ * @example
+ * ```ts
+ * formatForPyTypeName('pokedex') // Pokedex
+ * formatForPyTypeName('pokemon_center') // PokemonCenter
+ * formatForPyTypeName('victory-road') // VictoryRoad
+ * formatForPyTypeName('pokemon league') // PokemonLeague
+ * ```
+ */
+function formatForPyClassName(name: string): string {
+  return name
+    .split(/[^a-zA-Z0-9]/)
+    .map((word) => `${word[0].toUpperCase()}${word.slice(1)}`)
+    .join('')
+}
+
+/**
+ * Converts a Postgres name to snake_case.
+ *
+ * @example
+ * ```ts
+ * formatForPyTypeName('Pokedex') // pokedex
+ * formatForPyTypeName('PokemonCenter') // pokemon_enter
+ * formatForPyTypeName('victory-road') // victory_road
+ * formatForPyTypeName('pokemon league') // pokemon_league
+ * ```
+ */
+function formatForPyAttributeName(name: string): string {
+  return name
+    .split(/[^a-zA-Z0-9]+/) // Split on non-alphanumeric characters (like spaces, dashes, etc.)
+    .map(word => word.toLowerCase()) // Convert each word to lowercase
+    .join('_'); // Join with underscores
+}
+
+
+function generateTableStruct(
+  schema: PostgresSchema,
+  table: PostgresTable | PostgresView | PostgresMaterializedView,
+  columns: PostgresColumn[] | undefined,
+  types: PostgresType[],
+  operation: Operation
+): string {
+  // Storing columns as a tuple of [formattedName, type, name] rather than creating the string
+  // representation of the line allows us to pre-format the entries. Go formats
+  // struct fields to be aligned, e.g.:
+  // ```go
+  // type Pokemon struct {
+  //   id   int    `json:"id"`
+  //   name string `json:"name"`
+  // }
+  const columnEntries: [string, string, string][] =
+    columns?.map((column) => {
+      let nullable: boolean
+      if (operation === 'Insert') {
+        nullable =
+          column.is_nullable || column.is_identity || column.is_generated || !!column.default_value
+      } else if (operation === 'Update') {
+        nullable = true
+      } else {
+        nullable = column.is_nullable
+      }
+      return [
+        formatForPyAttributeName(column.name),
+        pgTypeToPythonType(column.format, nullable, types),
+        column.name,
+      ]
+    }) ?? []
+
+  // Pad the formatted name and type to align the struct fields, then join
+  // create the final string representation of the struct fields.
+  const formattedColumnEntries = columnEntries.map(([formattedName, type, name]) => {
+    return `  ${formattedName}: Annotated[${type}, Field(alias="${name}")]`
+  })
+
+  return `
+class ${formatForPyClassName(schema.name)}${formatForPyClassName(table.name)}${operation}(BaseModel):
+${formattedColumnEntries.join('\n')}
+`.trim()
+}
+
+function generateTableStructsForOperations(
+  schema: PostgresSchema,
+  table: PostgresTable | PostgresView | PostgresMaterializedView,
+  columns: PostgresColumn[] | undefined,
+  types: PostgresType[],
+  operations: Operation[]
+): string[] {
+  return operations.map((operation) =>
+    generateTableStruct(schema, table, columns, types, operation)
+  )
+}
+
+function generateCompositeTypeStruct(
+  schema: PostgresSchema,
+  type: PostgresType,
+  types: PostgresType[]
+): string {
+  // Use the type_id of the attributes to find the types of the attributes
+  const typeWithRetrievedAttributes = {
+    ...type,
+    attributes: type.attributes.map((attribute) => {
+      const type = types.find((type) => type.id === attribute.type_id)
+      return {
+        ...attribute,
+        type,
+      }
+    }),
+  }
+  const attributeEntries: [string, string, string][] = typeWithRetrievedAttributes.attributes.map(
+    (attribute) => [
+      formatForPyAttributeName(attribute.name),
+      pgTypeToPythonType(attribute.type!.format, false),
+      attribute.name,
+    ]
+  )
+
+  const [maxFormattedNameLength, maxTypeLength] = attributeEntries.reduce(
+    ([maxFormattedName, maxType], [formattedName, type]) => {
+      return [Math.max(maxFormattedName, formattedName.length), Math.max(maxType, type.length)]
+    },
+    [0, 0]
+  )
+
+  // Pad the formatted name and type to align the struct fields, then join
+  // create the final string representation of the struct fields.
+  const formattedAttributeEntries = attributeEntries.map(([formattedName, type, name]) => {
+    return `  ${formattedName.padEnd(maxFormattedNameLength)} ${type.padEnd(
+      maxTypeLength
+    )} \`json:"${name}"\``
+  })
+
+  return `
+class ${formatForPyClassName(schema.name)}${formatForPyClassName(type.name)}(BaseModel):
+${formattedAttributeEntries.join('\n')}
+`.trim()
+}
+
+const PY_TYPE_MAP = {
+  // Bool
+  bool: 'bool',
+
+  // Numbers
+  int2: 'int',
+  int4: 'int',
+  int8: 'int',
+  float4: 'float',
+  float8: 'float',
+  numeric: 'float',
+
+  // Strings
+  bytea: 'bytes',
+  bpchar: 'str',
+  varchar: 'str',
+  string: 'str',
+  date: 'datetime.date',
+  text: 'str',
+  citext: 'str',
+  time: 'datetime.time',
+  timetz: 'datetime.time',
+  timestamp: 'datetime.datetime',
+  timestamptz: 'datetime.datetime',
+  uuid: 'uuid.UUID',
+  vector: 'list[Any]',
+
+  // JSON
+  json: 'Json[Any]',
+  jsonb: 'Json[Any]',
+
+  // Range types (can be adjusted to more complex types if needed)
+  int4range: 'str',
+  int4multirange: 'str',
+  int8range: 'str',
+  int8multirange: 'str',
+  numrange: 'str',
+  nummultirange: 'str',
+  tsrange: 'str',
+  tsmultirange: 'str',
+  tstzrange: 'str',
+  tstzmultirange: 'str',
+  daterange: 'str',
+  datemultirange: 'str',
+
+  // Miscellaneous types
+  void: 'None',
+  record: 'dict[str, Any]',
+} as const
+
+type PythonType = (typeof PY_TYPE_MAP)[keyof typeof PY_TYPE_MAP]
+
+function pgTypeToPythonType(pgType: string, nullable: boolean, types: PostgresType[] = []): string {
+  let pythonType: PythonType | string | undefined = undefined
+
+  if (pgType in PY_TYPE_MAP) {
+    pythonType = PY_TYPE_MAP[pgType as keyof typeof PY_TYPE_MAP]
+  }
+
+  // Enums
+  const enumType = types.find((type) => type.name === pgType && type.enums.length > 0)
+  if (enumType) {
+    pythonType = formatForPyClassName(String(pgType))
+  }
+
+  if (pythonType) {
+    // If the type is nullable, append "| None" to the type
+    return nullable ? `${pythonType} | None` : pythonType
+  }
+
+  // Fallback
+  return nullable ? String(pgType)+' | None' : String(pgType)
+}

test/server/typegen.ts

@@ -2492,3 +2492,198 @@ test('typegen: swift w/ public access control', async () => {
     }"
   `)
 })
+
+test('typegen: python', async () => {
+  const { body } = await app.inject({
+    method: 'GET',
+    path: '/generators/python',
+    query: { access_control: 'public' },
+  })
+  expect(body).toMatchInlineSnapshot(`
+"from pydantic import BaseModel, Json
+from typing import Any, Annotated
+
+import datetime
+
+class UserStatus(Enum):
+
+class MemeStatus(Enum):
+
+class PublicUsersSelect(BaseModel):
+  id: Annotated[int, Field(alias="id")]
+  name: Annotated[str | None, Field(alias="name")]
+  status: Annotated[UserStatus | None, Field(alias="status")]
+
+
+class PublicUsersInsert(BaseModel):
+  id: Annotated[int | None, Field(alias="id")]
+  name: Annotated[str | None, Field(alias="name")]
+  status: Annotated[UserStatus | None, Field(alias="status")]
+
+
+class PublicUsersUpdate(BaseModel):
+  id: Annotated[int | None, Field(alias="id")]
+  name: Annotated[str | None, Field(alias="name")]
+  status: Annotated[UserStatus | None, Field(alias="status")]
+
+
+class PublicTodosSelect(BaseModel):
+  details: Annotated[str | None, Field(alias="details")]
+  id: Annotated[int, Field(alias="id")]
+  user_id: Annotated[int, Field(alias="user-id")]
+
+
+class PublicTodosInsert(BaseModel):
+  details: Annotated[str | None, Field(alias="details")]
+  id: Annotated[int | None, Field(alias="id")]
+  user_id: Annotated[int, Field(alias="user-id")]
+
+
+class PublicTodosUpdate(BaseModel):
+  details: Annotated[str | None, Field(alias="details")]
+  id: Annotated[int | None, Field(alias="id")]
+  user_id: Annotated[int | None, Field(alias="user-id")]
+
+
+class PublicUsersAuditSelect(BaseModel):
+  created_at: Annotated[datetime.datetime | None, Field(alias="created_at")]
+  id: Annotated[int, Field(alias="id")]
+  previous_value: Annotated[Json[Any] | None, Field(alias="previous_value")]
+  user_id: Annotated[int | None, Field(alias="user_id")]
+
+
+class PublicUsersAuditInsert(BaseModel):
+  created_at: Annotated[datetime.datetime | None, Field(alias="created_at")]
+  id: Annotated[int | None, Field(alias="id")]
+  previous_value: Annotated[Json[Any] | None, Field(alias="previous_value")]
+  user_id: Annotated[int | None, Field(alias="user_id")]
+
+
+class PublicUsersAuditUpdate(BaseModel):
+  created_at: Annotated[datetime.datetime | None, Field(alias="created_at")]
+  id: Annotated[int | None, Field(alias="id")]
+  previous_value: Annotated[Json[Any] | None, Field(alias="previous_value")]
+  user_id: Annotated[int | None, Field(alias="user_id")]
+
+
+class PublicUserDetailsSelect(BaseModel):
+  details: Annotated[str | None, Field(alias="details")]
+  user_id: Annotated[int, Field(alias="user_id")]
+
+
+class PublicUserDetailsInsert(BaseModel):
+  details: Annotated[str | None, Field(alias="details")]
+  user_id: Annotated[int, Field(alias="user_id")]
+
+
+class PublicUserDetailsUpdate(BaseModel):
+  details: Annotated[str | None, Field(alias="details")]
+  user_id: Annotated[int | None, Field(alias="user_id")]
+
+
+class PublicEmptySelect(BaseModel):
+  pass
+
+
+class PublicEmptyInsert(BaseModel):
+  pass
+
+
+class PublicEmptyUpdate(BaseModel):
+  pass
+
+
+class PublicTableWithOtherTablesRowTypeSelect(BaseModel):
+  col1: Annotated[PublicUserDetailsSelect, Field(alias="col1")]
+  col2: Annotated[PublicAViewSelect, Field(alias="col2")]
+
+
+class PublicTableWithOtherTablesRowTypeInsert(BaseModel):
+  col1: Annotated[PublicUserDetailsSelect, Field(alias="col1")]
+  col2: Annotated[PublicAViewSelect, Field(alias="col2")]
+
+
+class PublicTableWithOtherTablesRowTypeUpdate(BaseModel):
+  col1: Annotated[PublicUserDetailsSelect, Field(alias="col1")]
+  col2: Annotated[PublicAViewSelect, Field(alias="col2")]
+
+
+class PublicTableWithPrimaryKeyOtherThanIdSelect(BaseModel):
+  name: Annotated[str | None, Field(alias="name")]
+  other_id: Annotated[int, Field(alias="other_id")]
+
+
+class PublicTableWithPrimaryKeyOtherThanIdInsert(BaseModel):
+  name: Annotated[str | None, Field(alias="name")]
+  other_id: Annotated[int | None, Field(alias="other_id")]
+
+
+class PublicTableWithPrimaryKeyOtherThanIdUpdate(BaseModel):
+  name: Annotated[str | None, Field(alias="name")]
+  other_id: Annotated[int | None, Field(alias="other_id")]
+
+
+class PublicCategorySelect(BaseModel):
+  id: Annotated[int, Field(alias="id")]
+  name: Annotated[str, Field(alias="name")]
+
+
+class PublicCategoryInsert(BaseModel):
+  id: Annotated[int | None, Field(alias="id")]
+  name: Annotated[str, Field(alias="name")]
+
+
+class PublicCategoryUpdate(BaseModel):
+  id: Annotated[int | None, Field(alias="id")]
+  name: Annotated[str | None, Field(alias="name")]
+
+
+class PublicMemesSelect(BaseModel):
+  category: Annotated[int | None, Field(alias="category")]
+  created_at: Annotated[datetime.datetime, Field(alias="created_at")]
+  id: Annotated[int, Field(alias="id")]
+  metadata: Annotated[Json[Any] | None, Field(alias="metadata")]
+  name: Annotated[str, Field(alias="name")]
+  status: Annotated[MemeStatus | None, Field(alias="status")]
+
+
+class PublicMemesInsert(BaseModel):
+  category: Annotated[int | None, Field(alias="category")]
+  created_at: Annotated[datetime.datetime, Field(alias="created_at")]
+  id: Annotated[int | None, Field(alias="id")]
+  metadata: Annotated[Json[Any] | None, Field(alias="metadata")]
+  name: Annotated[str, Field(alias="name")]
+  status: Annotated[MemeStatus | None, Field(alias="status")]
+
+
+class PublicMemesUpdate(BaseModel):
+  category: Annotated[int | None, Field(alias="category")]
+  created_at: Annotated[datetime.datetime | None, Field(alias="created_at")]
+  id: Annotated[int | None, Field(alias="id")]
+  metadata: Annotated[Json[Any] | None, Field(alias="metadata")]
+  name: Annotated[str | None, Field(alias="name")]
+  status: Annotated[MemeStatus | None, Field(alias="status")]
+
+
+class PublicTodosViewSelect(BaseModel):
+  details: Annotated[str | None, Field(alias="details")]
+  id: Annotated[int | None, Field(alias="id")]
+  user_id: Annotated[int | None, Field(alias="user-id")]
+
+
+class PublicUsersViewSelect(BaseModel):
+  id: Annotated[int | None, Field(alias="id")]
+  name: Annotated[str | None, Field(alias="name")]
+  status: Annotated[UserStatus | None, Field(alias="status")]
+
+
+class PublicAViewSelect(BaseModel):
+  id: Annotated[int | None, Field(alias="id")]
+
+
+class PublicTodosMatviewSelect(BaseModel):
+  details: Annotated[str | None, Field(alias="details")]
+  id: Annotated[int | None, Field(alias="id")]
+  user_id: Annotated[int | None, Field(alias="user-id")]
+"`)
+})