@javalabs/prisma-client 1.0.0

package/src/scripts/data-migration/batch-migrator.ts

@@ -0,0 +1,569 @@
+import { Logger } from "@nestjs/common";
+import { PrismaClient } from "@prisma/client";
+import { EntityType, ColumnSchema, EnumCastValue } from "./types";
+import { DataTransformer } from "./data-transformer";
+import { SchemaUtils } from "./schema-utils";
+import { DatabaseConnections } from "./types";
+import { DependencyResolver } from "./dependency-resolver";
+import { ForeignKeyManager } from "./foreign-key-manager"; // Assuming this exists and might be useful later
+
+export class BatchMigrator {
+  private readonly logger = new Logger("BatchMigrator");
+  private readonly BATCH_SIZE = 10; // Consider making this configurable
+  // Removed typecastManager as it wasn't used and casting is handled inline
+
+  constructor(
+    private readonly dataTransformer: DataTransformer,
+    private readonly schemaUtils: SchemaUtils,
+    private readonly connections: DatabaseConnections,
+    private readonly dependencyResolver: DependencyResolver, // Keep for dependency checks
+    private readonly schemaCache: Record<string, ColumnSchema[]> = {},
+    private readonly targetSchemaCache: Record<string, ColumnSchema[]> = {} // Removed ForeignKeyManager from constructor if not used directly here
+  ) {}
+
+  private async checkTableHasData(
+    tenantId: string,
+    tableName: string
+  ): Promise<boolean> {
+    // Keep this function as is
+    try {
+      const result = await this.connections.targetPool.query(
+        `SELECT EXISTS (SELECT 1 FROM "${tenantId}"."${tableName}" LIMIT 1)`
+      );
+      return result.rows[0]?.exists || false;
+    } catch (error) {
+      this.logger.warn(
+        `Error checking data existence for ${tableName}: ${error.message}`
+      );
+      return false;
+    }
+  }
+
+  async migrateEntityDataInBatches(
+    prisma: PrismaClient, // Prisma client (likely connected to target public schema)
+    entity: EntityType, // Config object for the table being migrated
+    providerId: number | null, // Provider ID to filter source data (null if not filtering)
+    targetSchema: string // The schema in the TARGET database where data should be inserted
+  ) {
+    try {
+      // Ensure target schema exists (important for tenant schemas if used)
+      if (targetSchema !== "public") {
+        await this.ensureSchemaExists(targetSchema);
+      }
+
+      // Optional: Dependency check (can be complex with dynamic filtering)
+      // Consider if this check is still reliable or needed with the new strategy
+      // const dependencies = await this.dependencyResolver.analyzeDependencies();
+      // ... dependency check logic ...
+
+      // Disable foreign key checks in the target schema for the duration of this batch
+      // Note: This applies to the connection used by targetPool, ensure it targets the correct DB if replicas are used
+      await this.connections.targetPool.query(
+        `SET session_replication_role = 'replica';`
+      );
+
+      const { name: tableName, idField, filterColumn, filterVia } = entity;
+
+      this.logger.log(
+        `Migrating ${tableName} -> target schema '${targetSchema}'. ${
+          providerId
+            ? `Filtering source by Provider ID: ${providerId}`
+            : "Migrating all source records."
+        }`
+      );
+
+      try {
+        // Get source and target schemas
+        const sourceTableSchema = await this.getSourceSchema(tableName);
+        const targetTableSchema = await this.getTargetSchema(
+          targetSchema,
+          tableName
+        );
+
+        // Validate schemas
+        if (!sourceTableSchema.length) {
+          this.logger.warn(
+            `Source table ${tableName} schema not found. Skipping.`
+          );
+          return;
+        }
+        if (!targetTableSchema.length) {
+          this.logger.warn(
+            `Target table ${tableName} schema '${targetSchema}' not found. Skipping.`
+          );
+          return;
+        }
+
+        // --- DYNAMIC SOURCE QUERY GENERATION ---
+        let selectQuery: string;
+        let queryParams: any[] = [];
+
+        // Base query selects all columns from the source table
+        // Using alias 't' for the primary table
+        const columnList = sourceTableSchema
+          .map((col) => `"${col.column_name}"`)
+          .join(", ");
+        let fromClause = `FROM "${tableName}" t`;
+        let whereClause = "";
+
+        if (providerId && filterColumn) {
+          // Filtering is needed
+          if (filterVia) {
+            // Filter through an intermediate table (JOIN needed)
+            // Assumes intermediate table links via its primary key (e.g., 'id') to filterColumn
+            // Assumes intermediate table links to providers via 'provider_id'
+            // Example: Migrating 'invoices' (t) via 'transactions' (j)
+            // Needs filterColumn='transaction_id' (linking t to j) and filterVia='transactions'
+            // Final filter is j.provider_id = $1
+            // NOTE: This makes assumptions! A more robust solution might need more config.
+            // Let's assume filterVia table has a primary key named 'id' and a 'provider_id' column
+            const joinTable = filterVia;
+            const joinCondition = `t."${filterColumn}" = j.id`; // Assuming PK of joinTable is 'id'
+            const providerFilter = `j.provider_id = $1`; // Assuming FK in joinTable is 'provider_id'
+
+            fromClause = `FROM "${tableName}" t JOIN "${joinTable}" j ON ${joinCondition}`;
+            whereClause = `WHERE ${providerFilter}`;
+            queryParams = [providerId];
+            this.logger.log(`Using JOIN filter: ${fromClause} ${whereClause}`);
+          } else {
+            // Direct filter on the table itself
+            whereClause = `WHERE t."${filterColumn}" = $1`;
+            queryParams = [providerId];
+            this.logger.log(`Using direct filter: ${whereClause}`);
+          }
+        } else {
+          this.logger.log(
+            `No providerId filter applied for ${tableName}. Selecting all records.`
+          );
+          // No filtering needed (e.g., migrating public tables or filteredPublic without providerId)
+        }
+
+        selectQuery = `SELECT t.* ${fromClause} ${whereClause}`;
+        // --- END DYNAMIC SOURCE QUERY GENERATION ---
+
+        // Execute query and process data
+        const sourceData = await this.executeSourceQuery(
+          selectQuery,
+          queryParams
+        );
+        const totalRecords = sourceData.rows.length;
+
+        this.logger.log(
+          `Found ${totalRecords} ${tableName} records in source to migrate to '${targetSchema}'.`
+        );
+
+        if (totalRecords === 0) {
+          this.logger.log(
+            `No records to migrate for ${tableName} with current filter. Skipping processing.`
+          );
+          // Re-enable FK checks before returning
+          await this.connections.targetPool.query(
+            `SET session_replication_role = 'origin';`
+          );
+          return;
+        }
+
+        // Use the specific idField from the entity config
+        const primaryKeyField = idField;
+
+        await this.processRecords(
+          prisma, // Pass the main Prisma client
+          targetSchema, // Pass the specific target schema for insertion
+          tableName,
+          primaryKeyField,
+          sourceData.rows,
+          sourceTableSchema,
+          targetTableSchema
+        );
+      } catch (error) {
+        this.logger.error(
+          `Error during migration step for ${tableName} to schema '${targetSchema}': ${error.message}`
+        );
+        // Log context for better debugging
+        console.error(`Entity Config:`, JSON.stringify(entity));
+        console.error(`Provider ID used for filter:`, providerId);
+        // console.error(`Generated Select Query:`, selectQuery); // selectQuery might be out of scope here
+        // console.error(`Query Params:`, JSON.stringify(queryParams)); // queryParams might be out of scope here
+        // Rethrow to be caught by the outer try/catch in DataMigrationTool
+        throw error;
+      }
+    } catch (error) {
+      // Catch errors from initial checks or schema operations
+      this.logger.error(
+        `Error preparing migration for ${entity.name} to schema '${targetSchema}': ${error.message}`
+      );
+      throw error; // Rethrow to be caught by DataMigrationTool
+    } finally {
+      // ALWAYS re-enable foreign key checks, even if errors occurred
+      try {
+        await this.connections.targetPool.query(
+          `SET session_replication_role = 'origin';`
+        );
+      } catch (finallyError) {
+        this.logger.error(
+          `Failed to reset session_replication_role: ${finallyError.message}`
+        );
+      }
+    }
+  }
+
+  // --- Helper Functions (ensureSchemaExists, getSourceSchema, getTargetSchema, executeSourceQuery, getPrimaryKeyField) ---
+  // Keep these mostly as they are, but ensure getPrimaryKeyField uses the correct idField from EntityType if needed
+  // (Current implementation queries information_schema, which is fine, but ensure it uses the correct targetSchema)
+
+  private async ensureSchemaExists(schemaName: string): Promise<void> {
+    // Check if schema exists
+    const schemaExistsResult = await this.connections.targetPool.query(
+      `SELECT schema_name FROM information_schema.schemata WHERE schema_name = $1`,
+      [schemaName]
+    );
+
+    if (schemaExistsResult.rows.length === 0) {
+      this.logger.log(`Schema '${schemaName}' does not exist. Creating...`);
+      // Create schema if it doesn't exist - Needs structure copied from public
+      await this.schemaUtils.createSchema(schemaName);
+      this.logger.log(`Schema '${schemaName}' created.`);
+
+      // Optional: Baseline migration record for the new schema if using Prisma Migrate
+      try {
+        await this.connections.targetPool.query(`
+                    INSERT INTO "${schemaName}"."_prisma_migrations" (id, checksum, finished_at, migration_name, logs, rolled_back_at, started_at, applied_steps_count)
+                    SELECT id, checksum, finished_at, migration_name, logs, rolled_back_at, started_at, applied_steps_count
+                    FROM "public"."_prisma_migrations" WHERE migration_name LIKE '%_init'
+                    ON CONFLICT DO NOTHING;
+                `);
+        this.logger.log(
+          `Attempted to copy initial migration record to schema '${schemaName}'.`
+        );
+      } catch (migrationError) {
+        this.logger.warn(
+          `Could not copy baseline migration to schema '${schemaName}': ${migrationError.message}`
+        );
+      }
+    }
+  }
+
+  private async getSourceSchema(tableName: string): Promise<ColumnSchema[]> {
+    const cacheKey = `source.${tableName}`;
+    if (!this.schemaCache[cacheKey]) {
+      this.logger.debug(`Cache miss for source schema: ${tableName}`);
+      this.schemaCache[cacheKey] = await this.schemaUtils.getTableSchema(
+        tableName,
+        "source",
+        "public" // Source is always public schema in this context
+      );
+    }
+    return this.schemaCache[cacheKey];
+  }
+
+  private async getTargetSchema(
+    schema: string, // Target schema name (could be 'public' or tenantId)
+    tableName: string
+  ): Promise<ColumnSchema[]> {
+    const cacheKey = `${schema}.${tableName}`;
+    if (!this.targetSchemaCache[cacheKey]) {
+      this.logger.debug(`Cache miss for target schema: ${cacheKey}`);
+      this.targetSchemaCache[cacheKey] = await this.schemaUtils.getTableSchema(
+        tableName,
+        "target",
+        schema // Use the provided target schema name
+      );
+    }
+    return this.targetSchemaCache[cacheKey];
+  }
+
+  private async executeSourceQuery(query: string, params: any[]): Promise<any> {
+    try {
+      this.logger.debug(
+        `Executing source query: ${query.replace(
+          /\s\s+/g,
+          " "
+        )} || PARAMS: ${JSON.stringify(params)}`
+      );
+      const result = await this.connections.sourcePool.query(query, params);
+      this.logger.debug(`Query returned ${result.rows?.length || 0} rows`);
+      return result;
+    } catch (error) {
+      this.logger.error(`Error executing source query: ${error.message}`);
+      this.logger.error(`Query was: ${query}`);
+      this.logger.error(`Params were: ${JSON.stringify(params)}`);
+      throw error;
+    }
+  }
+
+  private async getPrimaryKeyField(
+    schemaName: string,
+    tableName: string
+  ): Promise<string | null> {
+    // This function remains useful for verifying the actual PK if needed,
+    // but we primarily rely on the idField from the config now.
+    try {
+      const result = await this.connections.targetPool.query(
+        `
+        SELECT kcu.column_name
+        FROM information_schema.table_constraints tc
+        JOIN information_schema.key_column_usage kcu
+          ON tc.constraint_name = kcu.constraint_name
+          AND tc.table_schema = kcu.table_schema
+        WHERE tc.constraint_type = 'PRIMARY KEY'
+          AND tc.table_schema = $1
+          AND tc.table_name = $2
+      `,
+        [schemaName, tableName]
+      );
+      return result.rows[0]?.column_name || null;
+    } catch (error) {
+      this.logger.error(
+        `Error getting primary key for ${schemaName}.${tableName}: ${error.message}`
+      );
+      return null;
+    }
+  }
+
+  // --- processRecords ---
+  // This function needs significant changes to correctly handle the targetSchema
+
+  private async processRecords(
+    prisma: PrismaClient, // Main prisma client
+    targetSchema: string, // The ACTUAL schema to insert/update into
+    tableName: string,
+    idField: string, // Use idField from config
+    records: any[],
+    sourceSchema: ColumnSchema[],
+    targetSchemaInfo: ColumnSchema[] // Renamed to avoid confusion
+  ): Promise<void> {
+    if (!idField) {
+      this.logger.error(
+        `Cannot process records for ${tableName}: idField is missing in configuration.`
+      );
+      return; // Or throw error
+    }
+
+    for (const record of records) {
+      let recordId = record[idField];
+      try {
+        this.logger.debug(
+          `Processing record ${
+            recordId ?? "(no id found)"
+          } for ${targetSchema}.${tableName}`
+        );
+
+        if (!record || Object.keys(record).length === 0) {
+          this.logger.warn(`Empty record found for ${tableName}, skipping`);
+          continue;
+        }
+        if (!recordId) {
+          this.logger.warn(
+            `Record missing configured ID field '${idField}' in source data, skipping: ${JSON.stringify(
+              record
+            )}`
+          );
+          continue;
+        }
+
+        // Transform data using the target schema *info* for type checking etc.
+        const transformedData = await this.dataTransformer.transformRecord(
+          record,
+          sourceSchema,
+          targetSchemaInfo, // Use the schema structure info
+          targetSchema // Pass target schema for enum validation context etc.
+        );
+
+        // Prepare data for raw query, ensuring correct types and casting strings
+        const processedData = Object.entries(transformedData).reduce(
+          (acc, [key, value]) => {
+            const columnSchema = targetSchemaInfo.find(
+              (col) => col.column_name === key
+            );
+            if (!columnSchema) return acc; // Skip columns not in target schema
+
+            const columnName = `"${columnSchema.column_name}"`; // Quote column names
+
+            // --- Helper function to escape values for SQL E'' strings ---
+            const escapeValue = (val: any): string => {
+              if (val === null || val === undefined) return "NULL";
+              if (typeof val === "boolean") return val ? "TRUE" : "FALSE";
+              if (typeof val === "number") return String(val);
+              // Escape single quotes and backslashes for E'' syntax
+              return String(val).replace(/'/g, "''").replace(/\\/g, "");
+            };
+            // --- End escapeValue ---
+
+            if (value === null || value === undefined) {
+              acc[columnName] = "NULL";
+              return acc;
+            }
+
+            // Special handling for EnumCastValue objects from transformer
+            if (
+              typeof value === "object" &&
+              value !== null &&
+              value["needsEnumCast"]
+            ) {
+              const enumValue = value as EnumCastValue;
+              const schemaPrefix =
+                targetSchema === "public" ? '"public".' : `"${targetSchema}".`;
+              const quotedEnumType = `"${enumValue.enumType}"`;
+              const escapedEnumValue = escapeValue(enumValue.value);
+              // Enum values are typically strings, use E''
+              acc[
+                columnName
+              ] = `CAST(E'${escapedEnumValue}' AS ${schemaPrefix}${quotedEnumType})`;
+              return acc;
+            }
+
+            // Handle standard types - Use data_type primarily, fallback to udt_name for enums/user-defined
+            let targetType = columnSchema.data_type.toLowerCase();
+            let udtName = columnSchema.udt_name;
+            let sqlValue: string;
+            let requiresQuotes = false;
+
+            // Determine if quotes are needed based on type category
+            if (
+              [
+                "text",
+                "varchar",
+                "character varying",
+                "char",
+                "timestamp with time zone",
+                "timestamptz",
+                "timestamp without time zone",
+                "timestamp",
+                "date",
+                "uuid",
+                "json",
+                "jsonb",
+              ].includes(targetType) ||
+              targetType.includes("enum") ||
+              (targetType === "user-defined" && udtName?.startsWith("enum_"))
+            ) {
+              requiresQuotes = true;
+            }
+
+            // Escape the value appropriately
+            const escaped = escapeValue(value);
+            if (escaped === "NULL") {
+              // Use SQL NULL keyword directly
+              sqlValue = "NULL";
+            } else if (requiresQuotes) {
+              sqlValue = `E'${escaped}'`; // Use E'' for strings, dates, enums, json, etc.
+            } else {
+              sqlValue = escaped; // Use raw value for numbers, booleans
+            }
+
+            // Determine necessary casting based on target type
+            let castExpression = "";
+            if (targetType.includes("timestamp"))
+              castExpression = "::timestamp with time zone";
+            else if (targetType === "date") castExpression = "::date";
+            else if (
+              targetType === "integer" ||
+              targetType === "int" ||
+              targetType === "int4"
+            )
+              castExpression = "::integer";
+            else if (targetType === "bigint" || targetType === "int8")
+              castExpression = "::bigint";
+            else if (targetType === "smallint" || targetType === "int2")
+              castExpression = "::smallint";
+            else if (targetType === "numeric" || targetType === "decimal")
+              castExpression = "::numeric";
+            else if (targetType === "real" || targetType === "float4")
+              castExpression = "::real";
+            else if (
+              targetType === "double precision" ||
+              targetType === "float8"
+            )
+              castExpression = "::double precision";
+            else if (targetType === "boolean" || targetType === "bool")
+              castExpression = "::boolean";
+            else if (targetType === "json" || targetType === "jsonb")
+              castExpression = `::${targetType}`;
+            else if (targetType === "uuid") castExpression = "::uuid";
+            else if (targetType === "text" || targetType.includes("char"))
+              castExpression = "::text";
+            else if (
+              targetType === "user-defined" &&
+              udtName?.startsWith("enum_")
+            ) {
+              const schemaPrefix =
+                targetSchema === "public" ? '"public".' : `"${targetSchema}".`;
+              castExpression = `::${schemaPrefix}"${udtName}"`;
+            }
+
+            acc[columnName] = `${sqlValue}${castExpression}`;
+
+            return acc;
+          },
+          {} as Record<string, string> // Accumulator holds SQL value strings
+        );
+
+        // Filter out entries where processedData might be undefined/invalid if needed
+        const validProcessedData = Object.entries(processedData).reduce(
+          (acc, [key, val]) => {
+            if (
+              val !== undefined &&
+              val !== "NULL" &&
+              val !== "E''" &&
+              val !== "E'undefined'" &&
+              val !== "E'null'"
+            ) {
+              // Additional checks for empty/invalid strings
+              acc[key] = val;
+            }
+            return acc;
+          },
+          {}
+        );
+
+        const columns = Object.keys(validProcessedData);
+        const valuesString = Object.values(validProcessedData).join(", "); // Values are already SQL strings
+
+        if (columns.length === 0) {
+          this.logger.warn(
+            `Record ${recordId} for ${tableName} resulted in no valid columns to insert/update after processing. Skipping.`
+          );
+          continue;
+        }
+
+        // Construct the SET clause for UPDATE
+        const updateSetClauses = columns
+          .filter((col) => col !== `"${idField}"`) // Don't update the PK itself
+          .map((col) => `${col} = EXCLUDED.${col}`) // Use EXCLUDED to get the value proposed for insertion
+          .join(", ");
+
+        // Ensure target schema and table name are quoted
+        const quotedSchemaTable = `"${targetSchema}"."${tableName}"`;
+        const quotedIdField = `"${idField}"`;
+
+        // Only include DO UPDATE clause if there are columns to update
+        const conflictClause = updateSetClauses
+          ? `ON CONFLICT (${quotedIdField}) DO UPDATE SET ${updateSetClauses}`
+          : `ON CONFLICT (${quotedIdField}) DO NOTHING`;
+
+        const query = `
+          INSERT INTO ${quotedSchemaTable} (${columns.join(", ")})
+          VALUES (${valuesString})
+          ${conflictClause}
+        `;
+
+        // Execute using targetPool connection for raw SQL flexibility
+        this.logger.debug(`Executing Upsert: ${query.replace(/\s\s+/g, " ")}`);
+        await this.connections.targetPool.query(query);
+      } catch (error) {
+        // Improved error logging
+        this.logger.error(
+          `Error processing record ID '${
+            recordId ?? "(unknown)"
+          }' for ${targetSchema}.${tableName}: ${error.message}`
+        );
+        this.logger.error(`Record data: ${JSON.stringify(record)}`);
+        // Consider logging transformedData and processedData as well for deep debugging
+        // this.logger.error(`Transformed data: ${JSON.stringify(transformedData)}`);
+        // this.logger.error(`Processed data (SQL values): ${JSON.stringify(processedData)}`);
+        // console.error("Underlying Error Stack:", error); // Log the original error stack
+        throw error; // Re-throw to stop the batch or be handled by the caller
+      }
+    }
+  }
+}